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
元の動画を見る

ハイライト

Apple Silicon の GPU は ML 訓練に天然の優位性——強力な並列計算、統合メモリアーキテクチャで GPU が大量メモリに直接アクセス、複数マシン間でのモデル重み分散が不要。今年 Apple は PyTorch MPS backend と JAX Metal backend の両方をアップグレード。重点は PyTorch。

主要内容

大規模モデル訓練の最初の壁はメモリ。NVIDIA GPU は独立 VRAM を持ち、パラメータが増えると複数カードに分散訓練が必要——エンジニアリング複雑度が急上昇。Apple Silicon の統合メモリは CPU と GPU が同一メモリを共有——128GB メモリの Mac Studio なら GPU が 128GB 全体に直接アクセス。CPU と GPU 間のパラメータコピーが不要。

セッションの核心更新は PyTorch MPS backend に集中。昨年 WWDC23 で MPS backend が beta 入り、今年 3 つの改善が transformer モデルを直接ターゲット：int8/int4 量子化で大規模モデルをデバイスメモリに収める；fused scaled dot product attention で multi-head attention の行列乗算、スケーリング、softmax を 1 回の GPU kernel 呼び出しに融合、dispatch オーバーヘッド削減；統合メモリサポートで CPU-GPU 間の冗長 tensor コピーを排除。この 3 改善で HuggingFace Top-50 人気 transformer モデルが MPS 上で開箱即使用可能。

JAX Metal backend も更新：BFloat16 データ型サポート、NDArray インデックス操作、padding/dilation 設定。MuJoCo（物理シミュレーション）と AXLearn（大規模深層学習）はすでに JAX Metal backend を本番利用。

詳細

PyTorch MPS backend の 3 改善

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

訓練時は通常 32-bit または 16-bit 浮動小数点。訓練後、量子化でパラメータを 8-bit または 4-bit 整数へ変換——メモリ使用量が半分または 1/4 に。モデルが小さく、計算スループット向上、精度損失は通常小さい。7B パラメータモデルも 16GB メモリ Mac で実行可能。

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

Scaled dot product attention は transformer の核心演算：入力テキストを query、key、value の 3 tensor に分割し、行列乗算、スケーリング、softmax などを順次実行。各ステップを個別 dispatch するとオーバーヘッドが蓄積。fused SDPA は一連の操作を 1 回の kernel 呼び出しに融合、GPU dispatch オーバーヘッド削減。

3. 統合メモリサポート（07:02）

従来アーキテクチャでは CPU と GPU が独立メモリ——tensor のコピーが必要。Apple Silicon 統合メモリでは tensor はメインメモリに 1 回存在し、CPU と GPU が直接アクセス——コピー不要。

エンドツーエンドワークフロー：OpenLLaMA v2 3B LoRA ファインチューニング

セッションはモデルダウンロードからファインチューニング、デプロイまでの完全フローをデモ（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 はファインチューニング対象の attention 層
  • .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 と base model をマージして保存（10:58）：

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

デプロイ：ExecuTorch + MPS Partitioner

セッションは 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 の今年の 3 更新（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)])  # 移除第一个元素

セッションは AXLearn の Fuji 7B モデルで JAX 推論デモ。BFloat16 データ型で Metal backend 出力は CPU と完全一致。

重要ポイント

  • 何をするか：Mac 上で PyTorch + LoRA によりオープンソース LLM をファインチューニング。なぜ価値があるか：統合メモリで単一マシンが 3B-7B パラメータモデルのファインチューニング可能——マルチ GPU 分散不要、エンジニアリング複雑度が 1 桁下がる。始め方：pip install torch transformers peft、OpenLLaMA または LLaMA モデルをダウンロード、.to("cuda") を .to("mps") に変更、PEFT の LoRA config でファインチューニング。

  • 何をするか：訓練済みモデルを int8/int4 量子化して端末デプロイ。なぜ価値があるか：量子化後メモリ使用量が半分または 1/4——7B モデルが 16GB デバイスで実行可能、推論速度も向上。始め方：PyTorch の torch.quantization API で post-training 量子化、許容可能な精度損失を検証、Core ML または ExecuTorch でデプロイ。

  • 何をするか：MLX フレームワークで Apple Silicon ネイティブのモデル実験。なぜ価値があるか：MLX は Apple Silicon 専用設計。統合メモリと JIT コンパイルをネイティブサポート、Python/Swift/C++ バインディング、NumPy 風 API で学習コスト低。始め方：pip install mlx、MLX GitHub リポジトリの fine-tuning と image generation サンプルを参照。

関連セッション

コメント

GitHub Issues · utterances