「PyTorch 2.13」リリース ─ FlexAttentionのApple Silicon対応など

　PyTorch Foundationは7月8日（現地時間）、機械学習フレームワーク「PyTorch 2.13」をリリースした。

　「PyTorch」は、Python向けのオープンソース機械学習フレームワーク。深層学習モデルの研究開発から本番環境での学習・推論などに利用されている。

　「PyTorch 2.13」では、FlexAttentionのApple Silicon対応、Inductor向けCuTeDSLバックエンド、メモリ使用量を削減するnn.LinearCrossEntropyLoss、新しい分散通信バックエンドtorchcommsなど、多数の機能追加と性能改善が行われている。

　「PyTorch 2.13」のハイライトは以下の通り。
〇FlexAttentionがApple SiliconのMetal／MPSに対応
〇CUDA版FlexAttentionで再現可能な勾配計算を行うdeterministic backward pathを追加
〇InductorにCuTeDSLベースの新しいGPUコード生成バックエンドを追加
〇大規模語彙モデル学習向けにnn.LinearCrossEntropyLossを追加
〇safetensorsファイルをtorch.load()で直接読み込み可能に
〇分散学習向けの新しい通信バックエンドtorchcommsを追加
〇FSDP2でreduce-scatterとall-gatherの通信重畳に対応
〇Linux向けにPython 3.15およびfree-threaded 3.15t wheelを提供
〇ROCm、Arm、Intel XPUなどのプラットフォーム対応を強化
〇ExecuTorchをPyTorch Coreへ統合
〇named tensor機能を削除
など。

　「PyTorch 2.13」では、Apple SiliconのMetal／MPS上でFlexAttentionを利用できるようになった。長いシーケンスに対する疎なattentionパターンでは、SDPAと比較して大幅な高速化が得られる可能性がある。

　CUDA版FlexAttentionでは、deterministic backward pathが追加された。新しい経路では、書き込み順序を事前に計算することで、再現可能な勾配計算を行える。

　ニューラルネットワーク関連では、nn.LinearCrossEntropyLossが追加された。大規模語彙を持つ言語モデルの学習では、最終線形層の出力であるlogits行列が非常に大きくなる。nn.LinearCrossEntropyLossは、線形射影とクロスエントロピー損失の計算を融合し、語彙方向を分割して処理することで、完全なlogits行列をメモリ上に展開せずに損失を計算できる。

　モデル重みの読み込みでは、safetensorsファイルをtorch.load()で直接読み込めるようになった。これにより、safetensors形式の重みを扱う際、別途ライブラリを明示的に読み込む必要がなくなる。

　コンパイル関連では、InductorがCuTeDSLを新しいGPUコード生成バックエンドとして利用できるようになった。CuTeDSLはNVIDIAのCuTeを基盤とするPythonネイティブのDSLで、GEMMやRMSNormなど、Transformer学習で重要な処理に対してTritonとは別の高性能なコード生成経路を提供する。

　分散学習では、新しい通信バックエンドtorchcommsが追加された。torchcommsは、大規模クラスタにおける耐障害性、スケーラビリティ、デバッグ性を改善することを目的としており、構造化ログやcollective tracingなどを備える。

　FSDP2では、reduce-scatter用の専用プロセスグループを利用できるようになった。有効にすると、reduce-scatterとall-gatherの通信を重畳でき、Fully Sharded Data Parallelによる分散学習のスループット向上が期待できる。

　Python対応では、Linux向けにPython 3.15 wheelの提供が開始された。ただし、Python 3.15は現時点でプレリリースであり、対応はtorch wheelのみ、Linuxのx86_64およびaarch64向けに限定される。

　後方互換性に関する変更として、非推奨だったnamed tensor機能が削除された。Tensor.namesやTensor.rename()、Tensor.refine_names()などの関連APIは利用できなくなる。また、Bazelビルドサポートも削除され、PyTorchヘッダを利用する場合はC++20が必要となった。

　「PyTorch 2.13」は、Webサイトから入手できる。

(川原 龍人/びぎねっと)

[関連リンク]
リリースアナウンス(Blogによる記事)
リリースノート
PyTorch Foundation