WWDC Quick Look 💓 By SwiftGGTeam
中文 EN 日本語
Train your machine learning and AI models on Apple GPUs

Train your machine learning and AI models on Apple GPUs

WWDC 2024
观看原视频

Highlight

Apple Silicon 的 GPU 在机器学习训练上有天然优势：强大的并行计算能力、统一内存架构让 GPU 直接访问大量内存、不需要在多台机器间分布式部署模型权重。今年 Apple 对 PyTorch MPS backend 和 JAX Metal backend 都做了升级，但重点是 PyTorch。

核心内容

训练大模型的第一道门槛是显存。NVIDIA 的 GPU 有独立的显存，模型参数多了就得拆到多张卡上做分布式训练，工程复杂度陡然上升。Apple Silicon 的统一内存架构让 CPU 和 GPU 共享同一块内存——一台 128GB 内存的 Mac Studio，GPU 可以直接访问全部 128GB，不需要把参数在 CPU 和 GPU 之间来回拷贝。

Session 的核心更新集中在 PyTorch MPS backend 上。去年 WWDC23 MPS backend 进入 beta，今年三项改进直接瞄准 transformer 模型：int8/int4 量化让大模型能塞进设备内存；fused scaled dot product attention 把多头注意力的矩阵乘法、缩放、softmax 融合成一次 GPU kernel 调用，减少 dispatch 开销；统一内存支持消除了 CPU 和 GPU 之间冗余的 tensor 拷贝。这三项改进让 HuggingFace Top-50 最热门的 transformer 模型在 MPS 上开箱即用。

JAX Metal backend 也有更新：新增 BFloat16 数据类型支持、NDArray 索引操作、以及 padding/dilation 配置。MuJoCo（物理仿真）和 AXLearn（大规模深度学习）已经在生产中使用 JAX Metal backend。

详细内容

PyTorch MPS backend 的三项改进

1. int8 和 int4 整数量化（05:41）

训练时通常使用 32-bit 或 16-bit 浮点数。训练完成后，通过量化把参数从浮点数转为 8-bit 或 4-bit 整数，内存占用减半甚至减至四分之一。好处是模型更小、计算吞吐更高，精度损失通常很小。这让 7B 参数量的模型也能在 16GB 内存的 Mac 上运行。

2. Fused scaled dot product attention（06:26）

Scaled dot product attention 是 transformer 的核心运算：输入文本被拆成 query、key、value 三个 tensor，然后依次经过矩阵乘法、缩放、softmax 等操作。如果每步操作都单独 dispatch 到 GPU，开销会累积。fused SDPA 把这一系列操作融合成单次 kernel 调用，减少 GPU dispatch 开销。

3. 统一内存支持（07:02）

在传统架构下，CPU 和 GPU 各有独立内存，tensor 需要在两者之间拷贝。Apple Silicon 的统一内存让 tensor 只需要存在于主内存中，CPU 和 GPU 都能直接访问，省去了拷贝。

端到端工作流：OpenLLaMA v2 3B LoRA 微调

Session 演示了从下载模型到微调再到部署的完整流程（07:34）。核心代码如下：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 下载 OpenLLaMA v2 3B 模型和对应的 tokenizer
model = AutoModelForCausalLM.from_pretrained("openlm-research/open_llama_3b_v2")
tokenizer = AutoTokenizer.from_pretrained("openlm-research/open_llama_3b_v2")

# 用 PEFT 库挂载 LoRA adapter
from peft import LoraConfig, get_peft_model
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)

# 把模型送到 MPS 设备
model = model.to("mps")

关键点：

  • from_pretrained 直接从 HuggingFace 下载模型权重和 tokenizer，不需要手动转换格式
  • LoraConfig 定义 LoRA adapter 的参数：r=8 是 LoRA 的秩，target_modules 指定要微调的注意力层
  • .to("mps") 把模型送到 Apple Silicon GPU，这是唯一需要改的设备相关代码

训练部分使用 HuggingFace 的 Trainer 类（09:01）：

from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling

training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=4,
    num_train_epochs=10,
)
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
    train_dataset=tokenized_dataset,
)
trainer.train()

关键点：

  • per_device_train_batch_size=4 利用统一内存的大容量来设置较大的 batch size
  • num_train_epochs=10 跑 10 个 epoch，演示中从”像字典条目”的输出变成了有语境的回复
  • mlm=False 表示用 causal LM（自回归）模式，不是 masked LM

训练完成后，合并 adapter 和基础模型并保存（10:58）：

model = model.merge_and_unload()
model.save_pretrained("./fine_tuned_model")
tokenizer.save_pretrained("./fine_tuned_model")

部署：ExecuTorch + MPS Partitioner

Session 展示了用 ExecuTorch 部署 Meta LLaMA2 4-bit 量化模型到 iPad 的流程（11:25）：

# 克隆并安装 ExecuTorch，启用 MPS bindings
git clone https://github.com/pytorch/executorch.git
cd executorch
git submodule update --init --recursive
./install_requirements.sh --mps

ExecuTorch 使用 MPS Partitioner 分析计算图，把能识别的模式自动交给 MPS 设备加速。

JAX Metal backend 更新

JAX Metal backend 今年的三项更新（13:24）：

BFloat16 支持（14:52）：

import jax
import jax.numpy as jnp

# 创建 BFloat16 tensor
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.bfloat16)

NDArray 索引（15:19）：

# NumPy 风格的数组索引和更新
arr = jnp.ones((2, 2))
arr = arr.at[:, 0].divide(10)  # 第一列除以 10

Padding 和 dilation（15:40）：

# 正 padding：在元素之间插入间隔（dilation）
padded = jax.lax.pad(tensor, 0.0, [(0, 0, 1)])  # (前, 后, dilation)

# 负 padding：移除元素
trimmed = jax.lax.pad(tensor, 0.0, [(-1, 0, 0)])  # 移除第一个元素

Session 用 AXLearn 的 Fuji 7B 模型做了 JAX 推理演示，BFloat16 数据类型下 Metal backend 输出与 CPU 完全一致。

核心启发

  • 做什么：在 Mac 上用 PyTorch + LoRA 微调开源大语言模型。为什么值得做：统一内存让单机就能跑 3B-7B 参数量模型的微调，不需要多卡分布式，工程复杂度低一个量级。怎么开始：pip install torch transformers peft，下载一个 OpenLLaMA 或 LLaMA 模型，把 .to("cuda") 改成 .to("mps")，用 PEFT 的 LoRA config 微调。

  • 做什么：对已训练好的模型做 int8/int4 量化再部署到设备端。为什么值得做：量化后内存占用减半甚至减至四分之一，7B 模型能在 16GB 设备上运行，推理速度也有提升。怎么开始：在 PyTorch 中使用 torch.quantization API 对模型做训练后量化，验证精度损失在可接受范围内，再用 Core ML 或 ExecuTorch 部署。

  • 做什么：用 MLX 框架做 Apple Silicon 原生的模型实验。为什么值得做：MLX 专为 Apple Silicon 设计，原生支持统一内存和 JIT 编译，Python/Swift/C++ 多语言绑定，API 风格类似 NumPy 学习成本低。怎么开始：pip install mlx，参考 MLX GitHub 仓库中的 fine-tuning 和 image generation 示例。

关联 Session

评论

GitHub Issues · utterances