Optimize machine learning for Metal apps

ハイライト

PyTorch 2.0 の MPS バックエンドが Beta 段階に入り、Top 60 の常用オペレーターと自動混合精度に対応しました。TensorFlow Metal プラグインが 1.0 安定版をリリース。JAX が初めて Metal GPU 加速を獲得し、M2 Max 上で CPU 比平均 10 倍の高速化を実現。MPSGraph にはシリアル化形式 MPSGraphPackage と Int8 量子化サポートが追加されました。

主要内容

PyTorch 2.0 MPS バックエンド: 実験段階から実用へ

以前、Mac 上で PyTorch を使ってモデルをトレーニングする場合、MPS バックエンドのオペレーターサポートは限定的で、多くのモデルが動かないか頻繁に CPU にフォールバックしていました。PyTorch 2.0 の MPS バックエンドは Beta 段階に入り、grid sampler、triangular solve、topk など最もよく使われる 60 の Torch オペレーターをカバーします。テストカバレッジも大幅に向上し、勾配テストや ModuleInfo テストを含みます。

複数の有名モデルが macOS バックエンドとして MPS を正式採用しています: WhisperAI（音声テキスト変換）、YOLO（物体検出）、Stable Diffusion（画像生成）。デモでは M2 Max 上の YOLOv5 が MPS バックエンドでリアルタイム物体検出を実行し、CPU 版より明らかに高いフレームレートを達成しました。

Metal System Trace で性能ボトルネックを特定

PyTorch nightly 版では MPS オペレータープロファイリングサポートが追加されました。OS signpost でオペレーター実行時間、CPU-GPU データコピー、CPU フォールバック操作を Metal System Trace に記録します。

デモでは 7 層の Sequential モデルで、Softshrink 活性化関数が MPS 非対応のため CPU にフォールバックしています。トレースでは GPU タイムライン上に大きな空白が見えます—GPU が CPU の Softshrink 完了を待っている状態です。解決策はカスタム Metal オペレーターで置き換えることです。

カスタム GPU オペレーターの 4 ステップ

（06:34）

ステップ 1: Objective-C と Metal でオペレーターを実装。get_command_buffer で MPSStream のコマンドバッファを取得し、get_dispatch_queue でシリアルディスパッチキューを取得、キュー内でカスタム kernel をエンコードします。

ステップ 2: pybind11 で Python バインディングを作成。

ステップ 3: torch.utils.cpp_extension.load() で拡張をコンパイル。

ステップ 4: トレーニングスクリプトでインポートして使用。

カスタム Softshrink オペレーターに置き換えると、モデルの実行効率が大幅に向上し、CPU フォールバックとデータコピーがすべて消えます。

自動混合精度トレーニング

MPSGraph に bfloat16 サポートが追加されました。bfloat16 は符号 1 ビット、指数 8 ビット、仮数 7 ビットで、標準 IEEE float16 よりディープラーニングに適しています。PyTorch MPS バックエンドは float16 と bfloat16 の両方の混合精度モードをサポートします。

有効化は簡単です。autocast をコンテキストマネージャーとしてトレーニング領域を囲むと、MPS が各レイヤーに適した精度を自動選択します。畳み込み層と線形層は低精度、還元層は高精度を維持します。PyTorch 2.0 + macOS Sonoma では、MPS バックエンドは前世代比最大 5 倍高速です。

TensorFlow Metal 1.0 安定版

TensorFlow Metal プラグインが 1.0 版をリリースしました。grappler remapping optimizer が追加され、畳み込み+加算、行列乗算、オプティマイザー操作、RNN cell などの一般的な計算パターンを自動識別・融合します。これらの最適化は計算グラフ作成時に自動で発生し、手動介入は不要です。

混合精度はグローバル設定に対応し、tf.keras.mixed_precision.set_global_policy('mixed_float16') の 1 行で有効化できます。インストールも簡素化され、pip で TensorFlow と tensorflow-metal プラグインをインストールするだけで加速が自動有効になります。

JAX が Metal に初登場

JAX は NumPy ベースの高性能数値計算ライブラリで、自動微分（grad）、ベクトル化（vmap）、JIT コンパイルをサポートします。今年 JAX が Metal GPU 加速バックエンドを獲得しました。

M2 Max MacBook Pro 上では、JAX Metal 加速は CPU 比平均 10 倍高速です。インストールと環境設定の詳細は Metal Developer Resources のページを参照してください。

MPSGraphPackage: 起動遅延の解決

複雑な MPSGraph は初回コンパイルに長時間かかり、アプリ起動を遅くする可能性があります。MPSGraphPackage は新しいシリアル化形式で、ビルド時に MPSGraphExecutable を事前コンパイルし、ランタイムでファイルから直接ロードできます。

作成方法: MPSGraphExecutable.serializationDescriptor で記述子を作成し、serialize メソッドでファイルに書き込みます。ロード時: コンパイル記述子とファイルパスで MPSGraphExecutable を初期化します。

