2.78兆パラメータのLLMを8GB RAMのシングルCPUで実行する

ポータブルなC99と巧妙なメモリストリーミングを使用して、2.78兆パラメータのKimi K3モデルがわずか8.24 GBのRAMでシングルCPU上で動作する方法を発見します。

AIの世界はスケールに夢中ですが、もしあなたがすでに所有しているマシンで2.78兆パラメータのモデルを実行できたらどうでしょうか?それを実現するのがkimi-k3-in-cプロジェクトです。2.78TパラメータのKimi K3モデルが、GPUもBLASもフレームワークもなしに、シングルCPU上でわずか8.24 GBのRAMで推論を実行します。エンジン全体は176 KBのC99バイナリで、モデルをディスクからストリーミングするため、最先端のAIを最先端とは言えないハードウェアで実行できます。

これはおもちゃでも、大幅に量子化された近似でもありません。出力は、8 GB RAMのラップトップで実行しても、224 GBのワークステーションで実行しても、バイト単位で同一です。唯一の違いは速度です。低スペックでは1トークンあたり32.69秒、高スペックでは1トークンあたり19.21秒です。このプロジェクトは、巧妙なエンジニアリング、モデルアーキテクチャの深い理解、そしてLLM推論に必要なものについての従来の前提に挑戦する意欲の組み合わせによってこれを達成しています。

問題:収まらないモデル

Kimi K3は、2.78兆パラメータの混合専門家(MoE)モデルで、1.56 TBのチェックポイントとして提供されています。コンシューマーマシンではそれをメモリに保持できません。素朴なアプローチ(すべてをRAMにロードする)では、bfloat16精度で5.56 TBが必要になります。これは非現実的なだけでなく、データセンターを持たない人には不可能です。

重要な洞察は、MoEモデルはすべてのパラメータを同時にアクティブにする必要がないということです。任意のトークンに対して、各層の896の専門家のうち16だけが発火します。これは2.78兆のうち約1040億のアクティブパラメータ、わずか3.7%です。残りは、必要になったときにアクセス可能であれば、ディスク上に置いておくことができます。

4つの削減

このプロジェクトは、4つの主要な削減を通じてメモリフットプリントを達成しています:

  1. MXFP4専門家:ルーティングされた専門家は4ビット浮動小数点形式で提供され、bfloat16の2バイトではなく、重みあたりわずか0.53125バイトしか占有しません。これだけで専門家の重みが5.45 TBから1.447 TBに削減されます。
  2. KDAアテンション:Kimi Delta Attentionは、コンテキスト長に応じて成長しない固定サイズのリカレント状態を使用します。つまり、93層のうち69層は、テキストをどれだけ入力してもメモリフットプリントが一定です。
  3. MLAアテンション:Multi-head Latent Attentionは、96の個別のキー/バリューヘッドの代わりに、位置ごとに単一の576次元の潜在変数をキャッシュし、KVキャッシュを53倍削減します。
  4. トランクのストリーミング:高密度層(「トランク」)は、ピン留めされたプレフィックスと単一のリングバッファを使用して、ディスクから層ごとにストリーミングされます。これにより、メモリ要件が固定された下限から調整可能なダイヤルに変わります。

仕組み:アーキテクチャ

エンジンはポータブルなC99で書かれており、libmとOpenMP以外の外部依存関係はありません。コードベースは非常に小さく、6つのCファイルが単一の176 KBバイナリにコンパイルされます。これは、プロジェクトがフレームワークのオーバーヘッドを回避し、すべてのカーネルをゼロから実装しているため可能です。

チェックポイントの読み取り

1.56 TBのチェックポイントは、96個のsafetensorsファイルで構成されています。エンジンはJSONヘッダーを読み取って497,220個すべてのテンソルのインデックスを構築し、必要なバイトだけをオンデマンドで読み取ります。これはO_DIRECT読み取りで行われ、ページキャッシュを完全にバイパスします。プロジェクトは、テストハードウェアではバッファリングされた読み取りよりも高速であることを発見しました。

専門家キャッシュ

