Mixture-of-Expertsのルーティングでローカルモデルの計算量はどう変わるのか？

Mixture-of-Experts（MoE）ルーティングでは、すべてのエキスパートを有効化する代わりに、各トークンを少数のエキスパート・フィードフォワードネットワークだけに通すことで、ローカルモデルの計算方法が変わります。

そのため、大規模なMoEモデルはホームサーバー上で一見変わった存在になります。チェックポイント全体は非常に大きくても、1トークンに対して実行されるエキスパート計算は、はるかに小さなアクティブモデルに近い場合があるためです。利点はスパースな演算であり、メモリが不要になることではありません。エキスパートの重みは保存しておく必要があり、複数デバイスでのルーティングでは、削減できた計算の一部が通信処理に置き換わることもあります。

通常のTransformerはすべてのトークンで同じフィードフォワードパラメータを有効化する

通常の密なTransformerレイヤーでは、すべてのトークンが同じフィードフォワードネットワークを通過します。そのため、フィードフォワードのパラメータ数を増やすと、モデルサイズだけでなく、各トークンで実行される演算量も増加します。

Switch Transformerはこれと対照的に、トークンごとに異なるパラメータサブセットを選択できるスパースに有効化されるMixture-of-Expertsレイヤーを採用しています。

ローカル推論では、ここからMoEに関する最初の区別が生まれます。総パラメータ数は保存すべきモデル状態の量を決め、アクティブパラメータ数は、トークンごとのエキスパート計算量の大部分を決めます。

ルーターは各トークンを利用可能なエキスパートと照合してスコアリングする

MoEレイヤーでは、小さなルーターがトークン表現を調べ、候補となるエキスパートに対するスコアまたは確率を生成します。実行時には、アーキテクチャに応じて通常top-1またはtop-kの限られた数が選択されます。

NVIDIA Megatron Coreには、各トークンをどのエキスパートに送るかを決定するルーターが用意されています。ルーティングの決定によって、そのトークンを処理するエキスパート行列が変わります。

したがって、同じプロンプト内の2つのトークンは、アテンションやその他の密なモデル部分を共有していても、異なるエキスパート重みを有効化することがあります。

スパースな有効化はエキスパート計算を減らすが、重み全体の保存量は減らさない

1つのレイヤーに8つのエキスパートがあり、各トークンを2つにルーティングする場合、そのトークンに対して実行されるのは選択されたエキスパートの計算だけです。残りのエキスパート行列は、そのステップでは休止しています。

Expert Choiceの研究では、MoEを、パラメータ数を増やしながらトークンごとの計算量を比較的抑える手法として説明しています。これが、MoEのチェックポイントで総パラメータ数が非常に大きく示されていても、アクティブパラメータ数ははるかに小さくなり得る理由です。

ランタイムがエキスパートをストリーミングまたはオフロードしない限り、ホームサーバーは依然として、すべてのエキスパート重みをどこかで利用可能な状態にしておく必要があります。スパースな計算は、スパースな保存と同じではありません。

Featured

ZimaCube 2 パーソナルクラウドホームNAS

ZimaCube2

$799.00 - $2,499.00

Buy Now

エキスパートの配置によって、ルーティングがデバイス間通信に変わることがある

1台のGPUでは、エキスパートの選択によって主に読み込まれるローカルな重みブロックが変わります。複数のGPUにまたがる場合は、エキスパートが分散配置されるため、トークン表現を選択されたエキスパートを所有するデバイスへ送り、その後メインのレイヤー処理へ戻さなければならないことがあります。

DeepSpeedのMoE推論設計では、GPU間のエキスパート並列処理を使ってエキスパートのグループを割り当てます。したがって、ルーティングには行列乗算だけでなく通信も含まれます。

そのため、小規模なホームマルチGPUサーバーでは、各トークンの演算がスパースであっても、インターコネクトがボトルネックになる可能性があります。PCIeトポロジーやピア・ツー・ピア対応は、アクティブパラメータ数だけから想定される以上に重要です。

ルーティングの偏りによって、一部のエキスパートが過負荷になり、容量が遊休化する

実際のプロンプトでは、トークンが均等に分散されるとは限りません。多くのトークンが同じエキスパートを選択すると、そのデバイスやカーネルがホットスポットになり、他のエキスパートは待機することになります。

MoETunerは、不均衡なトークンルーティングとエキスパート有効化を、実運用上の提供ボトルネックとして挙げています。これは、一部のデバイスに他より多くのルーティング処理や通信が集中する可能性があるためです。

ローカル推論では、コンテキスト長やモデルサイズが似ていても、プロンプトによってトークン遅延が変動する形で、ルーティングの偏りが現れることがあります。モデルが動的に異なる計算経路を選択しているためです。

ローカルMoEの容量は、総重みとアクティブ経路を分けて計画する

ハードウェアを見積もる際は、モデルの総重みサイズ、トークンごとのアクティブなエキスパート数、密なレイヤーのコスト、KVキャッシュ、量子化、さらにエキスパートが1台のアクセラレーター、複数のGPU、CPU RAM、ストレージのどこに置かれるかを記録すると有用です。

ZimaSpaceのローカルAIのメモリとアクセラレーター計画に関する解説では、ホームサーバーの容量を考える際の幅広い背景を確認できます。MoEでは、常駐容量とアクティブな計算をさらに分けて考える必要があります。

ルーティングパターンは変化する可能性があるため、異なるタスクのプロンプトでベンチマークを行ってください。公称のアクティブパラメータ数がシステム全体の性能を予測すると決めつけず、1秒あたりのトークン数、デバイスごとの使用率、エキスパート負荷、通信トラフィック、メモリ常駐量を測定しましょう。

特に大規模なローカルチェックポイントについては、ZimaSpaceの大規模モデルの展開に関する分析が、隣接する容量上の論点になります。ルーティングによってアクティブな計算は変わりますが、総重みの常駐量とデバイス間の配置が、モデルを実用的に動かせるかどうかを左右します。