Edge0、学習型プリルーティングで35B MoEをSSDから実行

導入

大規模なMixture-of-Experts（MoE）モデルをコンシューマー向け機器で動かす際、問題は演算性能だけではない。MoEは各トークンで一部の専門家だけを選ぶため計算量を抑えられるが、モデルに含まれる専門家の重みをすべて保存する必要は残る。素材の例では、35B級モデルの4-bit重みが約19.5GBに達し、多くの機器にとってメモリ容量が大きな制約になる。

Edge0は、専門家重みの一部をSSDに置き、計算と読み出しを重ね合わせることでこの問題に対応する。報告によれば、単一の24GBマシン上で35B級MoEを20 token/sで動かし、ピーク時のアクティブメモリを3GiB未満に抑える。5つの公開ベンチマークでは、FP16の教師モデルから平均で数ポイント以内の性能とされている。

中核となる仕組み

• SSDオフロードには時間的な依存関係がある。 通常、次の層で必要な専門家は、前の層の出力が得られるまで確定しない。実際のルーティングを待ってからSSDを読むと、読み出しを計算に隠す時間が不足する。
• プリルーターが次層を先読みする。 各層に設けた軽量な予測ヘッドが、1トークン先の専門家選択を予測し、重みのステージングを早める。
• 予測をそのままルーティングに使う。 予測後に元のルーターと照合し、間違いを修正する方式ではない。予測された専門家を直接実行するため、「100%精度」という説明は、予測と実行が内部的に一致するという意味であり、元のルーターへの一致率が100%という意味ではない。
• リカバリーLoRAで品質を補う。 4-bit量子化とルーティングの置き換えによる劣化を補償するため、実際の学生経路で学習した、重みに統合しないLoRAを利用する。

意義と残された課題

報告結果が再現できれば、Edge0はエッジ推論の設計を変える可能性がある。単にモデルを小さくするのではなく、SSD帯域、専門家のスケジューリング、ルーティング、量子化後の適応を一体のシステムとして扱うからだ。RAMが限られていても高速なNVMe SSDを持つ機器なら、これまで保持できなかった規模のモデルを動かせる可能性がある。

ただし、中心的なトレードオフは明確に評価する必要がある。素材には、層ごとの予測統計、誤予測の裾野、誤った専門家を読み込むコスト、SSDや系列長ごとの遅延分布が示されていない。予測を直接実行する構成では、誤りが元のルーターへの自動フォールバックで消えるわけではなく、スループットを維持したまま品質低下として現れる可能性がある。平均スコアや平均生成速度だけでなく、初回トークン遅延、トークンごとの揺らぎ、長文での安定性、SSD帯域の圧迫も重要になる。

Edge0の価値は、SSD利用だけでMoEの問題を解決したことではなく、学習型ルーティング予測、ストリーミング読み出し、量子化リカバリー、低メモリ実行を組み合わせた点にある。フレームワーク、チェックポイント、アダプターはオープンソースとされるため、今後は実装と詳細な測定による検証が焦点となる。

出典：Hugging Face Daily Papers