ルーティングされた専門家は、設定可能なサイズのLRUキャッシュにロードされます。しかし、プロジェクトは驚くべき結果を発見しました:専門家キャッシュは小さなサイズではほとんど役に立ちません。モデルのQuantile Balancingトレーニングにより、専門家の使用はプール全体で平坦化されるため、キャッシュが活用できるホットなサブセットがありません。キャッシュが効果を発揮し始めるのは、アリーナが約36 GBを超えてからです。

トランクのストリーミング

高密度トランク(108.81 GB)は層ごとに読み取られます。エンジンは、予算に収まるだけの層をピン留めし、残りを単一のリングバッファを通じてストリーミングします。エンジンは固定順序で層を処理するため、ピン留めされたプレフィックスはN/93の決定的なヒット率を達成し、循環スキャンでのLRUキャッシュよりもはるかに優れています。

検証:正しさの証明

このプロジェクトは単に動作することを主張するだけでなく、厳密な検証ラダーを通じてそれを証明しています:

  • ウェイトレステスト:実際のモデルと同じテンソルグラフを持つ13層のオラクルモデルを、PyTorchリファレンスと照合します。3つの実行パス(ティーチャーフォーシング、グリーディデコード、インクリメンタルデコード)すべてが同一のトークンIDを生成します。
  • 層の適合性:フルモデルの全93層がPyTorchリファレンスに対して検証され、最悪の場合の誤差は丸め予算の0.00倍です。
  • ロジットのパリティ:Cエンジンのロジットは、163,840語彙の出力全体でtorchリファレンスと要素ごとに一致し、最大差は7.87e-6です。
  • メモリラダー:8 GBから224 GBまでの12の異なるメモリ予算すべてが、バイト単位で同一のトークンIDを生成します。

パフォーマンス:期待できること

シングルCPU(AMD EPYC 7763、124コア)で、エンジンは以下を達成します:

  • 8 GB RAM:32.69秒/トークン
  • 32 GB RAM:31.44秒/トークン
  • 64 GB RAM:28.60秒/トークン
  • 128 GB RAM:29.40秒/トークン(ノイズに注意)
  • 224 GB RAM:19.21秒/トークン

パフォーマンスは非常にI/Oバウンドです。ウォールクロック時間の41%から61%はディスクの待機に費やされます。つまり、CPUよりもストレージデバイスの方が重要です。高速なNVMeドライブが良好なパフォーマンスには不可欠です。

はじめに

これを自分で実行するには、以下が必要です:

  • AVX2とFMAを備えたLinux x86-64
  • 少なくとも8 GBのRAM
  • 約1.7 TBの空きディスク容量
  • GCC ≥ 9またはClang ≥ 10
  • Python 3.9以上(ダウンロードおよびパッキングツール用)

リポジトリをクローンし、make -jでビルドし、make testを実行して、1.56 TBのチェックポイントをダウンロードする前にエンジンが動作することを確認します。次に、モデルをダウンロードし、トランクをパックして、以下を実行します:

./bin/k3 ~/k3model --trunk ~/k3trunk --preset laptop \
  --tok ~/k3model --prompt "The capital of France is" --gen 8 --incremental

結論

kimi-k3-in-cプロジェクトは、システムエンジニアリングのマスタークラスです。適切なアプローチを使えば、最大のAIモデルでも一般のハードウェアでアクセス可能にできることを示しています。重要なポイントは次のとおりです:

  • メモリはダイヤルであり、下限ではない:トランクをストリーミングすることで、速度とメモリ使用量をトレードオフできます。
  • モデルを理解する:アーキテクチャの選択(KDA、MLA、MXFP4)を最大限に活用しています。
  • すべてを測定する:プロジェクトの厳密な検証と測定方法論により、すべての主張がデータに裏付けられています。

これは単なる技術的な好奇心ではなく、AIをよりアクセスしやすくするための青写真です。自分のマシンで1兆パラメータのモデルを実行したいと思ったことがあるなら、このプロジェクトはそれが可能であることを示しています。そして、それはオープンソースです。

ソース

FareedKhan-dev/kimi-k3-in-c: 8.24 GBのRAMでシングルCPU上で推論を実行する2.78兆パラメータのKimi K3。ポータブルC99:BLASなし、フレームワークなし、GPUなし。