PyTorch 2.14、コンパイラと分散実行、ハードウェア対応を一段強化

概要

PyTorch 2.14の特徴は、単独の演算子を増やすことだけではない。GPUカーネルの選択、分散通信、障害からの復旧、動的なモデル表現、複数アクセラレーターへの対応を同時に進め、モデルを実際の運用環境で動かすための共通基盤を強化している。PyTorch 2.xが研究向けフレームワークから、大規模な学習と推論を支える統合プラットフォームへ移行する流れが、今回も明確になった。

主な変更点

• NVGEMMがInductorに加わった。 CuTeDSLで生成されたCUTLASSカーネルを利用する新しいGEMMバックエンドで、TritonやATenと並べて自動チューニングできる。epilogue fusion、scaled GEMM、NVFP4 GEMM、grouped-reduction epilogueに対応し、行列演算に対するコンパイラの選択肢を増やしている。
• 分散通信にnccl2を追加した。 torchcommsから移植されたバックエンドで、collectiveの契約全体、ノンブロッキング通信器、即時のcommunicator splittingを実装する。大規模クラスタでは通信の準備や再構成もジョブ全体の効率に関わるため、実行時の柔軟性が重要になる。
• 耐障害性をc10dの概念に引き上げた。 プロセスグループをインプレースで再構成でき、片方向のRMAウィンドウも利用できる。さらにFlight RecorderはNCCLだけでなく任意のバックエンドで動作する。障害対応と観測性が、特定の通信実装だけに依存しにくくなった。
• Apple Siliconの線形代数を強化した。 JacobiカーネルによるSVD、eigh、QR、Choleskyなどのネイティブ経路を追加。MPSGraphから手書きのMetalカーネルへ移行する演算も増え、リダクション、インデックス、畳み込み、活性化などでグラフコンパイルの負担を減らす。単一トークンのデコードでは、F.linearの経路修正とGEMVカーネルも行われた。
• 動的な処理を表現しやすくした。 torch.switchはtorch.condを多方向分岐へ一般化し、torch.while_loopはCUDA Graphに取り込めるようになった。@dynamic_specは変化するテンソル次元を宣言的に記述し、その情報をtorch.compile、torch.export、make_fxで共有する。
• 対応するハードウェアを拡大した。 ROCm 7.14のwheelはTheRock pip SDKから生成され、Intel XPUにはネイティブなグラフキャプチャが追加された。InductorはNVIDIA Rubinのsm_107をターゲットにする。また、実験的な複素テンソル対応では、対応する演算を実部と虚部へ分解して最適化する。

意義と注意点

開発者にとっては、ハードウェアごとの専用コードを書く必要を減らせる可能性がある。より多くのGEMMカーネルを自動選択でき、通信と計算のオーバーラップも標準で有効になるため、モデルコードを大きく変更せずに重要経路を改善しやすい。動的形状や分岐、ループをコンパイルとエクスポートの間で一貫して扱える点も、実用的な利点だ。

大規模学習では、nccl2とc10dの耐障害性が特に重要になる。ノンブロッキング通信器とプロセスグループの再構成により、ノード障害時に常に最初から再実行する必要があるとは限らなくなる。Apple Silicon向けMetal移行も、デスクトップ級GPUをより統一的な最適化体系に取り込む取り組みといえる。

一方、MPSの一部はAPI不安定で、複素テンソルのコンパイル対応は実験段階だ。導入時には対象モデルとアクセラレーターでのベンチマークが欠かせない。それでも2.14は、性能、移植性、信頼性を同じランタイムの中で扱おうとする方向性を示すリリースである。

出典：PyTorch Blog