「PyTorch 2.14」リリース ─ NVGEMMと分散学習の耐障害性を強化

　PyTorch Foundationは9月2日(現地時間)、オープンソースの機械学習フレームワーク「PyTorch 2.14.0」をリリースした。

　PyTorchは、ニューラルネットワークの構築や学習、推論などに利用される機械学習フレームワーク。Pythonを中心としたAPIを備え、CPUのほか、NVIDIA、AMD、Intel、Appleの各種アクセラレータに対応する。

　PyTorch 2.14では、GPU上の行列演算、分散学習、コンパイル処理、動的な制御フローなどを強化した。PyTorch 2.13以降、487人の貢献者による2995件のコミットが取り込まれている。

　PyTorch 2.14のハイライトは次の通り。
〇Inductor向けNVGEMMバックエンドでエピローグ融合やNVFP4などの低精度演算に対応
〇PyTorch Distributedに新しい「nccl2」バックエンドを追加
〇プロセスグループの再構成や片方向RMAを含む分散処理の耐障害機能を強化
〇NCCL以外のバックエンドにも対応するFlight Recorderを導入
〇複数分岐を表現する高階演算子「torch.switch」を追加
〇「torch.while_loop」のCUDA Graph取り込みに対応
〇動的なテンソル形状を宣言する「@dynamic_spec」を追加
〇複素数テンソルに対する「torch.compile」を実験的にサポート
〇Apple SiliconでSVD、QR、Choleskyなどの線形代数処理をネイティブ化
〇AMD ROCm 7.14、Intel XPUのネイティブグラフ取り込み、NVIDIA Rubinをサポート
〇Python 3.15とFree-Threaded版のPython 3.15t向けバイナリを提供

　NVGEMMは、CuTeDSLで生成したCUTLASSカーネルをInductorから利用する行列演算バックエンド。TritonやATenの候補と比較して適切なカーネルを選択し、バイアス加算や活性化関数など後続処理との融合にも対応する。

　分散学習では、通信と計算を重ね合わせるInductorの最適化が標準で有効になった。新しいnccl2バックエンドやプロセスグループの再構成機能により、大規模な学習処理で障害が発生した際の復旧や状態維持を行いやすくしている。

　Python 3.15およびGILを無効化したPython 3.15t向けwheelは、Linuxのx86-64／AArch64、Windowsのx86-64、Apple Silicon搭載macOS向けに提供される。ただし、Python 3.15上の「torch.compile」は未対応で、今回はEagerモードのみを利用できる。

　同時に公開された「TorchVision 0.29.0」はPyTorch 2.14に対するABI安定性を導入し、将来のPyTorch 2.15以降とも互換性を維持する方針となった。一方、「torchvision.io」の画像エンコーダーとデコーダーは非推奨となり、TorchCodec 0.16以降への移行が案内されている。

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

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

[関連リンク]
リリースアナウンス
リリースノート(GitHub)
リリースノート(TorchVision 0.29.0)