MPSGraphTool コマンドラインツールは Core ML の ML Program や ONNX モデルを MPSGraphPackage に直接変換でき、推論コードを手動で書き直す必要がありません。

Int8 量子化でメモリ使用量を削減

MPSGraph に対称・非対称 Int8 量子化 API が追加されました。推論時に浮動小数点の代わりに 8 ビット整数を使うことで、メモリ帯域幅を節約し、モデルのメモリ使用量を削減できます。

フロー: Int8 の活性化と重み入力を dequantizeTensor で浮動小数点に逆量子化し、畳み込みなどの操作に投入、出力を quantizeTensor で Int8 に再量子化します。MPSGraph はこれらの操作を単一 kernel に自動融合し、メモリ移動を削減します。

詳細

PyTorch MPS プロファイリングの有効化

（04:31）

import torch
import torch_mps

# MPS プロファイラを起動
torch.mps.profiler.start(mode="interval", wait_until_completed=False)

# あなたのトレーニングコード
model = MyModel().to("mps")
for batch in dataloader:
    output = model(batch.to("mps"))
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

# プロファイラを停止
torch.mps.profiler.stop()

キーポイント:

• torch.mps.profiler.start() で OS signpost 記録を有効化
• Instruments の Metal System Trace で os_signpost タイムラインを確認
• タイムラインに各オペレーターの実行時間、データ型、コピー長が表示されます
• GPU タイムライン上の空白を探す—通常は CPU フォールバックに対応します

カスタム Metal オペレーターの実装

（07:00）

// softshrink_metal.mm
#import <ATen/mps/MPSStream.h>
#import <ATen/mps/MPSAllocator.h>

void customSoftshrink(const at::Tensor& input, at::Tensor& output, float lambda) {
// MPSStream のコマンドバッファを取得
    id<MTLCommandBuffer> commandBuffer = at::mps::get_command_buffer();
    id<MTLDevice> device = at::mps::GetMPSDevice();
    
// シリアルディスパッチキューを取得
    dispatch_queue_t queue = at::mps::get_dispatch_queue();
    
    dispatch_sync(queue, ^{
        id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
        
// Metal shader を読み込む
        id<MTLLibrary> library = [device newDefaultLibrary];
        id<MTLFunction> function = [library newFunctionWithName:@"softshrink_kernel"];
        id<MTLComputePipelineState> pipeline = [device newComputePipelineStateWithFunction:function error:nil];
        
        [encoder setComputePipelineState:pipeline];
        [encoder setBuffer:input.storage().data() offset:0 atIndex:0];
        [encoder setBuffer:output.storage().data() offset:0 atIndex:1];
        [encoder setBytes:&lambda length:sizeof(float) atIndex:2];
        
        MTLSize gridSize = MTLSizeMake(input.numel(), 1, 1);
        MTLSize threadGroupSize = MTLSizeMake(256, 1, 1);
        [encoder dispatchThreads:gridSize threadsPerThreadgroup:threadGroupSize];
        [encoder endEncoding];
    });
    
// 同期的に待機（シリアライズが必要な場合）
    at::mps::synchronize();
}

キーポイント:

• get_command_buffer() で現在の MPSStream コマンドバッファを取得
• get_dispatch_queue() でシリアルキューを取得し、マルチスレッド送信を直列化
• dispatch_sync 内で kernel をエンコードし、競合状態を回避
• synchronize() で現在のコマンドバッファの完了を待機。直列化不要なら commit() を使用

Python バインディングとコンパイル

（07:58）

# softshrink_binding.cpp
#include <torch/extension.h>
#include <pybind11/pybind11.h>

void customSoftshrink(const at::Tensor& input, at::Tensor& output, float lambda);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("custom_softshrink", &customSoftshrink, "Custom Softshrink kernel");
}

# 拡張をコンパイルして読み込む
from torch.utils.cpp_extension import load

custom_ops = load(
    name="custom_softshrink",
    sources=["softshrink_binding.cpp", "softshrink_metal.mm"],
    extra_compile_args=["-std=c++17"]
)

# トレーニングスクリプトで使用
from custom_ops import custom_softshrink

class CustomSoftshrink(torch.nn.Module):
    def __init__(self, lambd=0.5):
        super().__init__()
        self.lambd = lambd
    
    def forward(self, input):
        output = torch.empty_like(input)
        custom_softshrink(input, output, self.lambd)
        return output

# 元のモデル内の Softshrink を置き換える
model = nn.Sequential(
    nn.Linear(784, 256),
    CustomSoftshrink(),
    nn.Linear(256, 10)
).to("mps")

キーポイント:

• PYBIND11_MODULE で C++ 関数を Python に公開
• torch.utils.cpp_extension.load() で共有ライブラリを自動コンパイル・ロード
• extra_compile_args で追加のコンパイラ引数を渡せます
• 元のオペレーターを置き換えると、モデルは MPS 上で完全に GPU 実行され、CPU フォールバックはなくなります

