混合エキスパートモデルは推論実行で密なモデルと比べてどうか？

Gradient Updates は著者個人の見解を述べるニュースレターで、Epoch AI 全体の見解とは限りません。

先週の Gradient Updates では、GPT-4o と Claude 3.5 Sonnet のパラメータ数が GPT-4 よりかなり少ないと推測できる理由を論じた。読者から最も多く寄せられた質問は、おおむね次のようなものだった。

推論実行の経済性には、総パラメータ数よりもアクティブパラメータ数のほうが重要なのではないか。そうだとすると、推論実行のコストと速度を見るだけで、どうやってモデルの総パラメータ数を推測できるのか。

これに答えるため、本号では混合エキスパート（MoE）モデルの推論実行という問題を取り上げることにした。基本的な結論は、MoEは総パラメータ数が同じ密なモデルより推論実行の効率が高いが、アクティブパラメータ数が同じ密なモデルよりは効率が低い、ということだ。大まかな目安として、8分の1スパースのモデルは、短いコンテキストでのデコードの経済性が、半分の大きさの密なモデルと同じになる。

簡潔さを優先して、Transformer の推論実行の基本概念は既知のものとして話を進める。そのため、この前提知識のない読者には本号はわかりにくいと思う。そうした読者には、本稿の前に Transformer の推論実行についての Pope らの論文、推論実行の計算についての kipply の記事、フォワードパスをどれだけ速く実行できるかについての Jacob Steinhardt の記事を読むことを勧める。

MoEモデルの利点

総パラメータ数を固定すると、推論実行でMoEモデルが密なモデルより優れている点は主に4つある。

MoEモデルは密なモデルよりアクティブパラメータが少ない。フォワードパスをゆっくり行う余裕がある場合には、これが優位の源になる。
MoEは密なモデルより浅く、幅が広い傾向がある。
MoEは、大きさと深さが同じ密なモデルと比べて、フォワードパス1回あたりのネットワーク通信が少なくて済む傾向がある。
MoEはアテンションブロックが小さい傾向がある。つまり、アテンションヘッドの数とヘッドの次元の積が小さい。これにより長いコンテキストでの提供コストが下がるかもしれないが、そうなるかどうかは、グループ化クエリアテンション（GQA）とマルチクエリアテンション（MQA）のどちらを使うと想定するかによる。

以下、これらの利点を順に論じる。全体として、(1) はフォワードパスをゆっくり行える場合には支配的だが、十分速く行わなければならない場合には関係がなくなる。(2) と (3) は、フォワードパスを速く行わなければならない場合に重要になる。(4) によって、長いコンテキストでの推論実行では、MoEは同程度の大きさの密なモデルより安くなる。

MoEは密なモデルよりアクティブパラメータが少ない

MoEモデルが密なモデルより優れている点として最もよく論じられるのは、処理または生成するトークン1個あたりの演算（FLOP）が少なくて済むことである。各トークンについて、モデルのパラメータのごく一部だけがアクティブになる、つまり乗算と累積に使われるからだ。演算がフォワードパスのボトルネックになっているなら、同じ大きさの密なモデルの代わりにMoEモデルを使えば、MoEのパラメータはすべてがアクティブになるわけではないので、推論実行は安く速くなる。

ここで重要な但し書きは、演算がボトルネックになっていなければならないということだ。実際には、これが当てはまるのは一部の場合だけである。

プレフィル（入力の事前処理）では、すべての入力トークンを並列に処理できるという利点がある。そのため、入力系列のすべてのトークンを一度に処理する1回のフォワードパスは、かなり遅くてもかまわない。たとえば、1000トークンの入力系列の処理に1秒かけても許容される。この条件では、ネットワークの遅延は問題にならず、処理するトークン単位で見たバッチサイズが大きいため、メモリ帯域幅もボトルネックにならない。また、フォワードパスにかかる時間に比べてネットワークの遅延が小さいので、通信を演算とより効果的に重ね合わせられ、より高帯域の通信プロトコルも使えるため、ネットワーク通信を演算の陰に隠すこともできる。これらすべての結果、プレフィルのコストは必要な演算の量で決まり、モデルの大きさが同じなら、MoEモデルは密なモデルに対して大きな優位を持つ。

デコードでも、大きなバッチサイズとインスタンスあたり少数のGPUを使ってゆっくり進めるなら、以上のことはすべて当てはまる。しかし、速く進めるには、インスタンスの規模を大きくして演算とメモリの負荷をより多くのGPUに分散させる必要があり、それによってネットワーク通信のオーバーヘッドが増える。ネットワーク通信にかかる時間を減らすには、バッチサイズを臨界バッチサイズ、つまり演算にかかる時間とメモリ読み出しにかかる時間が等しくなる点よりも小さくする必要がある。そうなると、行列積のためのメモリ読み出しのほうが演算より時間がかかるようになり、この状況ではMoEはアクティブパラメータの点で密なモデルに対して持っていた優位を失う。こうなるバッチサイズは、GPUの FLOP/s をTB/s単位のメモリ帯域幅で割ることで見積もれる。H100 では約300である。

