AI計算の通信を速める新技術

GPUの通信オーバーヘッドは、実際のAIワークロードにおいて測定可能なボトルネックとなっています。mKernelプロジェクトが引用したデータによると、通信はフォワードパスの43.6%、エンドツーエンドのトレーニング時間の32%を消費することがあります。一般的なMixture-of-Experts（MoE）モデルでは、デバイス間の通信が総実行時間の最大47%を占めることがあります。UC BerkeleyのUCCLプロジェクトの研究者たちは、ノード内NVLink通信、ノード間RDMA、および計算を単一のCUDAカーネルに融合する永続的なCUDAカーネルのライブラリであるmKernelをリリースしました。 従来のマルチGPU通信の標準モデルはホスト駆動型です。CPUが制御パスを実行し、NCCLやNVSHMEMのようなライブラリを呼び出します。ライブラリは、AllReduceやAllGatherなどの集合演算をGPU間で発行します。計算と通信は別々のCUDAストリームで実行され、カーネル境界でオーバーラップします。研究チームは、このアプローチに2つの問題があることを指摘しています。 (1) CPUがGPUの計算能力に追いついていません。GB300 NVL72ラックは、72基のBlackwell Ultra GPUと36基のGrace CPUを統合し、FP8/FP6で720 PFLOP/s、FP4 Tensor Coreで1.44 EFLOP/sの性能、および130 TB/sの全対全ノード内NVLink帯域幅を提供します。これらの速度では、cudaLaunchKernelの呼び出し、CPU側の「すべての書き込みが完了した」チェック、ストリーム間イベントといったマイクロ秒スケールのホストオーケストレーションのオーバーヘッドが、パイプラインのバブルとして直接現れます。 (2) ホスト駆動型システムでは、計算と通信が粗いカーネル境界でオーバーラップします。タイルまたはチャンクレベルでのより細かい粒度のオーバーラップは、ホスト側からは不可能です。 代替案はGPU駆動型通信です。GPU自体が転送をトリガーし、通信が計算と同じカーネルに融合されます。既存のほとんどの融合カーネルライブラリは、単一ノード内または単一GPU内で動作します。mKernelはマルチノードケースを対象としています。 mKernelは永続的なCUDAカーネルのライブラリです。各カーネルはノード内NVLink通信、ノード間RDMA、および