Tech カテゴリのイメージ

LLM推論のTensor・Pipeline・Data Parallelismとは?GPUの増やし方を分けて考える

LLM推論でGPUを増やす目的は一つではない。modelが1台へ収まらないために分割する場合と、同じmodelのrequest処理能力を増やす場合では、選ぶparallelismが違う。

Tensor Parallelismは一つのlayer内の行列を分け、Pipeline Parallelismはlayer群を段階へ分け、Data Parallelismはmodel全体を複製してrequestを分ける。GPU数だけを指定しても、どこを分割したかで通信量、KV Cache容量、latency、障害時の影響が変わる。

三つのparallelismは分割する軸が違う

方法分割するもの各GPUが持つもの主な目的
Tensor Parallelism一つのlayer内のtensor計算各layerのweightの一部大きなlayerを複数GPUで計算
Pipeline Parallelismmodelのlayer列担当する連続layer群深いmodelを段階へ分割
Data Parallelismrequestやbatchmodel全体のreplicarequest throughputを増やす

単純化すれば、必要なGPU数は次の積で表せる。

[ N_{GPU} = TP \times PP \times DP ]

TP 4、PP 2、DP 3なら、4 GPUで一つのpipeline stageを構成し、それを2段つないだmodel replicaを3組作るため24 GPUになる。Context ParallelismやExpert Parallelismも併用する構成では、さらに別の軸が加わる。

Tensor Parallelismはlayerの中で毎回協調する

TransformerのAttentionやMLPには大きなmatrix multiplicationがある。Tensor Parallelismはweight matrixを列方向や行方向へ分け、各GPUが部分結果を計算する。Megatron-LMはTransformer layerへ少数のcollective communicationを挿入するintra-layer model parallelismを示した。

各GPUはmodelのすべてのlayerに参加するが、各layerのweightは一部だけ持つ。1 tokenを進めるたびにGPU間でpartial resultを集約するため、All-ReduceやAll-Gatherなどのcollectiveが頻繁に発生する。

このためTP groupは、NVLinkやNVSwitchで高速接続された同一node内に置くのが一般的である。nodeをまたいでTPを広げると、Decodeの各stepでnetwork latencyを受けやすい。modelが動いたからといって効率よく動いているとは限らない。

TP sizeを増やせば1 GPU当たりのweightは減るが、通信参加者が増え、各GPUへ割り当てる計算が小さくなる。小さなmodelへ過剰なTPを使うと、計算よりcollectiveの比率が大きくなる。

Pipeline Parallelismはlayerを前から後ろへ渡す

Pipeline Parallelismは、たとえば80 layerを前半40 layerと後半40 layerに分ける。requestのhidden stateはstage 0で前半を通り、stage 1へ送られて後半を通る。

TPほどlayerごとのcollectiveは必要ないが、stage境界でactivationを転送する。modelをnode単位で分けやすく、vLLMの公式資料もmulti-node構成の一般例として、node内GPU数をTP、node数をPPに対応させている。

問題はstageの待ち時間である。一つのrequestだけなら、後段は前段の出力を待ち、前段は次の処理がなければ空く。複数requestやmicrobatchを流せばpipelineを埋められるが、online servingではrequest長と到着時刻がそろわない。

layer数や計算量がstage間で偏ると、最も遅いstageが全体を決める。均等なlayer数でも、embedding、output head、MoE layer、cross attentionがあると計算量やmemoryは均等にならない。

Data Parallelismはmodelを複製してrequestを分ける

Data Parallelismでは、各replicaがmodel全体を持ち、独立したrequest batchを処理する。1 replicaでmodelが収まるなら、request throughputを増やす最も分かりやすい方法である。

replica間で1 tokenごとのmodel計算を同期する必要がないため、TPやPPよりcommunicationは少ない。ただしweightを複製するので、modelを1 GPUへ載せられない問題は解決しない。その場合は、TPやPPで作ったmodel replicaをDPで複製する。

各DP replicaは独立したKV Cacheを持つ。同じprefixのrequestを別replicaへ送れば、片方のCacheをそのまま使えない。routerはqueue長だけでなく、Prefix Cacheのlocalityも考慮する必要がある。

一つのrequestを速くする方法と、全体を増やす方法を混同しない

TPやPPはmodelを分割して一つのrequestを複数GPUで進める。DPはrequest群を複数replicaへ分ける。

modelが1 GPUへ収まり、単一requestのlatencyも目標内なら、まずDPを増やす方が通信を抑えやすい。modelが収まらなければTPかPPが必要になる。単一requestの計算量を分けるためTPを使っても、collectiveが重ければlatencyが改善しない場合がある。

