技術的サマリー：マルチノード LLM 推論における通信パフォーマンスの理解と改善

問題定義

大規模言語モデル（LLM）の規模が拡大するにつれ、そのメモリフットプリントは単一 GPU の容量を超え、複数のノードにまたがる分散推論が必要となります。テンソル並列化（TP）やパイプライン並列化（PP）といったモデル並列化戦略は単一ノード推論においてよく研究されていますが、マルチノード環境におけるそのパフォーマンスは十分に探求されていません。マルチノード推論は、ノード間接続（例：NVLink）と比較してノード間遅延が大きいことから、重大な課題をもたらします。

著者らは、TP がメモリ束縛型でデコード負荷の重いワークロードにおいてしばしば優れている一方で、頻繁な all-reduce 操作に起因する大幅な通信オーバーヘッドに悩まされていると特定しました。一方、TP と PP を組み合わせたハイブリッド並列化（HP）は通信オーバーヘッドを削減しますが、行列乗算の性質上、デコード負荷の重い領域では計算時間の効果的な削減に失敗します。特定された重要なボトルネックは、LLM のデコードフェーズに典型的な小規模メッセージサイズ（128 KB から 2 MB）における標準的な NCCL all-reduce 実装の非最適性能であり、特にノード間でのスケーリング時に顕著です。

手法

著者らは、既存の最先端推論エンジンとカスタム研究プロトタイプの組み合わせを用いて、体系的なパフォーマンス調査を実施しました。

1. 実験設定:

   • ハードウェア: 実験は Perlmutter システム（NVIDIA A100 GPU、HPE Slingshot-11 相互接続）および Vista システム（NVIDIA GH200 GPU、InfiniBand 相互接続）上で実施されました。
   • モデル: bf16 精度の Llama 3.1 70B および 405B（Instruct）。
   • ワークロード: 2 つの主要な領域がテストされました。「プレフィル負荷型」（大規模なプロンプト長）と「デコード負荷型」（長い生成シーケンス）です。
   • エンジン: 本研究では、vLLMおよびSGLangに加え、マルチノード HPC 環境での制御された実験を容易にするために著者らによって開発されたカスタムオープンソース推論エンジンYALISが利用されました。YALIS は、統合されたモデル定義レイヤー、Torch Compile ベースの実行、および AxoNN ベースのテンソル並列化を特徴としています。

2. 分析アプローチ:

   • 著者らは、GPU 数を単一ノードから 32 ノード構成まで増加させた際の強スケーリング実験を行い、エンドツーエンドのバッチ遅延を測定しました。
   • 実行時間を計算（Matmul、Other Comp.）、通信、アイドル時間に分解し、ボトルネックを特定しました。
   • 合成 GEMM ベンチマークを用いて、マイクロバッチング（PP）とテンソルスプリッティング（TP）が行列乗算カーネルに及ぼす影響を分離しました。

3. 提案ソリューション（NVRAR）:

   • 通信ボトルネックに対処するため、著者らは NVSHMEM を用いた再帰的倍増に基づく階層的 all-reduce アルゴリズムNVRARを開発しました。
   • 設計: NVRAR は 3 つのフェーズで動作します。
     1. ノード内 Reduce-Scatter: NCCL を用いてノード内のデータを削減します。
     2. ノード間再帰的倍増: NVSHMEM のノンブロッキング RMA プリミティブ（put_nbi）を用いて、再帰的倍増トポロジーでノード間での all-reduce を実行します。
     3. ノード内 All-Gather: NCCL を用いて削減されたデータをノード内の全 GPU に集めます。
   • 最適化:
     • チャンク化ノンブロッキング通信: データをチャンク単位で処理することで計算と通信をオーバーラップさせます。
     • 融合ペイロード: Slingshot/InfiniBand 上の高価な明示的シグナリングプリミティブを回避するため、データと同期フラグを単一の 8 バイトペイロードに結合します。
     • シーケンス番号同期: グローバルバリア（fence/quiet）の遅延オーバーヘッドを回避しつつ正しさを保証するため、グローバル同期に一意のシーケンス番号を使用します。

主要な貢献

1. 体系的なパフォーマンス調査: 本論文は、マルチノード環境における TP と HP のスケーリング挙動の詳細な分解を提供し、TP が一般的にデコード負荷の重いワークロードで優れているが通信オーバーヘッドによって制限される一方、HP はこれらの領域において計算時間の削減に苦戦することを特定しました。
2. YALIS エンジン: マルチノード HPC 環境での計測可能性と制御された実験のために設計された、研究志向の推論エンジン YALIS の開発。
3. NVRAR アルゴリズム: LLM デコードフェーズに特徴的な小メッセージ領域向けに最適化されたカスタム all-reduce アルゴリズム NVRAR の提案と実装。これは NVSHMEM を活用し、Slingshot および InfiniBand ネットワーク上で標準的な NCCL 実装よりも低い遅延を実現します。
4. 統合と評価: NVRAR を YALIS および vLLM に統合し、大規模モデル（Llama 3.1 405B）および混合専門家（MoE）モデル（Qwen3-235B）におけるエンドツーエンドのパフォーマンス向上を実証しました。

結果

• スケーリング挙動: TP と HP の両方とも、マルチノード環境では強スケーリングが不良です。HP は計算束縛型（プレフィル負荷型）の領域で TP よりも優位ですが、TP はメモリ束縛型（デコード負荷型）の領域で HP よりも優位です。
• 通信ボトルネック: NCCL all-reduce は、小規模メッセージ（128 KB–2 MB）においてノード間でスケーリングが不良であり、これらの特定の領域では MPI ベースのソリューションよりも性能が劣ることが多いです。
• NVRAR のパフォーマンス:
  • マイクロベンチマーク: NVRAR は、HPE Slingshot および InfiniBand 相互接続上で 128 KB から 2 MB のメッセージサイズにおいて、NCCL よりも1.9 倍–3.6 倍低い遅延を達成します。
  • エンドツーエンド推論: YALIS および vLLM に統合された NVRAR は、テンソル並列化を使用したデコード負荷型ワークロードにおける Llama 3.1 405B モデルのエンドツーエンドバッチ遅延を最大1.72 倍削減します。
  • MoE モデル: NVRAR は、並列化戦略の TP 成分を高速化することで MoE モデル（例：Qwen3-235B）のスループットも向上させ、NCCL ベースの構成と比較して最大1.14 倍の高いスループットを達成します。

重要性

本論文は、LLM がより大規模化し、推論がマルチノードクラスターへ移行するにつれて、通信効率がパフォーマンスの主要な決定要因となると主張しています。著者らは、標準的な通信ライブラリ（NCCL）が LLM デコードフェーズ特有のメッセージサイズと遅延制約に対して最適化されていないことを実証しました。NVRAR を導入することで、彼らは通信オーバーヘッドを大幅に削減するターゲット型のソリューションを提供し、大規模モデルに対するテンソル並列化のより効率的なスケーリングを可能にしました。この研究は、次世代分散推論システムにおけるカスタム通信プリミティブの必要性を浮き彫りにし、特定の HPC 相互接続環境において標準的な集合通信ライブラリよりも階層的で NVSHMEM ベースのアプローチの有効性を検証しています。