PyTorch 2.15がMPSでFP8に対応

Hugging FaceのPyTorchチームは、AppleのMPSバックエンドをモデルの実行と学習の第一級の場にするための意図的な取り組みを説明している。2026年初頭には、MPSで小さなモデルを動かすだけでも未対応の演算や性能の急落、予想外のメモリ増加を回避する作業が必要になることが多かったという。チームの手法は、言語・画像・動画にわたる実際のモデルをHugging Face Hubから動かし、MPSが不足する点を見つけてPyTorch本体で修正するというもので、MPSGraphを専用のMetalカーネルに置き換える方向が強まっている。

この方針は2.15のリリース全体に表れている。FP8はスケール付き行列積とともにMPSに登場し、別途の取り組みは言語モデルのデコード段階を対象とする。線形層には専用の行列ベクトルカーネルが導入され、大規模語彙のsoftmaxは専用のMetalカーネルで動く。チームによると、[1, 128256]のロジット行に対するFP32 softmaxは8倍高速になった。言語以外では、定数パディングのMetal移行によりタイル型動画VAEデコードのドライバメモリの暴走が解消され、Qwen-Image-2.1のエンコード・デコード往復はPSNR 29.95 dBとなり、10.06 dBからCPUと同等になった。

投稿は今後の予定にも触れている。PyTorch 2.16はすでに進行中で、scaled dot-product attentionがMPSGraphを完全に廃してすべての形状をMetalカーネルで動かすほか、MPSGraphからMetalへの移行がさらに進んでいる。チームはまた、グラフブレークをなくしてより多くのモデルがエンドツーエンドでコンパイルできるよう、torch.compileの改善に取り組んでいるとしている。