Edge0:把 MoE 專家留在 SSD,讓 35B 模型以約 3 GiB 活躍記憶體推論
Edge0 以 SSD expert offload、提前預測路由的 prerouter,以及補回 int4 與路由替換損失的 Recover-LoRA,將 35B sparse MoE 的峰值活躍記憶體壓到約 2.9 GiB。這不等於整個模型只有 3 GB:checkpoint 仍約 19.5–23 GB,長上下文還會增加 KV cache,而且跨平台引擎仍屬 preview。
大型 MoE 每個 token 只啟用少數專家,但傳統推論仍常把全部權重放進記憶體。Edge0 的做法是讓大部分 expert 權重留在 SSD,需要時才 mmap/stream 進來,再用預測式路由把讀取延遲藏在運算後面。
先校正「35B 只要 3GB」
2.9 GiB 指的是短上下文測試下的 peak active memory,不是模型檔案總大小,也不是整台裝置只需要 3 GB RAM。
- 論文中的 35B int4 checkpoint 約 19.5 GB;官方 README 的實際下載說明約 23 GB。
- 仍需預留作業系統、tokenizer、runtime 與 KV cache;上下文越長,KV cache 越大。
- 官方主要 benchmark 使用 24 GB Mac mini M4 Pro,不是 3 GB RAM 裝置。
三個核心機制
1. SSD expert offload
Edge0 將量化後的 expert 權重以 safetensors 留在儲存裝置,透過 mmap 與 bounded streaming pool 只載入目前路由到的專家。MoE 的峰值記憶體因此由「總參數量」改為接近「活躍專家集合」。
35B tier 有 40 層、256 個 experts,K=4;每個 token 活躍約 3B 參數。SSD 並沒有消除權重,只是成為較低層的 weight tier。
2. Prerouter:提前一個 token 預測下一層專家
單純 on-demand offload 會在每層等待 SSD I/O。Edge0 訓練小型 prerouter,提前一個 token 預測下一層會使用的 experts,並把預測結果直接當成實際 routing,讓載入與當前 forward pass 重疊。
代價是原始模型的 routing 被替換,預測錯誤不再由 runtime fallback 修補,而是透過訓練把品質損失吸收掉。
3. Recover-LoRA:補回 int4 與 routing replacement 的能力
4-bit 量化和替代 routing 都會降低品質。Edge0 冻結 int4 base,以 teacher distillation 訓練 LoRA,部署時保持 unmerged,平行加回 LoRA delta。論文指出若先 merge 再重新量化,許多細小 LoRA 變化會被 4-bit 步階吃掉。
官方測試結果怎麼讀
| 項目 | 35B tier | 8B tier |
|---|---|---|
| 活躍參數/token | 約 3B | 約 1.2B |
| 官方 README decode | 14.9–17.7 tok/s | 23.9–25.3 tok/s |
| 論文 v2 配對測量 | 20.4 tok/s | 28.0 tok/s |
| 峰值活躍記憶體 | 2.9 GiB | 1.0–1.5 GiB,依設定而異 |
| 測試機 | Mac mini M4 Pro,24 GB | |