The Other Half of the Memory Wall: Menyajikan MoE 35B dari SSD dengan Prediksi Routing Terlatih

Panduan membaca

Masalah dan motivasi

Model Mixture-of-Experts (MoE) menskalakan total parameter sekaligus mempertahankan parameter aktif yang rendah per token (misalnya, mengaktifkan ~3B parameter per token pada model 35B). Namun, inferensi pada perangkat keras tingkat konsumen tetap dibatasi oleh memori bobot statis: MoE 35B memerlukan sekitar 19,5 GB bahkan setelah dikuantisasi ke int4. Pada mesin 24 GB, hal ini sepenuhnya menyita ruang untuk sistem operasi dan KV cache, yang memicu memory paging parah serta penurunan throughput decode yang drastis (misalnya, 3,9 tok/s untuk model yang sepenuhnya menetap di memori pada M4 Pro Mac mini).

Melakukan offloading bobot pakar ke SSD dapat mengurangi konsumsi memori, tetapi pemuatan langsung (on-demand) yang naif memicu leher botol serialisasi yang kaku: pemilihan pakar untuk lapisan $N+1$ bergantung pada luaran dari lapisan $N$. Pembacaan SSD tidak dapat dimulai cukup awal untuk bersembunyi di balik komputasi, menghentikan eksekusi karena latensi disk satu kali per lapisan per token.

Ide-ide utama

• SSD sebagai Tingkatan Penyimpanan Bobot: Bobot pakar MoE tetap berada di SSD dalam format int4 dan dipetakan ke memori (mmap) ke dalam streaming pool yang ringkas. Memori puncak ditentukan oleh working set aktif, bukan total ukuran parameter model.
• Prediction as Routing: Alih-alih menunggu lapisan $N$ selesai, sebuah prerouter per lapisan memprediksi routing lapisan $N+1$ satu token penuh sebelumnya. Yang terpenting, model memperlakukan prediksi tersebut sebagai routing itu sendiri pada saat decode. Ini secara inheren menghilangkan risiko token terbuang (dropped tokens) atau pemuatan cadangan (fallback loads) karena himpunan yang disiapkan identik dengan himpunan yang diarahkan.
• Unmerged Recovery LoRA: Kuantisasi 4-bit dan routing aproksimasi menyebabkan penurunan kualitas. Edge0 melatih adapter berperingkat rendah (low-rank adapter) pada jalur routing student melalui distilasi. Karena menggabungkan LoRA ke dalam basis int4 lalu menguantisasinya kembali akan menghapus sebagian besar nilai delta (hanya tersisa retensi 18% pada tingkat logits), Edge0 menyajikan adapter sebagai delta paralel tanpa penggabungan (unmerged) dengan overhead yang sangat kecil.

Metode

Arsitektur Head Prerouter

Setiap head prerouter yang dimiliki oleh lapisan $N$ berjalan pada token $t$ pada luaran layer norm pasca-atensi dan memprediksi routing untuk lapisan $N+1$ pada token $t+1$. Fitur masukan menggabungkan hidden state dengan dua vektor one-hot routing top-$k$ (token saat ini dan sebelumnya):

$$\text{dim}=d_{\text{model}}+2E$$

Untuk tingkat 35B, nilainya adalah $2048+2\times 256=2560$. Head menghitung:

$$\text{head}(x)={\mathrm{fc}}_2(\mathrm{erf}\text{-}\mathrm{gelu}({\mathrm{fc}}_1(x)))+\ell (x)$$

di mana $\ell (x)$ adalah jalur residual linier yang diinisialisasi (warm-started) dari bobot router lapisan berikutnya.

Eksekusi Adapter Tanpa Penggabungan (Unmerged)

Lapisan MoE yang disajikan mempertahankan bita basis terkuantisasi dan menjalankan delta LoRA secara paralel:

$$y=W_{int4}(x)+\frac{\alpha }{r}BAx$$

Pelatihan Tiga Fase

1. Head Distillation: Head prerouter dilatih untuk meniru keputusan router sebenarnya pada lapisan berikutnya dari teacher, sementara sisa jaringan dibekukan (frozen).
2. Student-Path SFT: Adapter LoRA pemulihan (dipasang pada proyeksi attention, linear-attention, dan shared-expert, tetapi tidak pada pakar streaming) dilatih menggunakan teks hasil generasi teacher sambil memaksakan routing melalui prerouter.
3. On-Policy Distillation: Student menghasilkan teks di bawah eksekusi yang diarahkan prerouter, teacher basis fp16 menilai generasi tersebut, dan parameter LoRA dioptimalkan melalui divergensi Kullback–Leibler terbalik dengan objektif probabilitas top-$k$ yang peka terhadap ekor distribusi (tail-aware).

Hasil

• Throughput & Memori (Mac mini M4 Pro, 24 GB):
  • edge0-35b ($K=4$): Mencapai decode 20,4 tok/s hanya menggunakan memori aktif puncak 2,9 GiB dari checkpoint int4 19,5 GB (dibandingkan dengan 3,9 tok/s dan 18,2 GiB saat sepenuhnya dimuat ke memori).
  • edge0-8b ($K=8$): Mencapai decode 28,0 tok/s dalam memori puncak 1,5 GiB.
• Peningkatan Kecepatan Prefetching (MacBook M2, 16 GB):
  • Dibandingkan dengan streaming on-demand yang murni reaktif, prefetch prerouter 1 token meningkatkan kecepatan decode sebesar +80% hingga +84% di berbagai variasi lebar routing ($K=2$: 4,8 ke 8,6 tok/s; $K=4$: 3,5 ke 6,4 tok/s; $K=8$: 1,8 ke 3,3 tok/s).
  • Waktu muat yang memblokir thread utama (main-thread blocked load time) turun drastis (misalnya, dari 244,0 ms menjadi 101,9 ms per langkah pada $K=4$).
• Kualitas Model (OpenCompass):
  • Di lima pengujian tolok ukur (AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench), edge0-35b mencetak skor rata-rata 79,2 dibandingkan 83,2 untuk basis Qwen3.6 fp16 (selisih 3,9 poin).
  • edge0-8b mencapai 69,9 dibandingkan 72,7 untuk basis Ling 3.0 tiny fp16 (selisih 2,8 poin).

Batasan dan pertanyaan terbuka

• Konkurensi Penyajian: Edge0 beroperasi secara serial FIFO untuk permintaan tunggal; batching mengubah dinamika working set pakar dan caching dengan cara yang belum tercakup dalam profil single-stream saat ini.
• Overhead Pembentukan Graph CPU: Kecepatan decode dibatasi oleh overhead kerangka kerja di sisi CPU (~44 ms/langkah dalam pembentukan graph MLX di seluruh 40 lapisan), yang tidak dapat diringankan oleh optimasi SSD.
• Penurunan Kualitas pada Penalaran: Penurunan kinerja terkonsentrasi pada tolok ukur penalaran yang kompleks (misalnya, penurunan 6,1 poin pada AIME untuk 35B dan 10,0 poin untuk 8B).
• Ketergantungan Perangkat Keras: Keuntungan throughput relatif dari prerouter bergantung pada latensi disk dan akan mengecil ketika cache bobot sudah hangat (warm) atau saat berjalan pada tingkatan media penyimpanan yang sangat cepat.