PyTorch 2.15、分散学習と低精度推論、多様なGPU対応を強化

概要

PyTorch 2.15は、2.x系列が研究向けのフレームワークから、大規模な学習・推論を支える本番基盤へ発展していく流れを引き継ぐリリースです。583人のコントリビューターによる3,117件の変更を含み、単一の演算子だけでなく、分散通信の診断性、メモリ管理、通信と計算の重複、ハードウェア別の最適化を幅広く扱っています。

主な変更点

• 刷新版NCCLバックエンドを標準化：ncclを使うすべてのプロセスグループで新バックエンドが標準になります。条件を満たす環境では、縮小・拡張を使った再構成により、障害後の再初期化を効率化できます。従来版へ戻す場合は、TORCH_DIST_USE_NCCL2=0またはnccl-legacyを利用できます。
• 通信ハングの調査を改善：バックエンドに依存しないフライトレコーダーが追加され、collectiveが停止した際に各rankの事後情報を保存できます。あるrankのタイムアウトをきっかけに、全rankで記録を取得できるため、多数のGPUを使うジョブの原因究明に役立ちます。
• 一方向データ転送APIを追加：torch.distributedに新しいTransport APIが加わり、最初のバックエンドとしてNIXLを採用しました。転送はCUDA streamに投入でき、CUDA Graphにも取り込めます。プロセスグループやcollectiveを必要としないデータ移動の基盤となります。
• FSDP2とパイプライン並列を細かく制御：精度が重要なパラメーターだけをfp32に保ち、他をbf16で学習できます。勾配蓄積全体のCUDA Graph化、勾配削減と計算の重複、受信バッファーの寿命管理なども改善されました。
• ハードウェア対応を拡大：NVIDIA BlackwellではNVFP4デコードを高速化し、Apple SiliconではFP8、grouped_mmのデコード経路、公開MPS stream APIを追加しました。CUDA 13.2を標準wheelとし、CUDA 13.4、ROCm 10.0のwheelも提供されます。AMDとIntel GPU向けの機能も拡張されています。

意味と影響

今回の本質は、PyTorchを「動く」だけでなく、「観測し、調整し、障害から復旧しやすい」基盤へ近づける点にあります。分散学習では通信停止の調査とFSDP2の制御性が運用負担の軽減につながる可能性があります。一方、BlackwellやApple Silicon向けの専用経路は、推論ワークロードをハードウェアごとに最適化する方向を示しています。

一部の新APIはAPI Unstableとされているため、本番導入ではドライバー、アクセラレーター、既存の分散学習構成との互換性を確認する必要があります。

出典：PyTorch Blog