TEMPO: Makespan-Aware Expert-Parallel Load Balancing Across Memory- and Compute-Bound Regimes
TEMPOは、メモリ拘束および計算拘束のレジームにおける非線形なエキスパート実行時間をモデル化し、トークン分布を動的に最適化することで、従来の線形カウントベースの手法が失敗する混合レジームのシナリオにおいて、最大15.5%のスループット向上と大幅なレイテンシ削減を実現する、メイクスパンを考慮したエキスパート並列ロードバランシング・ディスパッチャを導入する。
原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。 免責事項の全文を読む
あなたは、眠らない街のための大規模で高速なピザデリバリーサービスを運営していると想像してください。あなたには、同一のデリバリードライバー（GPU）の艦隊と、何百もの異なる専門シェフ（AIモデルにおける「エキスパート」）を擁する中央キッチンがあります。顧客がピザを注文するたびに、システムはどのシェフがそのピザを作り、どのドライバーが完成したピザを運ぶかを決定しなければなりません。Mixture-of-Experts（MoE）と呼ばれるAIの世界では、まさにこれが起こっています。これらのモデルは、数千の小さな専門的な「小脳」で構成された巨大な脳のようなものです。AIが思考するとき、脳全体を一度に使用するのではなく、その仕事を処理するために特定の数名のエキスパートを選び出します。
大きな課題は、チーム全体を同じスピードで動かし続けることです。もし一人のドライバーが巨大で複雑な注文品を抱え込み、他の全員が手持ち無沙汰になってしまったら、デリバリー全体が遅れてしまいます。一連の注文を完了させる時間は、グループの中で最も遅い人に左右されます。長年、ワークロードをバランスさせるための標準的なルールは単純でした。「注文数を均等に分割する」というものです。もし100件の注文があれば、10人のドライバーにそれぞれ10件ずつ配分します。これは論理的に見えます。リンゴの山を平等に分け合うようなものです。しかし、もし中には重い岩のようなリンゴもあれば、軽い羽のようなリンゴもあるとしたらどうでしょう？あるいは、もしキッチンに「新しいシェフを呼び出すには、作るピザの量に関わらず一定の時間が必要である」というルールがあったらどうでしょうか？古いルールは、時間は常に注文の「数」に直接結びついていると仮定していました。しかし、この論文はこう問いかけます。「もしその仮定が間違っていたらどうなるだろうか？」と。
KlingAIで研究を行っているこの論文の著者たちは、従来の「注文数を数える」というルールが実は罠であることを発見しました。彼らは、現代のAIハードウェアにおいて、エキスパートを処理する時間は単にそれが扱うトークン（単語やデータの塊）の数によって決まるのではないことを突き止めました。それは「二面性を持つ獣」なのです。ある時は、メモリバンクからエキスパートの「レシピ（重み）」を読み込む作業の純粋な労力によって時間が支配されます。これは、注文がどれほど小さくても、読み込みにかかる固定の時間が発生することを意味します。また別の時には、一度レシピが読み込まれてしまえば、時間は注文数に応じて線形に増加します。単に注文の数だけを見ていた古い手法は、レシピを読み込むという隠れたコストを見落としていました。彼らは、重い岩と軽い羽が混ざったリンゴの山を、重さを量らずに数だけでバランスさせようとしていたのです。
これを解決するために、チームは新しいディスパッチャーであるTEMPO（Time-modeled Expert-Parallel Optimization）を構築しました。トークンを数える代わりに、TEMPOはキッチンの物理法則を理解するスマートな交通管制官のように振る舞います。TEMPOは、シェフのレシピを読み込むのにどれくらいの時間がかかるか、そしてピザを調理するのにどれくらいの時間がかかるかを正確に測定する特別な「コストモデル」を使用します。TEMPOは、「冷えた（しばらく使われていない）」エキスパートの場合、その小さな注文を2人のドライバーに分割してしまうことは、読み込み手数料を2回支払うことになるため災難であると理解しています。しかし、「熱い（頻繁に使われている）」エキスパートで注文が山積みになっている場合は、分割しても問題ありません。
論文は、TEMPOが単に推測しているのではなく、ミリ秒単位ですべてのリクエストに対して完璧なバランスを計算していることを示しています。彼らが実際のAIモデルでテストしたところ、古い手法はしばしば15%遅かったり、列の最後尾にいる数人の顧客に対して大幅な遅延を引き起こしたりすることが分かりました。しかし、TEMPOは列をスムーズに動かし続けます。それは、「全員に同じ数のリンゴを配る」というルールから、「全員に同じ量の仕事を与える」というルールへ切り替えるようなものです。そこには、リンゴが重いか、あるいはシェフが起きるのが遅いかといった事情も考慮されています。
研究者たちは、この新しい手法がどこで機能し、どこで機能しないのかを正確に示すことに細心の注意を払いました。彼らは、もし「熱い」エキスパートがあまりにも多すぎて、システムが単に膨大なデータ量によって圧倒されている場合（「計算量限定（compute-bound）」の領域）、従来のトークン計数法でも実は問題ないことを証明しました。しかし、現実の世界、つまり一部のエキスパートが忙しく、他のエキスパートが休んでいる状況で、「読み込み手数料」が高い場合には、TEMPOが輝きます。彼らはさらに、AIトラフィックの「気象図」のような「フェーズ・ダイアグラム（相図）」を作成し、新しい手法がいつ時間を節約できるのか、そして古い手法がいつ十分なものなのかを正確に予測しました。
結局のところ、この論文は単なる高速なアルゴリズムについての話ではありません。それは、AIにおけるワークロードのバランスの取り方に対する私たちの考え方を変えるものです。現代の複雑で高速なAIの世界では、単に物を数えるだけでは不十分であることを教えてくれます。データの移動に伴う隠れたコストと、仕事の具体的な形状を理解しなければなりません。実際の作業時間を測定することで、アイテムの数を数えるのではなく、仕事にかかる実時間を捉えることで、TEMPOはAIモデルをより速く、より効率的にし、将来の膨大な需要に対応できるものにします。それは混沌としたキッチンを、よく整備された機械へと変え、ピザが保温ランプの下で待たされることなく、どのドライバーも待ちぼうけを食らうことがないようにするのです。
自分の分野の論文に埋もれていませんか？
研究キーワードに一致する最新の論文のダイジェストを毎日受け取りましょう——技術要約付き、あなたの言語で。