大きなバッチサイズでネットワーク通信がどれだけの時間を占めうるかを示すため、8方向のテンソル並列で H100 の1ノードを使って Llama 3.3 70B を推論実行用に提供する場合を考えよう。フィードフォワードブロックだけでも、モデルの次元が8192なので、処理するトークン1個ごとに、各層で8192ワードのオールリデュースが必要になる。活性化の精度が16ビットなら、これはトークンあたり1 MBで、臨界バッチサイズの300では、フォワードパス1回あたり300 MBのデータをオールリデュースしなければならない。H100 の低遅延な NVLink オールリデュースプロトコルのオールリデュース帯域幅は112 GB/s なので、遅延を無視しても、フィードフォワードのオールリデュースだけでトークンあたり2.7ミリ秒かかる。これに対し、この構成で演算に必要な時間は約5ミリ秒である。参考までに、トークンあたり5ミリ秒というのは、Fireworks が自社版の Llama 3.3 70B を提供している速度とほぼ同じだ。臨界バッチサイズに近いところでこのような速度を達成することは、とてもできない。

MoEは密なモデルより浅く、幅が広い

アクティブパラメータの少なさはMoEが密なモデルより優れている主な点だが、それだけではない。パラメータ数が同じなら、MoEは密なモデルより浅く幅が広い傾向もあり、これが推論実行の経済性で優位をもたらす。直列の計算を増やした場合は、並列の計算を増やした場合よりも高速化がはるかに難しいからである。

この一般的な傾向の例として、Mixtral 8x22B は1,400億パラメータで56層のモデルだが、Llama 3 70B は700億パラメータで80層のモデルである。層の少なさがとくに重要になるのは、短いコンテキストで高速にデコードする場合である。この場合、フォワードパスの間にネットワークの遅延を被る回数がモデルの層数で決まるからだ。浅く幅の広いモデルはネットワーク通信も少なくて済み、これも同程度の大きさの密なモデルに対してMoEが持つことの多い利点である。

MoEの深さが、アクティブパラメータ数が同じ密なモデルと同じようにスケールすると仮定すれば、MoEの推論実行の経済性は、全体の大きさが同じ密なモデルよりは良く、アクティブパラメータ数が同じ密なモデルよりは悪いと予想される。推論実行の経済性には、モデルの深さと総パラメータ数の両方が効くからである。前者はフォワードパスの間に遅延を被る回数を決め、後者はトークンを生成する際に HBM から読み出す必要のある情報の量を決める。

深さが同じなら、MoEモデルは密なモデルよりネットワーク通信が少ない

バッチサイズとハードウェア構成を固定すると、処理または生成するトークン1個あたりにフィードフォワードブロックが必要とするネットワーク通信の量は、モデルの次元、モデルの深さ、アクティブなエキスパートの数の積に比例する。これを使えば、深さとパラメータ数が同じMoEモデルと密なモデルを、必要なネットワーク通信の量で比べられる。

まず GPT-4 から考えよう。層あたりのパラメータ数はモデルの次元の2乗に比例し、GPT-4 のエキスパートは全部で16個なので、GPT-4 と深さもパラメータ数も同じ密なモデルは、モデルの次元が4倍でなければならない。GPT-4 のアクティブなエキスパートは2個なので、上の積によるスケーリングから、GPT-4 は大きさと深さが同じ密なモデルより通信量が50%少なくて済むことがわかる。

この議論を一般化すると、MoEのアクティブなエキスパートの数がエキスパートの総数の平方根より少ない限り、そのMoEは総パラメータ数と深さが同じ密なモデルに対して、ネットワーク通信の点で優位を持つことがわかる。この条件は GPT-4 で成り立ち、エキスパートが全部で8個でそのうち2個がアクティブな Mixtral 8x22B のようなモデルでも成り立つ。

実際には、MoEは密なモデルに対して、ここで計算したよりもさらに大きな優位を持つ傾向がある。私の計算はモデルの深さを固定すると仮定しているが、実際には先に述べたようにMoEはより浅く幅も広いので、同じ大きさの密なモデルに対する通信面での優位はさらに大きくなる。

MoEは密なモデルよりアテンションブロックが小さい

MoEモデルのアテンションブロックは、フィードフォワードブロックがエキスパート1個だけからなるモデルにふさわしい大きさに設定されることが多い。たとえば GPT-4 のアテンションブロックは、合計550億パラメータで GPT-3 のものとほぼ同じ大きさだが、総パラメータ数で見ると GPT-4 は約10倍大きい。その一因は GPT-4 がマルチクエリアテンションを使っていることだが、それで説明できるのは3倍の差だけである。したがって GPT-4 のアテンションブロックは、同程度の大きさの密なモデルから予想されるものより、なお約3倍小さい。

