SSDを「メモリ階層」に使うColibrì、巨大MoEモデルを一般PCで推論

導入

超大規模言語モデルをローカルで動かす際、最大の壁は計算性能だけではありません。数千億から数兆規模の重みを、一般的なPCのRAMやGPUメモリに収められるかが問題になります。オープンソースのColibrìは、モデル全体を高速メモリに常駐させる代わりに、当面使わない重みをNVMe SSDへ置き、必要になった時点で読み込む方法を採用します。

Colibrìは純粋なC実装で、特定の推論エンジンに依存しません。提供された資料によれば、複数のモデル系列に対応しており、744BパラメータのGLM-5.2を使った実験を起点に開発されました。GitHub上では数万規模のStarを集めています。

MoEだから可能になる分割

この仕組みを支えるのがMixture-of-Experts、つまりMoEです。MoEモデルは全体として巨大でも、各トークンの生成時に全パラメータを計算するわけではありません。資料のGLM-5.2では、総パラメータ数が744Bである一方、1トークンあたりの有効パラメータは約40Bとされています。

Colibrìはモデルを複数の階層へ分けます。

• Attention、Embedding、共有エキスパートなどはRAMへ常駐させる。GLM-5.2のint4構成では約9.9GBです。
• 大量のルーティングエキスパートは主にSSDへ保存する。
• Routerが選んだエキスパートだけをRAMへ読み込み、GPUがあればVRAMも利用する。

これはモデル重みに対するJIT読み込みに近い考え方です。エキスパートがどの階層に置かれるかは速度に影響しますが、Routerの選択や重みそのものを変更するものではありません。

キャッシュと先読みが実用性を左右する

トークンごとにSSDから読み込めば、推論は極端に遅くなります。そこでColibrìはLRUキャッシュを使い、直近で利用したエキスパートをRAMに残します。さらに、実行中の利用回数を記録し、頻繁に選ばれるエキスパートを優先的に高速層へ置きます。

現在の層を計算している間に、次の層で使われそうなエキスパートを先読みする仕組みもあります。資料では、隣接層のルーティングに一定の相関があり、先読みの予測可能性は71.6％に達すると説明されています。SSDを2台使い、モデルの複製を分散して読み出す方法にも対応します。

プロジェクトのテスト結果では、RAM 25GBの環境でコールドキャッシュ時は約0.05～0.1 token/s、RAM 128GBのCPUのみの環境では約1.8 token/sとされています。6枚のRTX 5090を使い、多くのエキスパートを高速層に置いた場合は約5.8～6.8 token/sです。これらは特定条件でのプロジェクト側の数値であり、すべてのPCにそのまま当てはまるベンチマークではありません。

意義と限界

ColibrìはSSDをGPUメモリと同じ速度にする技術ではありません。SSDは容量、RAMはキャッシュ、VRAMは最も頻繁に使う重みの高速領域という役割分担を作ることで、超大規模モデルの試行に必要な初期投資を下げます。一方、起動時間、初回読み込み、持続的な生成速度は、SSDの帯域幅、RAM容量、エキスパートの再利用率、モデル構造に依存します。

したがって、Colibrìは高スループットの本番サーバーの代替というより、一般的なハードウェアで巨大モデルを検証するためのシステム設計として見るのが適切です。

出典：量子位