技術 約7分で読めます

397Bパラメータモデルを48GB MacBookで動かすFlash-MoE

Flash-MoEはQwen3.5-397B-A17B(4970億パラメータのMixture of Expertsモデル)をMacBook Pro M3 Maxで実行するC/Metal推論エンジンだ。48GBの統合メモリ環境で4.36トークン/秒以上を達成し、ツール呼び出しを含む本番品質の出力を生成できる。

GitHubで公開されており、実際に動くコードと詳細な58回の実験ログが含まれている。単に「動いた」という報告ではなく、何が効いて何が効かなかったかの試行錯誤の記録としても面白い。

モデルアーキテクチャ

Qwen3.5-397B-A17BはAlibaba発のMoE(Mixture of Experts)モデルだ。「397B」が総パラメータ数、「A17B」が1トークン処理時に実際に活性化されるパラメータ数(約170億)を表す。全パラメータを常時使うのではなく、入力に応じて必要なエキスパートだけを選択的に呼び出す設計がMoEの特徴だ。

このモデルの60層トランスフォーマーのうち:

  • 45層: GatedDeltaNet(線形注意)
  • 15層: 標準的なフル注意機構

各層に512個のエキスパートが配置され、1トークンごとに4個が活性化される。モデル全体の重みは4ビット量子化で209GBに収まるが、これは48GBの統合メモリには入り切らない。そこでSSDをフル活用する構成になっている。

コア技術: SSDエキスパートストリーミング

Flash-MoEの中核となる設計は「必要なエキスパートだけをSSDから随時読み込む」というアプローチだ。

各エキスパートの重みサイズ: 約6.75MB(4ビット量子化時)
1トークンで読み込む重み: 4個 × 6.75MB ≈ 27MB
SSD読み込み速度(実測): 17.5GB/s

実装ではpread()を並列呼び出しすることで複数エキスパートを同時取得する。OSのページキャッシュに任せることで「キャッシュヒット率が自然に約71%」に達し、カスタムキャッシュを別途実装する必要がなかった。

非エキスパート重み(全層共有の部分)は5.5GBでメモリに常駐させ、推論時のスクラッチ領域として200MBを確保するため、合計メモリ使用量は約6GB。OOMリスクはほとんどない。

Metal Compute Shaders

MacのGPUを使う計算にはAppleのMetalシェーダーを手書きで実装している。

カーネル内容
4ビット行列ベクトル積SIMD縮約・タイル処理・入力キャッシュ共有
2ビット行列ベクトル積同上(速度優先だが量子化誤差大)
RMS正規化2パス処理(二乗和縮約 → 適用)
GPU注意機構Q@K^T・ソフトマックス・scores@V
MoE結合エキスパート出力のWeighted sum

4ビット行列ベクトル積にはFMA(Fused Multiply-Add)の最適化が効いている。単純な実装では (nibble * scale + bias) * x という演算順だが、これを fma(nibble, scale*x, bias*x) に変形することでGPUのFMAユニットを活用し、12%の高速化を実現した。

Metal Performance Shaders(MPS)はCUDAに相当するAppleのGPUコンピューティングフレームワークだ。CUDAがNVIDIA GPU専用であるのに対し、MPSはApple Siliconの統合GPUで動作する。Flash-MoEはMetalシェーダーを直接実装することでMPSレベルの最適化を実現している。

GatedDeltaNetの高速化

45層に使われるGatedDeltaNetは線形注意機構の一種で、従来のTransformerより計算量が少ない。各層で「64ヘッド × 128×128の状態行列」を更新する再帰的な計算が中心だ。

この計算にはApple AccelerateライブラリのBLAS関数(cblas_sscalcblas_sgemvcblas_sger)を使い、スカラー実装と比較して64%の高速化を達成した。

遅延GPU実行(パイプライン最適化)

レイヤーごとの平均処理時間4.28msの内訳:

処理時間
GPU処理(注意/正規化等)1.22ms + 0.55ms
CPU処理(GatedDeltaNet)0.013ms
SSD I/O(エキスパート読み込み)2.41ms

I/Oが支配的なため、「GPU処理(CMD3: エキスパート順伝播)を非同期に投入し、CPUが次レイヤーの準備をしている間にGPUを走らせる」という遅延実行パターンを採用している。ただしApple SiliconはSSD DMAとGPU計算がメモリコントローラーを共有するため、過度な並列化はかえって遅くなる。シリアルなパイプラインが最適という結論になっている。

パフォーマンス結果

