Expert Parallelismとは？MoEモデルでtokenをGPU間に振り分ける仕組み

Mixture of Expertsは、多数のexpertを持ちながら、各tokenではその一部だけを使う。総parameter数が大きくても、token当たりの計算量を抑えられるのが特徴である。ただし、選ばれたexpertが別GPUにあれば、tokenをnetwork越しに送り、処理後に元の順序へ戻さなければならない。

Expert Parallelismは、MoE layerのexpertを複数GPUへ分散して保持する方法である。大きなMoE modelを載せやすくなる一方、性能はmatrix multiplicationだけでなく、routerの偏りとAll-to-All通信に左右される。

MoEはすべてのparameterを毎token使わない

dense TransformerのFFNでは、各tokenが同じweightを通る。MoE layerではrouterがtokenごとにscoreを計算し、top-kのexpertを選ぶ。top-1なら一つ、top-2なら二つのexpertがそのtokenを処理する。

64 expertを持っていてもtop-2なら、各tokenが使うexpertは2個である。総parameter数とactive parameter数を分けて考える必要がある。

Switch Transformerはtop-1 routingを使い、model capacityを増やしながらtoken当たりの計算を抑える構成を示した。GShardはtop-2 gatingを使うMoEを多数のdeviceへ分散した。

Expert Parallelismはexpert単位でweightを分ける

8 expertを4 GPUへ分けるなら、単純な配置では各GPUが2 expertずつ持つ。tokenのrouter結果に応じて、各GPUは担当expertへ送るtokenをまとめる。

処理は概ね次の流れになる。

各rankでtokenのrouter scoreと選択expertを計算する
tokenをexpert ownerごとに並べ替える
All-to-Allで担当GPUへtokenを送る
各GPUがlocal expertのFFNを実行する
結果をAll-to-Allで元のrankへ戻す
元のtoken順に並べ直し、選択weightを反映する

Tensor Parallelismが一つのmatrixを複数GPUへ分けるのに対し、Expert Parallelismは異なるexpert matrixをGPUごとに持つ。MoE以外のAttentionやshared layerは、Data ParallelismやTensor Parallelismで処理する構成がある。

All-to-Allはtoken数と偏りに影響される

All-to-Allでは、各rankが複数rankへ異なるtoken集合を送る。送信量はbatch内でどのexpertが選ばれたかによって毎step変わる。

routerが均等なら、各GPUへ近い量のtokenが届く。特定expertへ集中すると、そのGPUだけexpert計算が長くなり、ほかのGPUは同期点で待つ。平均GPU利用率が高くても、一つのhot expertがDecode step全体を遅らせることがある。

通信量は、選ばれたtoken数、hidden dimension、data type、top-kに依存する。top-2はtop-1より多くのexpert出力を使えるが、送信と計算も増える。

trainingのload balanceとinferenceのtraffic偏りは別である

MoE modelは学習時にload balancing lossやexpert capacityを使い、routerが一部expertへ偏りすぎないようにすることがある。Switch Transformerでは、expert capacityをbatch内token数、expert数、capacity factorから定める。

しかし、学習dataで均等だったmodelが本番trafficでも均等とは限らない。日本語、code、数学、特定業界の文書など、入力domainが変われば選択expertの分布も変わる。

online inferenceではbatchが小さく、token到着も不均一である。平均では均等でも、あるDecode stepだけ特定expertへ集中することがある。学習時のauxiliary lossだけを根拠にcapacityを決めず、本番requestのexpert histogramを測る。

expert capacity超過をどう扱うか

固定shapeを前提とする実装では、expertごとに処理できるtoken数の上限を持つ場合がある。上限を大きくすればoverflowは減るが、padding、memory、communicationが増える。小さくすれば効率は上がるが、偏ったbatchで超過しやすい。

trainingでは超過tokenをdropまたは別expertへ送る設計がある。しかし推論でtokenをdropすれば出力品質が変わる。runtimeがoverflowを許さない、bufferを拡張する、chunkへ分ける、fallback expertへ送るなど、実装ごとの挙動を確認する。

capacity超過をOOMやtimeoutとして初めて知る構成にしない。expert別queueと最大dispatch数を監視し、偏りを再現する負荷試験を行う。

expertの複製でhot spotを減らせる

頻繁に選ばれるexpertを複数GPUへ複製し、router結果を空いているreplicaへ送る方法がある。load balanceは改善するが、追加のweight memoryとrouting metadataが必要になる。

複製候補は長期平均だけで選ばない。時間帯やtenantでhot expertが変わる場合、静的な複製は外れる。expert利用率、queue、通信量を観測し、配置変更の効果をcanaryで確認する。

