PyTorchコミュニティがNCCL2メモリプールの修正を提案、後続の割り当てで対称通信カーネルが選ばれない可能性

PyTorchは10月13日、分散通信の回帰に関する報告を受けた。NCCL2バックエンドで、空の対称メモリプールを先に登録し、その後プールからテンソルを割り当てると、想定された対称通信経路が有効にならない場合があるという。報告された環境はGB200 2基、CUDA 13.0、NCCL 2.30.7で、開発版とnightlyの両方が対象となっている。報告者はこの問題を、2.15でのデフォルトバックエンド切り替えに関連付けているが、リリース版への影響は確認が必要だ。問題報告

問題は、メモリ領域の登録ライフサイクルにある。register_mem_pool(pool, symm=True)は呼び出し時点ですでに存在する領域にのみ対称ウィンドウを作成する。その後に追加された領域は通常の登録だけが行われ、対称属性は引き継がれない。報告に含まれるプロファイラーの結果では、そのためAllReduceがringカーネルを使用していた。また、zero-CTAを設定したAllGatherもCopy Engineを利用していなかった。再現手順と分析

この問題は、通信と計算のオーバーラップに関するエンジニアの判断に影響する。NVIDIAのドキュメントによると、対称ウィンドウでは各rankが同じバッファオフセットを使う必要があり、これはzero-CTA最適化の前提条件でもある。ハードウェア、ドライバー、集合演算の要件を満たす場合、Copy EngineはCTAを消費せずに通信を実行できる。このことから、登録漏れにより当初想定していたリソース配分が機能しなくなる可能性はあるが、現時点の情報から固定的な学習速度の低下を算出することはできない。NCCL技術ドキュメント

修正案では、プールごとのsymmフラグを保持し、新しい領域の割り当て時にウィンドウ登録を追加で行う。提案者による局所的な検証では、追加された4つのカーネル選択テストが通過した。修正を無効にした場合、NCCL2のテストは失敗した一方、旧バックエンドのテストは引き続き通過した。レビューでは、対称登録に失敗した際に暗黙にフォールバックさせず、エラーを直接報告することも求められた。ただし、ウィンドウ登録は集合操作であるため、各rankは引き続き同じ順序で領域を割り当てる必要がある。修正案とテスト

確認時点で提案はまだオープンであり、リリース済みの修正とは見なせない。対称メモリプールを使うチームは、登録前後でカーネル経路を比較し、報告で挙げられた旧バックエンドへの一時的な切り替えをテスト環境で評価できる。今後はマージとリリースへの反映に加え、実際のワークロードにおけるレイテンシ、帯域幅、オーバーラップへの影響を確認する必要がある。現時点の証拠が示しているのは通信経路の回帰であり、テンソル計算結果の誤りではない。提案の状態