0 / 4 節読了

MoEモデルにおけるGPU負荷分散の根本的な課題

MoE (Mixture of Experts) モデルは、大規模言語モデルの効率的なスケーリングを可能にする強力なアーキテクチャです。しかし、その性能を最大限に引き出すためには、GPU間の負荷分散、特に「エキスパート並列 (EP) サービング」における最適化が不可欠です。各レイヤーでの同期は、最も遅いGPUの処理時間に律速されるため、GPU間の処理時間の均一化がシステムの全体スループットを決定します。 従来の負荷分散手法、例えばEPLB (Expert-Parallel Load Balancing)、LPLB (Layer-Parallel Load Balancing)、UltraEP、そしてMETROなどは、エキスパートの処理時間がトークン数やアクティブなエキスパート数に対して単純な線形関係にあるという仮定に基づいて設計されてきました。これは直感的であり、多くの初期の最適化において有効でした。しかし、この仮定が現実のハードウェア特性とワークロードの複雑さを十分に捉えきれていないことが、TEMPOの研究によって明らかにされました。 TEMPOは、2世代のデータセンター向けGPUを用いて詳細な測定を行いました。その結果、エキスパートの処理時間とトークン数の関係は、単純な線形ではなく、特定のトークン数閾値（約156〜168トークン）を境に、その特性が大きく変化する二つの異なる「レジーム」が存在することを発見しました。この発見は、MoEモデルのGPUサービングにおける最適化の考え方を根本的に見直す必要性を示唆しています。

メモリバウンドと計算バウンド：二つの支配的なレジーム

TEMPOが特定した二つのレジームは、それぞれ異なるボトルネックによって支配されます。 メモリバウンド・レジーム: 約156〜168トークンを下回る少量のトークン処理の場合、HBM（高帯域幅メモリ）からのウェイトストリーミングが処理時間の主要因となります。この状況では、エキスパートの処理コストは、トークン数そのものよりも「アクティブ化されたエキスパートのレプリカ数」に強く依存します。つまり、メモリからのデータ転送がボトルネックとなるため、どれだけ多くのエキスパートがメモリにロードされているか、その数によって処理時間が決まります。トークン数が少ない場合、GPUの計算ユニットはアイドル状態になりがちで、メモリからのデータ供給が追いつかないのです。 計算バウンド・レジーム: 一方、この閾値を超える大量のトークン処理の場合、グループ化されたGEMM (General Matrix Multiply) 演算が支配的となります。このレジームでは、GPUの計算ユニットがフル稼働し、計算能力がボトルネックとなります。特に、エキスパートを複数のGPUに「分割」して処理する場合、GPUのハードウェア特性上、トークン数が128タイルM-タイルに丸められるため、不要なパディングによる計算コストが発生します。これは、GPUが効率的に処理できるデータブロックサイズに満たない部分が、無駄な計算として加算されることを意味します。 TEMPOは、この二つのレジームを「max-affineプロファイル」として数学的にモデル化しました。t = max(a+bG, c+βN)という形式は、メモリバウンド（a+bG）と計算バウンド（c+βN）の両方の特性を同時に捉えることで、より現実的なエキスパート処理時間の予測を可能にします。このモデルは、異なるワークロードやハードウェア条件下での挙動を正確に予測するための基盤となります。

TEMPOディスパッチャーの革新的なアプローチ

TEMPOは、この二つのレジームが同時に発生する現実のデコードバッチの課題に対処するために開発されました。実際のバッチには、線形レジームで動作する「ホットなエキスパート」（多くのトークンを処理する）と、フラットなレジームで動作する「コールドなエキスパート」（少量のトークンを処理する）が混在します。従来のディスパッチャーは、どちらか一方のレジームに最適化されているため、このような混合環境では非効率な負荷分散を引き起こしていました。 TEMPOは、バッチごとのディスパッチ問題を「固定チャージ付きメイクスパン問題」として形式化しました。メイクスパン問題とは、複数のタスクを複数のリソースに割り当て、全タスクが完了するまでの総時間を最小化する問題です。これは一般にNP困難な問題ですが、TEMPOは、クリティカルパスから外れた場所でミリ秒単位で解を導き出す効率的なアルゴリズムを開発しました。これにより、システムのレイテンシに影響を与えることなく、リアルタイムに近い最適な負荷分散を実現します。 技術的な実装として、TEMPOはSGLangとの統合を実現しています。これにより、ディスパッチロジックはプロセス外で実行され、ディスパッチとカウント収集を単一のイングラフカーネルに融合することで、オーバーヘッドを最小限に抑えています。これは、既存のシステムに容易に組み込み、高い効率で動作させるための重要な設計です。

実証された性能向上と私の分析

TEMPOの有効性は、複数のテストベッドとモデルで実証されました。8GPUのテストベッドAでのマイクロベンチマークでは、TEMPOは既存の最良の固定ベースラインと比較して、常に1%以内の性能差に留まるという堅牢性を示しました。さらに重要なのは、メモリバウンドと計算バウンドのレジームが混在する特定のワークロードにおいて、最大15.5%という顕著な性能向上を達成した点です。これは、従来のディスパッチャーが捉えきれていなかった複雑な相互作用をTEMPOが効率的に管理できることを明確に示しています。 エンドツーエンドの性能評価では、テストベッドB上でQwen3-235Bモデルを使用した場合、TEMPOは4〜6%のスループット向上と、p99レイテンシの約15.6%削減を実現しました。Qwen3-235Bは、TEMPOが性能向上を期待できる「win region」に位置するモデルであり、この結果はTEMPOの設計思想が現実のモデルで有効であることを裏付けています。 一方で、DeepSeek-V3モデルのように、通信が支配的なボトルネックとなるモデルでは、TEMPOによる性能向上は限定的であり、メカニズムコストのみが観測されました。これは、TEMPOが解決しようとしている問題が、エキスパートの計算とメモリアクセスに起因するものであり、通信ボトルネックを直接解決するものではないためです。 私の見方では、TEMPOの最も重要な貢献は、ユニバーサルな「銀の弾丸」を主張するのではなく、「フェーズダイアグラム」という概念を提示した点です。これは、デプロイ前に特定のモデルやワークロードに対してTEMPOがどれほどの性能向上をもたらすかを予測できることを意味します。このような予測可能性は、AIプロダクト開発において極めて価値が高いです。無駄な試行錯誤を減らし、最適なアーキテクチャと運用戦略を最速で確立するために不可欠な情報を提供します。 TEMPOの研究は、AIモデルの効率的な運用には、単にモデルの規模を大きくするだけでなく、基盤となるハードウェアとソフトウェアの相互作用を深く理解し、それに応じたきめ細やかな最適化が不可欠であることを改めて強調しています。私の経験上、この種の一次情報に基づいた深い洞察こそが、競争優位性を生み出す源泉となります。

書籍ゼロからはじめるCodex

Kindleで読む →
柴亮太
柴亮太の視点

MoEの負荷分散は、単純な線形モデルでは限界があります。メモリと計算、二つのボトルネックを同時に捉えるTEMPOは、まさに一次情報に基づいた最適化です。自分は、このフェーズダイアグラムを既存プロダクトに適用し、最速でボトルネックを特定します。