ホームへ戻る

AI 基礎設施

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

新版では、`nccl2`、Process Groupのインプレース再構成、バックエンド横断のFlight Recorderにより、大規模分散学習を強化。InductorにはNVGEMM、宣言的な動的形状、複素数テンソルのコンパイルも加わったが、多くのインターフェースは依然として不安定とされている。

Picture: M-L. Text: Pöllö. Please notice the original sources referred in some parts of the text. · CC BY-SA 2.5 · Image source
zh-Hant

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は削除された。また、公式には、あらゆるハードウェアとワークロードを包括する単一のエンドツーエンド高速化指標は提示されていない。エンジニアリングチームは、自社のモデル、トポロジー、精度設定を用いて改めてベンチマークを行うべきだ。

出典

  1. PyTorch 2.14 Release Blog
  2. PyTorch 2.14.0 Release