PyTorch 2.14、障害時の再構成をc10dに統合、Inductorの低精度・マルチハードウェアカーネルを拡充

PyTorch 2.14の主眼は、単一モデルの性能値ではなく、コンパイラ、分散通信、ハードウェアバックエンドを実運用に向けて前進させることにある。新版では、従来のtorchcommsの取り組みをnccl2 c10dバックエンドとして統合し、ノンブロッキングcommunicator、メモリの一時停止／再開、片方向RMA windowを実装した。さらに、BackendとProcessGroupにインプレース再構成インターフェースが追加され、ノード障害後にジョブ全体を必ずしも破棄せず、Process Groupを再構築できるようになった。従来はNCCLに依存していたFlight Recorderも、Process Groupのhookを介してデータを収集する方式に変更され、Glooやカスタムバックエンドでも同種のcollectiveトレースデータを利用できる。

コンパイル経路では、NVGEMMにより、CuTeDSL／CUTLASSが生成するカーネル候補がTritonやATenとともにmm、addmm、scaled_mmの自動チューニングに参加するようになった。bias、要素単位演算、一部のreductionも融合できる。NVFP4経路にはBlackwellが必要となる。ShapesSpecと@dynamic_specを使えば、同一の動的次元制約をtorch.compile、torch.export、make_fxで共有でき、3つのエントリーポイントごとに形状を個別記述する際の差異を減らせる。torch.switchではトレース可能な多分岐制御フローが追加され、複素数テンソルも実部と虚部に分割する方式で、実験的にtorch.compileへ取り込めるようになった。

ハードウェア対応もROCm 7.14、Intel XPUのgraph capture、Rubinのsm_107へ拡大した。Apple SiliconではネイティブのSVD、eigh、QR、Choleskyが追加され、より多くの演算子がMPSGraphからMetalへ移行した。ただし、アップグレード前には注意が必要だ。nccl2は現時点でeagerのみをサポートし、多くの新APIにはunstableの表示がある。Python 3.15向けwheelは提供されるものの、torch.compileは使用できない。TVM Relayバックエンド、一部のprofilerインターフェース、ROCm 7.1は削除された。また、公式には、あらゆるハードウェアとワークロードを包括する単一のエンドツーエンド高速化指標は提示されていない。エンジニアリングチームは、自社のモデル、トポロジー、精度設定を用いて改めてベンチマークを行うべきだ。