PyTorch 2.15 MPS リリースノート

概要

PyTorch 2.15 は、ネイティブな FP8 サポート、最適化された LLM デコードパス、および専用の Metal カーネルへの移行を導入することで、Apple Silicon ユーザーのパフォーマンスとメモリ効率を大幅に向上させます。これらのアップデートにより MPSGraph への依存が軽減され、密行列モデルや Mixture-of-Experts (MoE) モデルの推論が高速化されるほか、動画生成のような複雑なワークロードにおけるメモリ使用量がより安定します。

主な技術的改善

FP8 サポート

PyTorch 2.15 では、MPS バックエンドで torch.float8_e4m3fn テンソルがサポートされました。これにより、開発者は要素あたり 1 バイトのみを使用して値を格納できるようになり、FP16 や BF16 と比較してメモリフットプリントを実質的に半分に削減できます。今回のリリースには、GPU 上での直接的な FP8 行列乗算を行うための F.scaled_mm が含まれており、テンソル単位および行単位のスケーリングの両方をサポートしています。

LLM デコードの最適化

MPSGraph をバイパスしてネイティブな Metal カーネルを採用することで、密行列モデルと MoE モデルの両方でデコードパフォーマンスが大幅に向上しました。

  • 線形層: F.linear が専用の行列ベクトルカーネルを利用するようになり、一般的なデコード形状全体で 1.23 倍の幾何平均速度向上を実現しました。
  • MoE モデル: 新しいネイティブなグループ化行列乗算の実装により、トークンのルーティングが高速化されました。M3 Ultra で torch.compile を使用した場合、Qwen3.6-35B-A3B のような MoE モデルのスループットは 9.4 トークン/秒から 61.3 トークン/秒に向上しました。
  • Softmax: 専用の Metal カーネルにより、従来の MPSGraph 実装と比較して softmax および log_softmax 操作のレイテンシが約 8 倍削減されました。

動画ワークロードにおけるメモリ管理

MPSGraph ベースのパディングを専用の Metal カーネルに置き換えることで、PyTorch 2.15 は動画処理における重大なメモリオーバーヘッドの問題を解決しました。例えば、Wan2.2 のタイル化 VAE デコードプロセスでは、以前のバージョンで 270 GiB 必要だったメモリが最大 8 GiB にまで削減されました。

線形代数と演算

線形代数の機能が拡張され、以下が含まれるようになりました。

  • 複素数値の LU 分解、LU ソルブ、および三角行列ソルブ。
  • 小行列反転（最大 4x4）のための最適化されたカーネルと、SVD 操作の堅牢性の向上。
  • conv1d、定数/反射/複製パディング、およびアップサンプリングの逆伝播のための新しい Metal ベースのパス。

開発者の制御と要件

開発者は torch.mps.Stream および torch.mps.stream コンテキストマネージャーを通じて実行をより細かく制御できるようになり、明示的なストリーム管理と同期が可能になりました。

システム要件:

  • macOS: 15 以降。
  • Python: 3.11 以降。

今後の展望: PyTorch 2.16

2.16 リリースの開発では、MPSGraph への依存をさらに減らすことに重点を置いています。計画されている改善点は以下の通りです。

  • フル Metal アテンション: スケーリングされたドット積アテンションを完全に Metal カーネルへ移行します。
  • コマンドバッファの最適化: デコード中の GPU/CPU 並列実行を可能にするため、16 カーネルごとにコマンドバッファをコミットします。
  • バイナリサイズの削減: libtorch_cpu.dylib のサイズを削減するため、デフォルトで Metal 3.1 用のバンドルシェーダーをコンパイルします。

Sources