この差によって、MoEモデルは長いコンテキストでの推論実行が同じ大きさの密なモデルより安くなりうる。キーとバリューのヘッドのアテンショングループの大きさを固定すると仮定すれば、アテンションの演算コストも KV キャッシュの読み出しコストも、アテンションヘッドの数とヘッドの次元の積に比例するからである。しかし、アテンショングループの大きさをアテンションヘッドの数に比例して大きくしていくと考えるなら、スケールしてもトークンあたりの KV キャッシュの大きさは一定のままなので、アテンションブロックが小さくなることの利点は得られない。

実際には、ここで期待できる効果はかなり小さい。この節約について最も楽観的な場合、つまりグループの大きさが固定されたグループ化クエリアテンションを仮定しても、トークンあたりの KV キャッシュの大きさはモデルの大きさの立方根に比例するにすぎないと予想される。したがって、エキスパートが16個のMoEモデルは、同じ大きさの密なモデルより推論実行のコストがおそらく60%安くなる程度である。これは、グループ化クエリアテンションを使うことの効果よりはるかに小さい。グループ化クエリアテンションはそれだけで、Llama 3 70B と Llama 3 405B の KV キャッシュの大きさを8分の1にしている。

MoEの推論実行での優位を見積もる

では、以上の考察から、MoEの推論実行での優位の大きさを定量的に見積もってみよう。前と同じく、簡単のため GPT-4 を基準のモデルとし、密なモデルの深さと重み行列の2つの次元はそれぞれ、モデルの大きさの立方根に比例すると仮定する。

これらから、次のことが推測できる。

GPT-4 の演算コストは、2,750億パラメータの密なモデルに相当する。
GPT-4 のメモリ帯域幅のコストは、1.8兆パラメータの密なモデルに相当する。
GPT-4 の直列の深さは、4,000億パラメータのモデルとほぼ同じである（Llama 3 405B の126層に対して120層あるため）。
GPT-4 のGPU間のデータ移動コストは、1.8兆パラメータの密なモデルの半分で、上のスケーリングの仮定のもとでは約6,350億パラメータの密なモデルに相当する。
GPT-4 のアテンションブロックの大きさは、Llama 3 405B のアテンションブロックの約75%なので、3,000億パラメータの密なモデルに相当する。

プレフィルや、短いコンテキストでのゆっくりしたデコードでは、GPT-4 は実質的に2,750億パラメータの密なモデルのようにふるまう。これは驚くことではない。この領域では、使われるバッチサイズと系列長が大きいため、フォワードパスは演算に律速されるからだ。コストは、モデルのスパース性の倍率と同じだけ下がる。

短いコンテキストでの高速なデコードでは、バッチサイズが小さすぎる傾向があるため、演算コストは関係ない。メモリ読み出しの時間とネットワークの時間（遅延と帯域幅の両方）がほぼ同じくらい重要だと仮定しよう。すると、GPT-4 の「密なモデル換算」の大きさは、メモリ帯域幅のコストに0.5、直列の深さとネットワーク通信にそれぞれ0.25の重みを付けて幾何平均をとるのが自然な見積もり方である。後の2つの要因が合わさってネットワーク通信の総時間を決めるからだ。こうすると、短いコンテキストでの高速なデコードでの GPT-4 の「密なモデル換算の大きさ」は約9,500億パラメータになる。

一般に、上の計算をもっと詳しくたどると、ここで選んだ平均のとり方のもとでは、MoEの総パラメータ数を (エキスパートの数)^(0.44) / (アクティブなエキスパートの数)^(0.63) で割れば、密なモデル換算の大きさが得られることがわかる。たとえばこれによれば、1,400億パラメータの Mixtral 8x22B は、短いコンテキストでの高速なデコードでは約900億パラメータの密なモデルに相当するはずである。

長いコンテキストではアテンションの仕組みがコストの大半を占め、同様の「密なモデル換算の大きさ」を見つけるのは難しい。推論実行の経済性には、KV ヘッドをどうグループ化するかの細部のほうが、モデルそのものの大きさよりはるかに重要になるからである。

結論

混合エキスパートモデルは一般に、推論実行用に提供するコストが密なモデルより安い。しかし、プレフィルを除けば、それはアクティブパラメータが少ないことが直接の理由ではない。同じ大きさの密なモデルと比べて層が少なく、フォワードパスに必要なネットワーク通信も少ない傾向があるからである。したがって、アクティブパラメータ数だけを素朴に見てMoEの推論実行コストを見積もると、誤った結論に導かれる。

推論実行コストの大まかな見積もりが必要なら、検討している推論実行の領域ごとに「密なモデル換算の大きさ」を計算すれば、MoEを密なモデルと比べられる。短いコンテキストでの高速なデコードでは、4分の1スパースのモデルと8分の1スパースのモデルは、それぞれ総パラメータ数の65%と50%の密なモデルに相当する。そのため、モデルがどれだけ速く提供されているかを見るときには、アクティブパラメータ数よりも総パラメータ数のほうがおそらく推測しやすい。