Distributing LLM inference in DwarfStar

Salvatore Sanfilippo

DwarfStarにおけるLLM推論の分散化

原文は Salvatore Sanfilippo により に公開されました。 このブログを購読する

ハイエンドのNVIDIAカードと、それらを動かすために必要なサーバーや電力には莫大なコストがかかる。特に巨大なモデルを動かすのに十分なVRAMを確保しようとすればなおさらだ。これまでの代替策はAppleのハードウェアか、DGX Sparkだった。Sparkはメモリ帯域幅が著しく制限されているとはいえ、LLMのプロンプト処理(prefill)を十分な速さで実行することはできる。Mac Studioは最大512GBのユニファイドメモリを提供し、控えめながらもSparkよりはるかに優れたメモリ帯域幅と演算性能を、当時の状況を考えれば比較的妥当な価格で実現していた。

例えばDwarfStarを使えば、Mac Studio M3 Ultra 512GBモデルでDeepSeek v4 PROをprefill 150 t/s、デコード約10〜13 t/sで動かすことができる。決して速いとは言えないが、特定のユースケースでは実用的な水準だ。2ビット量子化してもDeepSeek v4 PROの耐性は非常に高く、同じ量子化のFlashと同様だ(今日はPROにCコンパイラを書かせてみた。近々動画を公開する予定だ)。約12kドルの総コストでフロンティアモデルを自宅で動かせるというのは、決して些細なことではないと思う。

この状況は今後どんどん良くなっていくと期待したくなるが、先行きは曇っている。NVIDIA構成が安くなる望みはほぼゼロだし、小さな企業でさえ、ローカル推論のために小さなデータセンターを気軽に購入して運用できるわけではない。同時にRAM不足により、Mac Studio M5 Ultraが登場する可能性は決して高くない。仮に出たとしても1.2T/sのメモリ帯域幅とさらなる演算性能を備えることになるのだろうが(M5 Maxはすでに演算性能ではより高速で、各GPUコア内に特定のモデルで威力を発揮するNeural Acceleratorを搭載している)。

つまり、ローカル推論における現状のベストマシンはおそらくラップトップだ。M5 Max 128GBモデルなら、DeepSeek v4 FlashやMimo V2.5を2ビット量子化で、十分に良好なprefillとデコード速度で動かせる。具体的にはprefillで約500 t/s、デコードで約35〜40 t/sで、コンテキスト長が伸びた際の性能低下の傾きも非常に緩やかだ。構成にもよるが6〜7kドルのコストで、現状では最もお得な選択肢の一つと言える。

こうした状況を踏まえると、ローカル推論プロジェクト一般、とりわけDwarfStarにとって、分散推論に目を向けることが面白くなってくる。MacBook M5 Maxが2台、3台、4台あったら何ができるだろうか。あるいは512GB RAMを積んだM3 Ultraが2台あったらどうか。

従来、分散推論には主に2つの方式がある。1つはメモリを倍増させるために、Transformerレイヤーの50%をコンピュータAに、残りの50%をコンピュータBにロードし、逐次的に推論を実行する方法だ。この場合、やり取りするのはアクティベーションだけで、概念的には非常にシンプルだ。さらにマイクロバッチの工夫を加えれば、単にメモリを2倍にするだけでなく、理論上はプロンプト処理速度を大幅に向上させることも可能になる(ただしデコードは別だ。1トークンを生成するたびにマシンAの前半レイヤー、次にマシンBの後半レイヤーと順に待たなければならない――とはいえ発熱は抑えられるので持続的な負荷での運用は可能になる)。これは決して悪くない。例えば512GBのMac Studioを2台持っている幸運な人なら、DeepSeek v4 PROのフルサイズ版を動かすことができ(もっとも2ビット量子化版でも非常によく動くのだが)、マイクロバッチを使えばより高速なprefillさえ期待できる。

もう1つのアプローチは、AppleのRDMAを使って2台のマシン間で実行を並列化する、いわば垂直分割だ。例えば同じ2ビット量子化モデルをマシンAとBの両方に収まるようにロードし、両方がすべてのrouted expertsを持つようにする。そして各レイヤーで、半分のエキスパートをマシンAで、残り半分をマシンBで実行するように協調させるのだ(両方のマシンがすべてのエキスパートを持っているため、ルーターが何を選択しても、計算の50%をもう一方のマシンに送ることができ、アクティベーションはごく小さい)。これはrouted expertsがはるかに大きいPROの方が、通信のペナルティが相対的に小さくなるため、より現実的だ。ただし、これがうまく機能するかどうかは、まだまったく未知数だ。

テンソル並列もあるじゃないか、と思っただろう?しかし2台のAppleコンピュータ間や2台のDGX Spark間などの通信速度では、これはまったく実現可能ではないと断言できる(NVLinkの速度を調べてみてほしい)。上記2つの方式の妙味は、送るデータがごくわずかで済むという点にある。

さて、ここまで読んで、これはLLMを並列実行することについて誰もが知っている同じ話じゃないか、と思ったはずだ。実際その通りだ。だがこの記事は、まさにこの地点にたどり着くために書かれた。もし2台のMacやDGXをまったく別の方法で並列化できるとしたらどうだろう?オープンウェイトモデルは今や黄金時代にあり、数は豊富で、その多くが非常に強力だ。128GBで動く2ビット量子化のクラスだけでも興味深いモデルは多い。Minimax M2.7、Mimo V2.5、DeepSeek v4 Flashなど、他にもいくつかある。同時に、最近LLMのアンサンブル(https://arxiv.org/abs/2502.18036)が、十分に研究されていない可能性として注目されている。これは2つのモデルを2台の異なるマシンで完全にシェアードナッシングな形で実行し、最後にlogitsを統合するか、最良の続きを選択するだけという手法だ。これにはいくつかのやり方があり、2つのモデルが異なる語彙を持っていても機能する。パープレキシティが低い方の続きを選ぶ(つまり、より確信を持っている方のモデルを選ぶ。暗黙的なルーティングを持つ2エキスパートのMoEのようなものだ)こともできるし、logitsを統合して(異なる語彙による複雑さはあるが)そこからサンプリングすることも可能だ。より最近の論文では、この2つの手法を組み合わせるのが最良だと示唆されている。いずれにせよ、これらの手法は実際に効果があるようで、モデルは単独よりも良い結果を出すように見える。それぞれが次に何を言うべきかについて独自の視点を持ち寄ることで、知識が向上するかのようだ。

おそらくこれが、最初の2つに次ぐ、試す価値のある最も理にかなった第3のアプローチなのだろう。今後数ヶ月のうちに、これらをもっといろいろといじってみる時間が取れることを、心から願っている。

この記事は「muse-spark-1.2-contributor」を使用して翻訳されました。

コメント