11月8日、Joab Jacksonが「PyTorch 2.15 tackles the 'one bad GPU' problem」と題した記事を公開した。この記事では、PyTorch 2.15における分散GPU通信バックエンドの刷新と、1台のGPU障害が訓練ジョブ全体を巻き込む問題への対処について詳しく紹介されている。
NCCLとは何か、なぜ重要か
NCCL（Nvidia Collective Communications Library）は「ニッケル」と読む。PyTorchに対してGPU間通信のプリミティブを提供するライブラリで、単一ノード内および複数ノードにまたがるGPUクラスター間でテンソルや勾配、重みなどをやり取りするために使われる。LLMの訓練や大規模推論における基盤技術だ。標準的なTCPベースのネットワークスタックではGPUクラスターにボトルネックが生じるため、NCCLはCPUオーバーヘッドを抑えた最適化済みの通信パスを使う。
NCCLはNvidiaがSC15で公開し、Metaが2016年にPyTorchを構築した際に採用した。以来、OpenAI・DeepMind・MetaといったメジャーなMLラボがPyTorchの分散訓練の基盤として使い続けており、拡散モデルからLLM、推薦システム、ロボティクスまで幅広い用途を支えている。
「1台の不良GPU」問題とは何か
大規模なモデル訓練を運用したことがあるエンジニアなら、心当たりがあるはずだ。数百台のGPUが協調して動く訓練ジョブの途中で、1台のGPUがクラッシュまたは著しく遅延すると、ジョブ全体がタイムアウトやdeadlockで止まる。そして再起動は最初からやり直し——というシナリオだ。
PyTorch 2.15は、この問題の根本にあるNCCLバックエンドを全面的に書き直したことで、このリスクを大幅に低減しようとしている。
旧バックエンドの何が問題だったか
旧来のProcessGroupNCCLは、1つのプロセスグループが数十のコミュニケーターを生成できる設計だった。これが断続的なデッドロックや予測不能なパフォーマンス低下を招いていた。PyTorchコア開発者でMetaのソフトウェアエンジニアであるTristan Riceは、YouTubeライブ配信の中で「旧インターフェースは技術的負債に溺れていた」と述べた。
新バックエンドの3つの改善点
1. 1プロセスグループ＝1コミュニケーターの厳格化
新バージョンでは、初期化時に1つのプロセスグループに対してコミュニケーターは1つのみというルールが課される。複数コミュニケーター間のクロストークが原因だった複雑な問題をデフォルトで排除できる。並列処理が必要な場合は、プロセスグループを追加で投入する設計になった。
2. 遅延GPUの動的な切り離しと代替割り当て
障害が起きたGPUや著しく遅延したGPUを、サードパーティのオーケストレーターが切り離し、代替ワーカーをアサインできるよう再設定が容易になった。いわゆる「straggler（落伍者）」問題への対処だ。
3. フライトレコーダーの追加
以前はGPUクラスターがスタックしても、ユーザーが受け取るのは汎用的なエラーメッセージだけだった。新バックエンドにはフライトレコーダーが追加され、全オペレーションのスナップショットを連続的に記録する。記録の対象はAllReduce・Broadcast・ReduceScatterといった集合通信オペレーションで、各オペレーションの開始・完了タイムスタンプや参加ランク情報が含まれる。障害発生後にさかのぼって「どのオペレーションのどのランクが詰まったか」を特定できるため、根本原因の調査が格段に効率化される。
移行スケジュールと後方互換性
旧来のProcessGroupNCCLはPyTorch 2.15にロールバック用として同梱される。新バックエンドが期待通りに動作すると確認できれば、2.16（来年初頭リリース予定）で旧バージョンは削除される予定だ。
Riceは、書き直しによってコードベースがクリーンになり、今後の機能追加も容易になると述べた。
その他のアップデート
NCCLバックエンドの刷新以外にも、PyTorch 2.15には以下の変更が含まれる。
- 非同期の一方向データ転送向けにPython Transport APIを新設。トレーナーとジェネレーター間で重みを同期する訓練ジョブに適している
- Apple製品向けの並列処理改善
- NvidiaのBlackwell GPUのサポート追加
今回のリリースは583人のコントリビューターによる3,117コミットで構成されている。
詳細はPyTorch 2.15 tackles the 'one bad GPU' problemを参照していただきたい。
