WWDC Quick Look 💓 By SwiftGGTeam
Dive into Core AI model authoring and optimization

Dive into Core AI model authoring and optimization

观看原视频

Highlight

Core AI 为 Apple Silicon 端侧 AI 部署提供了完整的 Python 工具链:从 PyTorch 模型导出、配置驱动压缩,到可视化调试和自定义 Metal 内核打包,全部在一个工作流内完成。

核心内容

端侧模型部署的老问题

把深度学习模型搬上 iPhone 或 Mac,开发者一直面临三个麻烦:模型太大装不下、量化后精度掉点找不到原因、自定义算子难以和模型一起打包分发。以前用 Core ML,转完模型上机一跑,精度掉了只能靠猜是哪一层出了问题,或者写大量脚本逐层 dump 张量对比。自定义算子更麻烦,需要额外管理 Metal 库文件和编译管线。(00:14

Core AI 的完整工作流

Apple 这次把 PyTorch 到 Apple Silicon 的部署链路彻底打通了。安装 pip install coreai-torch 后,你手头的 PyTorch 模型可以直接走一条完整的流水线:用 torch.export 捕获计算图,经 TorchConverter 转成 Core AI 内部表示,再经 optimize() 生成 .aimodel 资产,最后从 Python 直接加载推理。整个过程不需要离开熟悉的 Python 环境。(03:27

量化掉点怎么查

压缩是把大模型搬上设备的关键。以 SAM3(Segment Anything Model 3)为例,这个 8.5 亿参数的图像分割模型,原始大小超过 3GB。用 coreai-optpresets.w4 做 4-bit 全量化后,模型压到了 430MB,但一个被遮挡的花朵检测不到了。(06:02

问题出在哪?Core AI Debugger 登场。它在 Mac 上运行,能可视化模型结构、在真机上执行推理,还能把 Core AI 的中间张量和 PyTorch 参考运行结果逐层对比。Debugger 自动识别同步点,用 PSNR(峰值信噪比)指标给每层打分,绿色表示吻合,黄色表示轻度偏离,红色表示严重差异。排序后发现,低 PSNR 的层几乎全部来自 detector decoder —— 这个只占模型 4% 参数的模块对量化极其敏感。把它从量化配置里剔除后,精度恢复,模型仍然只有原来的很小一部分。(10:40

自定义 Metal 内核随模型打包

标准算子库永远覆盖不了所有需求。Core AI 允许你手写 Metal Shading Language(MSL)内核,和 PyTorch 参考实现一起注册到 TorchMetalKernel,转换时直接嵌入 .aimodel 资产。部署时 App 只加载一个文件,不需要额外管理 Metal 库。(21:12

模型重写让推理快 76%

把模型拆成多个独立函数是另一个高级技巧。SAM3 的图像编码器、文本编码器和检测头被拆成三个入口。用户换了一个提示词(从”花”换成”蝴蝶”),只需要重新运行文本编码器和检测头,图像编码结果直接复用。第二次推理比完整跑一遍快了 76%。(24:56

详细内容

PyTorch 模型导出与 Core AI 转换

03:27

第一步是用 torch.export 捕获模型的完整计算图,包括权重、操作和形状信息。

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(256, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = MLP().eval()
example_input = (torch.randn(1, 256),)
exported_program = torch.export.export(model, example_input)

关键点:

  • torch.export.export() 捕获静态计算图,权重、操作和形状全部冻结在 exported_program
  • example_input 用于推导张量形状,模型中不能有依赖输入数据的动态控制流
  • .eval() 关闭 dropout 等训练专用层,确保推理行为一致

第二步用 TorchConverter 转成 Core AI 表示,优化并保存为 .aimodel

import coreai
import coreai_torch
from coreai.runtime import NDArray

converter = coreai_torch.TorchConverter()
converter.add_exported_program(
    exported_program,
    input_names=["features"], output_names=["logits"])
core_ai_program = converter.to_coreai()

core_ai_program.optimize()
asset = core_ai_program.save_asset("mlp.aimodel")

specialized_model = await AIModel.load("mlp.aimodel")
specialized_function = specialized_model.load_function("main")
result = await specialized_function({"features": NDArray(example[0].numpy())})

关键点:

  • TorchConverter 是 PyTorch 到 Core AI 的桥梁,API 设计和 Core ML Tools 类似
  • optimize() 根据目标平台做算子融合和内存优化
  • save_asset() 生成 .aimodel 文件,这是 Apple Silicon 的原生执行格式
  • 推理输入用 NDArray 包装 numpy 数组,通过字典按名称映射

配置驱动的模型压缩

05:54

coreai-opt 支持 int4、int8、FP4、FP8 多种压缩粒度,通过配置文件决定哪些层压缩、哪些层保留原精度。

import coreai_opt

# 使用预设的 4-bit 每通道对称量化
config = coreai_opt.presets.w4
config.execution_mode = coreai_opt.ExecutionMode.EAGER

quantizer = coreai_opt.Quantizer(config)
quantizer.initialize(model, example_inputs)
compressed_model = quantizer.finalize()

关键点:

  • presets.w4 是一行配置的快捷方式,背后对应每通道 4-bit 对称量化
  • ExecutionMode.EAGER 适合权重量化,GRAPH 模式适合激活量化
  • initialize() 需要传入示例输入用于校准统计
  • 也可以传入大量数据做量化感知训练(QAT),精度损失更小

Core AI Debugger 定位量化问题

10:40

Debugger 的核心能力是对比 PyTorch 参考运行和 Core AI 端侧运行的中间张量。首先在 Python 中保存 PyTorch 的中间结果:

# 保存 PyTorch 中间张量
intermediates = coreai_torch.save_intermediates(
    original_model,
    quantized_model,
    example_inputs
)

然后在 Debugger 中加载这个文件作为参考运行,启动对比会话。Debugger 自动对齐两个计算图的节点,生成同步点,并计算 PSNR。绿色节点表示张量相似度高,黄色表示中度偏离,红色表示严重差异。按 PSNR 排序后,可以快速定位问题层。(15:00

关键点:

  • save_intermediates API 在 PyTorch 端捕获每个操作的输出张量
  • Debugger 自动识别同步点,不需要手动对齐节点
  • PSNR 是默认相似度指标,可以根据模型类型切换其他指标
  • 左侧导航器按 PyTorch 模块分组,大型模型的导航体验和代码结构一致

自定义 Metal 内核

21:12

以 SiLU(Sigmoid Linear Unit)激活函数为例,同时提供 PyTorch 参考实现和 MSL 内核:

import torch
from coreai_torch.dsl import TorchMetalKernel, MetalParameter

def silu_torch(x):
    return x * torch.sigmoid(x)

SILU_MSL = """
float val = float(x[gid]);
float sig = 1.0f / (1.0f + exp(-val));
y[gid] = TYPE(val * sig);
"""

silu_kernel = TorchMetalKernel(
    name="fused_silu",
    input_names=["x"],
    result_names=["y"],
    src=SILU_MSL,
    torch_defn=silu_torch,
    metal_params=[MetalParameter("gid", "uint", "thread_position_in_grid")],
    template_dtypes={"x": "TYPE"},
)

在模型中使用这个内核:

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(256, 256)

    def forward(self, x):
        h = self.linear(x)
        n = h.numel()
        return silu_kernel(
            h,
            threads_per_grid_size=(n, 1, 1),
            threads_per_thread_group=(min(n, 256), 1, 1),
            result_shapes=[h.shape],
        )

exported_program = torch.export.export(MyModel(), (torch.randn(1, 256),))

converter = coreai_torch.TorchConverter()
converter.register_custom_kernels([silu_kernel])
converter.add_exported_program(exported_program,
                               input_names=["x"], output_names=["y"])
deployable = converter.to_coreai()

关键点:

  • torch_defn 是给 torch.export 看的参考实现,用于形状推导和 Debugger 对比
  • src 是实际在 GPU 上执行的 MSL 代码,TYPE 宏由 template_dtypes 注入实际数据类型
  • MetalParameter 把 MSL 中的变量绑定到 Metal 的内置属性,如 thread_position_in_grid
  • result_shapes 在每次调用时传入,支持动态形状输入
  • MSL 代码随模型一起打包进 .aimodel,部署时无需额外文件

模型重写成多个独立函数

24:56

把 SAM3 拆成 image_encode、text_encode、detect 三个入口,每个可以独立压缩和调用:

# 三个模块分别导出
image_exported = torch.export.export(image_encoder, image_input)
text_exported = torch.export.export(text_encoder, text_input)
detect_exported = torch.export.export(detector, detect_input)

converter = coreai_torch.TorchConverter()
converter.add_exported_program(image_exported,
                               input_names=["image"], output_names=["image_emb"],
                               entrypoint_name="image_encode")
converter.add_exported_program(text_exported,
                               input_names=["text"], output_names=["text_emb"],
                               entrypoint_name="text_encode")
converter.add_exported_program(detect_exported,
                               input_names=["image_emb", "text_emb"], output_names=["masks"],
                               entrypoint_name="detect")

deployable = converter.to_coreai()
asset = deployable.save_asset("sam3_split.aimodel")

加载后按需调用:

model = await AIModel.load("sam3_split.aimodel")

# 第一次:完整流程
image_fn = model.load_function("image_encode")
text_fn = model.load_function("text_encode")
detect_fn = model.load_function("detect")

image_emb = await image_fn({"image": image_ndarray})
text_emb = await text_fn({"text": text_ndarray})
masks = await detect_fn({"image_emb": image_emb, "text_emb": text_emb})

# 换提示词后:只跑 text_encode + detect
text_emb_new = await text_fn({"text": new_text_ndarray})
masks_new = await detect_fn({"image_emb": image_emb, "text_emb": text_emb_new})

关键点:

  • 每个 add_exported_programentrypoint_name 定义独立入口
  • 一个 .aimodel 资产可以包含多个可调用函数
  • 图像编码结果 image_emb 可以缓存复用,换提示词时跳过最重的计算
  • 每个入口可以独立配置压缩策略,detector 保持原精度,两个 encoder 做 4-bit 压缩

核心启发

1. 给现有 App 加一个端侧图像分割功能

coreai-models 仓库里的 SAM3 示例,配合 coreai-opt 的混合精度配置,把模型压到 500MB 以内。用 Core AI Debugger 验证分割精度满足产品要求后,直接集成到照片编辑 App 中。入口 API 是 AIModel.load()load_function(),和加载普通模型没有区别。

2. 做一个”换提示词实时重分割”的交互体验

参考 Session 中的三函数拆分思路,把图像编码和文本编码拆开。用户画一个框或输入一个提示词,图像只编码一次,后续每次改提示词只跑文本编码和检测头。实现思路:缓存 image_encode 的输出张量,复用到后续的 text_encode + detect 调用中。

3. 把自定义激活函数打包进模型资产

如果你的模型用了论文里新提出的激活函数,而 Core AI 的标准算子库还没覆盖,用 TorchMetalKernel 手写一个 MSL 实现。PyTorch 参考实现负责形状推导和调试对比,MSL 负责实际执行。转换后一个 .aimodel 文件搞定部署。

4. 用 AI Skills 加速团队上手 Core AI

coreai-models 仓库附带了一套 Agent Skills,可以安装到 Cursor 或 Copilot 中。这些 Skills 包含 Apple 工程师的最佳实践,能直接把”我想把 SAM3 部署到 iPhone 上”这种自然语言需求翻译成具体的转换、压缩、拆分脚本。团队新人不需要翻完整文档就能产出可运行的代码。

5. 用 Debugger 建立量化回归测试流水线

在 CI 中集成 save_intermediates + Debugger 对比:每次模型更新后,自动对比 PyTorch 参考输出和量化后的 Core AI 输出,用 PSNR 阈值判断是否通过。把以前靠人眼检查的环节变成自动化的数值门禁。

关联 Session

评论

GitHub Issues · utterances