比較するときは、次を分ける。

  • 1 requestのTTFTとITL
  • replica全体のrequest throughput
  • GPU当たりのtoken throughput
  • 同じlatency目標で処理できるconcurrency
  • modelを載せた後に残るKV Cache容量

総token throughputだけが増えても、対話requestのp99 ITLが悪化すれば用途に合わない。

KV Cacheの置き方もparallelismで変わる

TPではattention headやhidden dimensionの一部を各rankが担当し、KV Cacheも担当部分へ分かれる。TP sizeを変えればCache layoutとblock当たりのbyte数が変わりうる。

PPでは各stageが担当layerのKVを持つ。requestのDecode状態はpipeline全体に分散する。一つのstageだけ再起動しても、単純に別GPUへrequestを引き継げるとは限らない。

DPでは各replicaに独立したCacheがある。Prefix Cacheを有効にする場合、同じtenantや共通prefixを同じreplicaへ寄せるとhitしやすいが、Hot Keyでqueueが偏る可能性がある。

KV Cache量子化、Offloading、PagedAttentionの設定はrankごとに一致させる。TPやPPのdegreeを変えたdeployment間で、保存済みKV blockを互換だと仮定しない。

topologyを無視してGPUを番号で選ばない

同じ8 GPUでも、GPU間の接続は同じとは限らない。NVLinkで直接つながる組、PCIe switchを共有する組、別NUMA nodeの組がある。multi-nodeではInfiniBand、RoCE、通常のEthernetでも差が出る。

TPのcollectiveを遅いlinkへまたがせるとDecodeごとに影響する。PPのstage境界は通信頻度を抑えやすいが、activation転送とnode障害を考える必要がある。DP replicaは比較的独立しているため、node単位の配置とscaleがしやすい。

配置前にGPU topologyとNIC affinityを確認し、benchmarkでは通信backendが意図した高速経路を使っているかを見る。異なるGPU世代やmemory容量を同じparallel groupへ混ぜると、遅いrankが同期点で全体を待たせる。

failure domainが広がる

TP 8のmodel replicaは、8 rankの一つが停止すると通常そのreplica全体を継続できない。PPも一つのstageが失われればpipelineが途切れる。DPなら故障replicaをload balancerから外し、ほかへ新規requestを送れるが、故障replicaの実行中requestとKV Cacheは失われる。

大きなparallel groupほど、一つのrequestが依存するprocess、GPU、linkが増える。health checkはAPI processだけでなく、collective timeout、rank heartbeat、GPU error、network errorを含める。

rolling deploymentでは古いgroupをdrainしてから停止する。groupの一部だけ新versionへ変えない。weight、tokenizer、parallel degree、KV layoutを一つのdeployment revisionとして扱う。

監視するmetric

  • TP、PP、DPのdegreeと実際のrank配置
  • rankごとのGPU memory、SM利用率、memory bandwidth
  • collectiveの時間、回数、message size
  • pipeline stageごとの処理時間とidle時間
  • stage間activationの転送時間
  • DP replicaごとのrunning、waiting request数
  • replicaごとのKV Cache使用率とPrefix Cache hit率
  • TTFT、ITL、E2E latency、token throughput
  • collective timeout、rank再起動、request失敗
  • node間network throughput、retransmit、NIC error

平均GPU利用率だけでは、通信待ちとstage imbalanceを見分けられない。rank別のtimelineを取る。

構成を選ぶ負荷試験

同じ総GPU数で複数構成を比較する。8 GPUならTP 8、TP 4とPP 2、TP 2のDP 4などが候補になる。ただしmodelが収まる構成に限る。

  1. 短いpromptと長いpromptを分ける
  2. 短い回答と長い回答を混ぜる
  3. concurrencyを1から飽和まで上げる
  4. Prefix Cacheをcoldとwarmで測る
  5. node内とnode間のTPを比較する
  6. PP stageごとのlayer配置を変える
  7. DP routerへ偏ったHot Keyを送る
  8. rank、node、network linkを一つずつ停止する
  9. rolling deployment中のdrainを確認する
  10. 同じqualityとlatency条件でGPU当たりcostを比較する

最速の平均値だけでなく、p99 ITL、minimum free KV block、collective tail latency、故障時に失うrequest数を見る。

GPU数ではなく分割軸を設計する

Tensor Parallelismはlayer内の大きな計算を分ける代わりに頻繁なcollectiveを使う。Pipeline Parallelismはlayer群をstageへ分け、activationを順に渡す。Data Parallelismはmodelを複製し、requestを独立に処理する。

modelを収めるためのTPとPP、request throughputを増やすDPを分けて考え、GPU topology、KV Cache容量、latency目標から組み合わせる。GPUを増やした事実ではなく、通信待ちを含むGPU当たりの有効処理量で判断する必要がある。

参考資料