配置を動的に変える場合は、全rankが同じowner mapを参照しなければならない。versionの異なるmapが混ざると、tokenを誤ったexpertへ送る。更新はgeneration単位または安全なbarrierで切り替える。

topology-aware placementが必要になる

頻繁にtokenを交換するEP groupは、高速なGPU間接続とnetworkを必要とする。node内NVLinkとnode間networkではlatencyと帯域が違う。

すべてのexpertを遠隔nodeへ均等配置することが最適とは限らない。よく同時選択されるexpert、shared expert、Attention rankとの距離を考える。Tensor Parallelismも併用する場合、TP collectiveとEP All-to-Allが同じlinkを奪い合う可能性がある。

vLLMのExpert Parallel Deploymentでは、MoE layerをEP rankへshardし、Attention layerはData ParallelとTensor Parallelを組み合わせられる。設定上のgroupサイズだけでなく、実際のrank、GPU、NIC対応を記録する。

small batchではexpert計算が細かくなりすぎる

多数expertへ少数tokenを分散すると、各expertが処理するmatrixが小さくなり、GPU kernelの効率が落ちる。通信と並べ替えの固定costも相対的に大きくなる。

online servingではContinuous Batchingで複数requestのtokenをまとめ、同じlocal expert向けtokenをgrouped GEMMで処理する。batchを大きくすれば効率は上がるが、待機時間とITLが伸びる可能性がある。

GPU利用率だけでなく、expert当たりtoken数、grouped GEMMのshape、dispatch待ち、p99 ITLを一緒に見る。

KV Cacheはexpert weightとは別に扱う

MoEのexpertは主にFFN部分であり、AttentionのKV Cacheとは役割が違う。EPでexpert weightを分散しても、KV Cacheが自動的に同じ割合で減るとは限らない。

AttentionをData Parallelにすれば、各DP rankはrequestごとのKV Cacheを独立に持つ。AttentionをTensor Parallelにすれば、KVはheadやdimensionに応じてshardされる。MoEの総parameter数だけからKV capacityを推定しない。

routerとexpert配置を変えても、実行中requestのKVを別rankへ安全に移せるとは限らない。scale変更やdrain時は、Attention側のstateも含めて設計する。

failure domainと縮退経路

一つのexpert ownerが停止すると、そのexpertへroutingされたtokenだけでなく、collectiveへ参加するEP group全体が止まる場合がある。processのhealthだけでなく、collective timeout、GPU error、rankごとの進捗を監視する。

故障時に別expertで代替すれば出力分布が変わる。通常は同じweightのreplicaへ切り替えるか、requestを健全なmodel replicaで最初から再実行する。途中から別versionのexpertへ送らない。

EP groupを大きくするほど一つのrequestが依存するGPUとlinkが増える。Data Parallel replicaを複数持ち、故障groupをload balancerから外せる構成も検討する。

監視するmetric

layer・expert別の選択token数と割合
top-1とtop-kのscore差
expert別の最大、平均、p95 token数
load imbalanceとhot expertの継続時間
dispatch、All-to-All、combineの時間とbyte数
local expert GEMMの時間とtoken数
padding、capacity overflow、fallback回数
rank別GPU利用率、memory、network帯域
TTFT、ITL、E2E latency、token throughput
language、task、tenant別のexpert分布
collective timeout、rank failure、request再実行

全expert平均だけではhot spotを隠す。最大値と分布をlayer別に持つ。

負荷試験

日本語、英語、code、数学、業務文書を分ける
短いpromptと長いpromptを混ぜる
concurrency 1から飽和まで上げる
特定expertへ偏る入力を探索して連続送信する
top-kとexpert capacityの設定を変える
EP sizeとTP sizeの組み合わせを比較する
node内とnode間のAll-to-Allを比較する
TP collectiveとEP通信を同時に飽和させる
expert ownerとnetwork linkを一つずつ停止する
model更新時のowner map切替とdrainを確認する

同じ総GPU数で、Expert Parallelを広げる構成とMoE model replicaを増やす構成を比較する。平均throughputだけでなく、GPU当たりtoken、p99 ITL、故障時の影響範囲を見る。

Expert Parallelismはparameter分散ではなくtoken routingまで含む

Expert Parallelismは、expert weightを複数GPUへ置くだけの仕組みではない。各stepでtokenを選択expertのownerへ送り、処理結果を元へ戻す分散routingである。

性能を決めるのは総parameter数ではなく、active expert数、routerの偏り、expert当たりbatch、All-to-All、GPU topologyである。学習時のload balanceを前提にせず、本番trafficでexpert分布とtail latencyを測る必要がある。