All-Reduce・All-Gather・All-to-Allとは？LLM分散推論のGPU通信を整理する

LLMを複数GPUへ分けたとき、各GPUは自分の計算だけで次tokenを確定できるとは限らない。Tensor Parallelismではpartial resultを合算し、Context Parallelismでは分割したKVを共有し、Expert Parallelismではtokenを担当expertへ送り直す。

このGPU間の定型的なdata交換をcollective communicationと呼ぶ。All-Reduce、All-Gather、Reduce-Scatter、All-to-Allは名前が似ているが、各rankが最後に持つdataが違う。通信を理解すると、GPU利用率が低い理由や、multi-nodeだけ急に遅くなる原因を切り分けやすい。

rankとcommunicatorを先に理解する

分散処理では、通信へ参加するprocessやGPUをrankとして識別する。4 GPUのcommunicatorならrank 0から3が一つのcollectiveへ参加する。

collectiveは、一部のrankだけが任意に呼ぶ通常のfunctionではない。NCCLの公式資料では、各rankが同じcountとdata typeで対応するcollectiveを呼ぶ必要があり、呼び出しがそろわなければhang、crash、data corruptionを起こしうるとしている。

groupの構成も重要である。8 GPUがあっても、TP group 0から3とTP group 4から7が別々にAll-Reduceする場合がある。world全体のrank番号だけでなく、どのcommunicatorへ属するかを見る。

All-Reduceは全rankの値を集約して全rankへ返す

4 rankが同じ形のarrayを持ち、要素ごとのsumを計算するとする。

rank	入力
0	`[1, 2]`
1	`[10, 20]`
2	`[100, 200]`
3	`[1000, 2000]`

sumのAll-Reduce後は、全rankが[1111, 2222]を持つ。Reduceは集約結果をroot rankだけへ置くが、All-Reduceは全rankへ同じ結果を返す。

LLM推論のTensor Parallelismでは、各GPUがmatrix multiplicationのpartial outputを計算し、その和を全rankで必要とする場面がある。Decodeの各layerでcollectiveが起きれば、少量のtoken生成でも通信latencyが積み重なる。

trainingのgradient同期でよく使われる操作だが、推論でもTPのpartial result集約に使う。用途をtrainingだけに限定しない。

All-Gatherは分割dataを連結して全rankへ配る

rank 0が[A]、rank 1が[B]、rank 2が[C]、rank 3が[D]を持つ場合、All-Gather後は全rankが[A, B, C, D]を持つ。

値を足すのではなく、rank順に集める。各rankの入力がN要素なら、各rankの出力はrank数倍の要素になる。そのため受信bufferと一時memoryが増える。

Tensor Parallelismでshardしたactivationやlogitsを全rankへ復元するとき、Context Parallelismで分割したKVやsequence情報を集めるときに使われる。ただし、実装はfull tensorを一度にmaterializeせず、ringのsend・receiveへ分解して計算と重ねることもある。

Reduce-Scatterは集約後の結果をrankへ分ける

Reduce-Scatterは、各rankの入力を要素ごとにreduceし、その結果を等しいchunkへ分けてrankごとに一部だけ返す。

先ほどのAll-Reduceでは全rankが[1111, 2222]を受け取った。2 rankへReduce-Scatterする単純な例なら、rank 0は集約結果の前半、rank 1は後半を持つ。

最終結果を全rankで複製する必要がない場合、All-Reduceより各rankの出力memoryを減らせる。NCCLの資料では、Reduce-Scatterの後にAll-Gatherを行うとAll-Reduceと同等のdata flowを構成できる。

Tensor ParallelismやSequence Parallelismでは、集約済みactivationをsequence chunkとしてrankへ残すために使う。次の処理がshardをそのまま受け取れるかが選択条件になる。

All-to-Allはrankごとに違う宛先へ違うdataを送る

All-to-Allでは、各rankが送信bufferを宛先rankごとのchunkへ分ける。rank 0からrank 1へ送る内容と、rank 0からrank 2へ送る内容は異なる。

Expert Parallelismでは、routerが選んだexpertのownerへtokenを送る。各rankはlocal batch内のtokenをexpert owner別に並べ替え、All-to-Allでdispatchする。expert計算後は逆方向のAll-to-Allで元のrankへ戻す。

送信量が均等とは限らない。hot expertへtokenが集中すれば、特定宛先のmessageだけ大きくなる。固定countのAll-to-Allと、宛先ごとにcountが異なるAll-to-All-vではbuffer管理と性能が変わる。

Broadcast、Gather、Scatterとの違い

Broadcastは一つのrootから同じdataを全rankへ送る
Gatherは全rankのdataをrootだけへ集める
Scatterはrootが持つdataをrankごとに分けて送る
All-Gatherは集めた結果を全rankが持つ
All-to-Allは全rankが各宛先へ別のchunkを送る

model weightや設定をrootから配るならBroadcast、最終結果を一つのprocessへ集めるならGatherが合う。必要のないrankまでfull dataを持つAll-Gatherを使うとmemoryとbandwidthを浪費する。

collective名と実際のnetwork algorithmは別である

All-Reduceという意味上の操作を、ring、tree、hierarchical algorithmなどで実装できる。message size、rank数、GPU topology、node数によって有利な方式が違う。