自動混合精度トレーニング

（10:57）

import torch

model = MyModel().to("mps")
optimizer = torch.optim.Adam(model.parameters())
scaler = torch.cuda.amp.GradScaler()  # MPS にも適用可能

for batch, target in dataloader:
    batch = batch.to("mps")
    target = target.to("mps")
    
# autocast が各層に適した精度を自動選択
    with torch.autocast(device_type="mps", dtype=torch.float16):
        output = model(batch)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

キーポイント:

• torch.autocast(device_type="mps", dtype=torch.float16) で混合精度を有効化
• float16 と bfloat16 の両モードをサポート
• 畳み込み層と線形層は自動で半精度、還元層は全精度を維持
• GradScaler で勾配のアンダーフローを防止

MPSGraphPackage シリアル化

（16:34）

// MPSGraphPackage を作成して保存
MPSGraphExecutable *executable = [graph compileWithDevice:device
                                               feeds:feeds
                                            targetTensors:targets
                                     targetOperations:nil];

MPSGraphExecutableSerializationDescriptor *descriptor = [[MPSGraphExecutableSerializationDescriptor alloc] init];
[descriptor setMinimumDeploymentTarget:@"16.0"];

NSURL *packageURL = [NSURL fileURLWithPath:@"model.mpsgraphpackage"];
[executable serializeToURL:packageURL descriptor:descriptor error:nil];

// 実行時に読み込む
MPSGraphExecutableCompilationDescriptor *compDescriptor = [[MPSGraphExecutableCompilationDescriptor alloc] init];
MPSGraphExecutable *loadedExecutable = [[MPSGraphExecutable alloc] initWithPackageURL:packageURL
                                                                        compilationDescriptor:compDescriptor];

キーポイント:

• serializeToURL:descriptor:error: で事前コンパイル済み executable をファイルに書き込み
• minimumDeploymentTarget でサポートする最低 OS バージョンを指定
• ランタイムでは initWithPackageURL で直接ロードし、コンパイル段階をスキップ
• MPSGraphTool で Core ML や ONNX モデルを直接変換可能

MPSGraphTool コマンドライン変換

# Core ML モデルを MPSGraphPackage に変換
mpsgraphtool --convert coreml --input model.mlpackage \
             --output model.mpsgraphpackage \
             --minimum-deployment-target "17.0"

# ONNX モデルを MPSGraphPackage に変換
mpsgraphtool --convert onnx --input model.onnx \
             --output model.mpsgraphpackage \
             --minimum-deployment-target "17.0"

キーポイント:

• --convert でソースモデル形式を指定: coreml または onnx
• --minimum-deployment-target で最低 OS バージョンを設定
• 変換後の package は App 内で直接ロード・実行できます

重要ポイント

1. Mac 上のローカル Stable Diffusion ジェネレーター

• 何を作るか: Mac 上で Stable Diffusion を実行して画像を生成。完全オフライン
• 価値がある理由: PyTorch 2.0 MPS バックエンドが Stable Diffusion を公式サポート。M2 シリーズチップで実用速度
• 始め方: diffusers ライブラリで Stable Diffusion モデルをロードし、device="mps" を設定、torch.autocast で加速

2. リアルタイムビデオ物体検出 App

• 何を作るか: カメラのリアルタイム映像で物体を検出・ラベル付け
• 価値がある理由: YOLOv5 が MPS バックエンドを公式サポート。リアルタイムに十分なフレームレート
• 始め方: torch.hub.load('ultralytics/yolov5', 'yolov5s') でモデルをロード、.to('mps') で GPU に送り、AVCaptureVideoDataOutput でフレーム取得

3. ニッチモデルのカスタムオペレーター加速

• 何を作るか: 研究モデルに MPS 非対応のオペレーターがある場合、カスタム Metal kernel で置き換え
• 価値がある理由: 1 つの CPU フォールバックオペレーターがモデル全体を遅くする可能性。カスタム kernel で GPU 全速を回復
• 始め方: まず torch.mps.profiler でフォールバックオペレーターを特定し、4 ステップで Objective-C + Metal 実装 + Python バインディング

4. モバイルモデル量子化デプロイ

• 何を作るか: トレーニング済み浮動小数点モデルを Int8 に量子化し、App サイズとランタイムメモリを削減
• 価値がある理由: MPSGraph の Int8 量子化 API が逆量子化-計算-量子化フローを自動融合し、推論が高速化
• 始め方: MPSGraphTool で ONNX モデルを MPSGraphPackage に変換し、graph に quantizeTensor と dequantizeTensor ノードを挿入

関連セッション

• Improve Core ML integration with async prediction — Core ML 非同期予測 API と性能最適化
• Accelerate machine learning with Metal — WWDC22 の Metal ML 加速の基礎
• Build customized ML models with Metal Performance Shaders Graph — MPSGraph でカスタム ML モデルを構築