AI・機械学習
上級

AllReduce(オールリデュース)

分散学習における集合通信オペレーションの一つ。各GPU(ワーカー)が保持する勾配テンソルを全ノードで集約(Reduce)し、結果を全ノードに配布(Broadcast)する操作。Ring-AllReduceが最も広く使われる実装。

0 回閲覧
0 いいね

AllReduceとは

AllReduceは、分散学習において各GPU(プロセス)が保持するテンソル(主に勾配)を全プロセスで集約し、集約結果を全プロセスに配布する集合通信オペレーションである。データ並列学習では各GPUが異なるミニバッチで計算した勾配を平均化するためにAllReduceを使用する。

AllReduceの基本動作

AllReduceは論理的には「Reduce + Broadcast」の2ステップで構成される。

  • Reduce: 全プロセスのテンソルを1つのプロセスに集約し、要素ごとの演算(通常は加算)を適用
  • Broadcast: 集約結果を全プロセスに配布
  • AllReduce: 上記2つを1回の操作で効率的に実行。最終的に全プロセスが同一の集約結果を保持する

AllReduceの入力は各プロセスのローカルテンソル、出力は全プロセスで同一の集約テンソルとなる。

主要なAllReduceアルゴリズム

アルゴリズム通信量レイテンシ特徴採用例
Ring-AllReduce2(N-1)/N × M2(N-1)×α帯域最適、大テンソル向きHorovod, NCCL
Tree-AllReduce2×M2×log₂N×αレイテンシ最適、小テンソル向きMPI, Gloo
Recursive Halving-Doubling2×M2×log₂N×αバランス型Intel MPI
Butterfly AllReduceM×log₂Nlog₂N×α全二重リンク活用研究段階
SHARP In-Network2(N-1)/N × Mα+ネットワーク遅延ネットワークスイッチ内で集約Mellanox SHARP

※ N=プロセス数、M=テンソルサイズ、α=レイテンシ

Ring-AllReduceの仕組み

Ring-AllReduceは最も広く使われるアルゴリズムで、N個のGPUをリング状に接続して動作する。

  1. ReduceScatter フェーズ: テンソルをN個のチャンクに分割。各ステップでGPU iがチャンクをGPU (i+1)に送信し、受信したチャンクとローカルチャンクを加算。N-1ステップ後、各GPUが1チャンクの最終結果を保持
  2. AllGather フェーズ: 各ステップで完成チャンクをリング上の次のGPUに転送。N-1ステップ後、全GPUが全チャンクの最終結果を保持

総通信量は 2×(N-1)/N × M であり、プロセス数Nに対してほぼ線形にスケールする(帯域律速時)。

NCCLにおけるAllReduce

NVIDIA NCCL(NVIDIA Collective Communications Library)はGPU向けに最適化された集合通信ライブラリである。

  • Ring アルゴリズム: 大きなテンソル(数MB以上)で帯域効率が最大化
  • Tree アルゴリズム: 小さなテンソル(数KB-数百KB)でレイテンシが最小化
  • NVLink活用: ノード内GPU間では最大900GB/s(H100 NVSwitch)の帯域でAllReduceを実行
  • InfiniBand/RoCE: ノード間では400Gbps NDR InfiniBandまたはRoCE v2で通信
  • CUDA Graph対応: AllReduce操作をCUDA Graphに統合して起動オーバーヘッドを削減
  • FP8 AllReduce: H100以降ではFP8精度でのAllReduceが可能、通信量を半減

勾配圧縮技術

AllReduceの通信ボトルネックを緩和する勾配圧縮技術がある。

  • Top-K スパース化: 勾配の上位K%のみを通信。精度低下を補償するためエラーフィードバックを使用
  • 1-bit Adam: 勾配を1ビットに量子化して通信量を32倍削減。DeepSpeedが実装
  • PowerSGD: 低ランク近似で勾配を圧縮。ランクrで通信量をr/d倍に削減(d=パラメータ次元)
  • FP8量子化: 勾配をFP8に変換して通信。FP16比で通信量半減、精度劣化は最小限

パフォーマンスチューニング

AllReduceの性能を最大化するためのテクニック。

  • バケット化(Bucketing): 複数の小さなテンソルを1つの大きなバケットにまとめてAllReduceを実行。PyTorchのDDPはデフォルト25MBバケット
  • 計算-通信オーバーラップ: 後方伝播の勾配計算と並行して、既に計算済みの層の勾配をAllReduceする。PyTorch DDPのfind_unused_parameters=Falseで最適化
  • ストリーム分離: 計算用CUDAストリームと通信用ストリームを分離して並行実行
  • NCCL環境変数チューニング: NCCL_ALGO=Ring, NCCL_MIN_NCHANNELS=4, NCCL_NET_GDR_LEVEL=5など

よくある質問(FAQ)

Q1: AllReduceとAllGatherの違いは何ですか? A: AllReduceは各プロセスのテンソルを集約(加算・平均など)して全プロセスに配布する。AllGatherは集約せず各プロセスのテンソルをそのまま連結して全プロセスに配布する。ZeRO-3ではAllGatherでパラメータを収集する。

Q2: Ring-AllReduceはなぜ大きなテンソルで効率的なのですか? A: 通信量が 2(N-1)/N × M でプロセス数Nにほぼ依存しないため。Nが増えても1プロセスあたりの通信量はほぼ2Mで一定。ただしN-1ステップのレイテンシは増加するため、小テンソルではTree型が有利。

Q3: AllReduceのボトルネックを特定する方法は? A: NCCL環境変数NCCL_DEBUG=INFOで通信ログを出力し、nsys profileでCUDA/NCCLのタイムラインを可視化する。torch.profilerのtrace exportでAllReduce操作の所要時間を計測できる。

まとめ

  • AllReduceは分散学習の勾配同期における最重要集合通信オペレーション
  • Ring-AllReduceが大テンソルで帯域最適、Tree-AllReduceが小テンソルでレイテンシ最適
  • NCCLがNVIDIA GPU向けの標準実装で、NVLink/InfiniBandを最大活用
  • 勾配圧縮(1-bit Adam, PowerSGD, FP8)で通信ボトルネックを緩和
  • バケット化と計算-通信オーバーラップが実用上の性能チューニングの鍵