ring All-Reduceではdataをchunkへ分け、隣接rankへ回しながらreduceと配布を進める。大きなmessageでbandwidthを使いやすい一方、小さなmessageでは複数stepのlatencyが目立つ。treeはstep数を抑えやすいが、link利用の特徴が変わる。

NCCLはGPU topologyと利用可能なtransportに応じてalgorithmやprotocolを選ぶ。application側はAll-Reduceを呼んでいても、実際にはNVLink、PCIe、shared memory、InfiniBand、RoCE、TCP socketなど経路が異なる。

node内で速くてもmulti-nodeでは別の結果になる

同一nodeのGPU間はNVLinkやNVSwitchで高bandwidth、低latencyに通信できる場合がある。node間ではNICとnetwork fabricを通る。

TPのAll-Reduceをnode境界へまたがせると、Decodeの各stepがnetwork latencyを受ける。EPのAll-to-Allは複数宛先へdataが広がり、network congestionの影響を受けやすい。CPのKV ringは長文Prefillで大きなmessageを流す。

同じGPU数でも、rank mappingが変わるだけで性能が変わる。communicator内のrankがどのGPU、PCIe switch、NUMA node、NICへ対応するかを記録する。

synchronizationが遅いrankを全体へ広げる

collectiveは参加rankがそろって進む。あるrankの前段計算が遅いと、ほかのrankはcollective開始地点で待つ。networkが原因に見えて、実際は一つのGPUの熱throttling、ECC error、CPU threadの遅延、expert load偏りが原因の場合がある。

反対にcollective自体が遅ければ、その後のGPU kernelも開始できない。GPU利用率の谷だけでは区別できないため、rank別timelineで「collectiveへ到着するまで」と「collective実行中」を分ける。

最も遅いrank、message size、collective duration、call siteを関連付ける。

communicationとcomputeを重ねる条件

非同期collectiveを別CUDA streamへ発行し、data dependencyのないkernelと重ねれば待ち時間を隠せる。ただし、同じGPU memory bandwidthやSM resourceを奪い合い、重ねた方が遅くなる場合もある。

次の計算がcollective結果をすぐ必要とするならoverlapの余地は小さい。tensorを小さなbucketへ分けすぎるとlaunchとlatencyが増え、大きすぎると最初の通信開始が遅れる。

overlapの有無はwall-clockで測る。通信時間と計算時間を足した理論値だけでspeedupを見積もらない。

data typeと圧縮は品質だけの問題ではない

通信tensorをFP16、BF16、FP8などへ小さくすればbyte数は減る。だが、quantizeとdequantize、scaleの交換、追加kernelが必要になる。小さなmessageでは変換costが削減時間を上回る場合がある。

partial resultを低precisionでreduceすると数値誤差も変わる。KV Cache量子化やweight量子化とは別に、通信中のdata typeと集約precisionを管理する。

まずmessage sizeとlink利用率を測り、communication-boundであることを確認してから圧縮を検討する。

hangをtimeoutだけで再起動しない

collective hangは、network障害だけでなくcall sequence不一致でも起きる。

一部rankだけ別branchへ入りcollectiveを呼ばない
rankごとにcountやdata typeが違う
communicatorのgroup構成が一致しない
あるrankが前段でOOMやexceptionになった
processが別requestのcollectiveを先に呼ぶ
CUDA stream間のdependencyが循環する

timeout後に全rankを再起動するだけでは再発する。各rankの最後に入ったcollective、sequence number、shape、dtype、communicator ID、前段errorを残す。

途中rankだけを復帰させても、ほかのrankのcommunicator状態と一致しない場合がある。model replicaまたはparallel group単位で隔離し、新規requestを健全なgroupへ送る。

監視するmetric

collective種別、call count、message byte数
communicatorと参加rank
call siteごとの平均、p95、p99 duration
rankがcollectiveへ到着した時刻の差
bus bandwidthとalgorithm bandwidth
node内・node間の送受信byte数
NIC throughput、retransmit、drop、error
NVLink、PCIe throughputとerror
computeとのoverlap時間
collective timeout、abort、rank再起動
TTFT、ITL、token throughputへの寄与

すべてのcollectiveを一つの平均へまとめない。TPの小さいAll-ReduceとCPの大きいKV通信では調整方法が違う。

負荷試験

All-Reduce、All-Gather、Reduce-Scatter、All-to-Allを単体benchmarkする
message sizeを小から大へ変える
rank数を2、4、8と増やす
node内とnode間を分ける
TP、CP、EPを実workloadで個別に有効化する
同じlink上で複数communicatorを同時実行する
一つのrankだけ計算を遅らせ、到着差を観測する
NIC、link、rankを一つずつ停止する
countやcall順序の不一致をtest環境で注入する
overlap有無をend-to-end latencyで比較する

microbenchmarkのpeak bandwidthだけで構成を決めない。実際のmessage size、call頻度、computeとの競合、tail latencyを再現する。

collectiveを理解するとparallelismのcostが見える

All-Reduceは集約結果を全rankへ返し、All-Gatherは分割dataを全rankへ連結し、Reduce-Scatterは集約結果を分割したまま残す。All-to-Allはrankごとに異なるdataを交換する。

これらはTensor、Context、Expert Parallelismを動かす通信部品である。GPUを増やしても、collectiveの頻度、message size、topology、遅いrankへの同期が増えれば速くならない。modelの計算時間と同じ粒度でcollectiveを観測する必要がある。