構成スループットツール呼び出しモデルサイズ
4ビット(FMA最適化)4.36 tok/s安定209GB
4ビット(基本実装)3.90 tok/s安定209GB
2ビット5.74 tok/s不安定120GB

2ビット量子化は速度で上回るが、JSON出力が破損する問題があった。具体的には "name"\name\ に変化し、ツール呼び出しのパース処理が失敗する。本番用途には4ビット構成を推奨している。

廃棄された最適化(58回の実験から)

試してうまくいかなかったアプローチの記録も公開されている。

アプローチ結果理由
LZ4圧縮-13%(悪化)解凍オーバーヘッドが読み込み削減を上回る
メモリマップ(mmap)-5倍(大幅悪化)ページフォルトのオーバーヘッドが高い
プリフェッチ-73%(大幅悪化)統合メモリの帯域幅制約で逆効果
予測的エキスパートルーティング精度31%(無効)事前予測の精度が低すぎる

特にプリフェッチの結果は興味深い。通常のNVMe SSDであれば有効な手法だが、Apple Siliconの「SSD・GPU・CPU・NPUが統合メモリバスを共有する」アーキテクチャでは、先読みが他の処理とバンド幅を奪い合い逆効果になった。

実行方法

cd metal_infer
make
./chat                                          # ツール呼び出し対応チャット
./infer --prompt "説明してください" --tokens 100  # 単発推論
./infer --prompt "Hello" --tokens 20 --timing   # タイミング計測付き

必要環境はMacBook Pro M3 Max(48GB統合メモリ、1TB Apple Fabric SSD)でmacOS 26.2以上。モデル重みはHugging FaceからQwen3.5-397B-A17Bを事前ダウンロードしておく必要がある。

オフロード戦略の比較: レイヤー/テンソル vs SSDストリーミング

ローカルLLM推論で「モデルがメモリに入り切らない」問題に対するアプローチは、大きく2系統ある。

BERT+Qwen OCR校正ツールの記事でllama.cppのオフロード機能を試した。llama.cppには --n-gpu-layers によるレイヤー単位のCPU/GPU振り分けと、--override-tensor によるテンソル単位の振り分けがある。FFN(フィードフォワード)テンソルをCPUに逃がしてattentionだけGPUに載せる手法で、Redditでは200%以上の速度向上が報告されている。

ただしこのアプローチは「モデル全体がCPUメモリ+VRAMの合計に収まる」ことが前提だ。RTX 3050 Ti(4GB)+ メインメモリ16GBの環境では4B〜9Bモデルが限界で、しかもA/B判定のような短いプロンプトではGPU並列計算の恩恵が小さく、テンソルオフロードで差が出なかった。

Flash-MoEのアプローチはこれとは根本的に異なる。209GBのモデルはCPUメモリにもVRAMにも全く入り切らないので、SSDを「超低速だが超大容量のVRAM」として使う発想だ。

llama.cpp オフロードFlash-MoE SSDストリーミング
対象メモリに収まるモデル(〜数十GB)メモリに収まらないモデル(209GB)
GPU活用レイヤー/テンソル単位で振り分けMetalシェーダーで全GPU処理
ボトルネックCPU-GPU間のデータ転送SSD読み込み帯域(17.5GB/s)
キャッシュなし(全重みがメモリ常駐)OSページキャッシュ(ヒット率71%)
MoE活用汎用(Dense/MoE問わず)MoE専用(活性化4/512の疎性が前提)

Flash-MoEがSSDストリーミングで成立するのは、MoEアーキテクチャの疎性(1トークンで512個中4個だけ活性化)があるからだ。Denseモデルの397Bパラメータを同じ方式で動かすのは不可能で、毎トークン全重みを読み込むとSSD帯域が完全にボトルネックになる。逆にllama.cppのレイヤーオフロードはDenseモデルでも使えるが、メモリに収まるサイズが上限になる。

Apple Siliconの統合メモリはこの2つのアプローチそれぞれに異なる影響を与えている。llama.cppのオフロードでは、CPU/GPUが同じメモリ空間を共有するため「オフロード」の概念自体が薄れる。OCR校正ツールをApple Silicon(M1 Max 64GB)で動かしたときは、BERTとLLMの両方がGPU上で自然に動き、CUDA環境で必要なVRAM管理の手間がなかった。一方Flash-MoEでは、統合メモリバスをSSD DMAとGPU計算が共有するため、プリフェッチが逆効果になるという統合メモリ特有の制約が発生している。同じ統合メモリアーキテクチャが、片方では恩恵、もう片方では制約として現れる。