Nửa còn lại của bức tường bộ nhớ: Phục vụ mô hình MoE 35B từ SSD với dự đoán định tuyến được huấn luyện

Hướng dẫn đọc

Vấn đề và động lực

Các mô hình Mixture-of-Experts (MoE) mở rộng tổng số tham số trong khi vẫn giữ số tham số kích hoạt ở mức thấp trên mỗi token (ví dụ: kích hoạt ~3B tham số mỗi token trong mô hình 35B). Tuy nhiên, việc suy luận trên phần cứng phổ thông vẫn bị hạn chế bởi bộ nhớ trọng số tĩnh: một MoE 35B đòi hỏi khoảng 19,5 GB ngay cả khi được lượng tử hóa sang int4. Trên một máy 24 GB, điều này chiếm trọn không gian của hệ điều hành và KV cache, dẫn đến hiện tượng phân trang bộ nhớ (memory paging) nghiêm trọng và làm giảm thông lượng giải mã (ví dụ: chỉ đạt 3,9 tok/s đối với mô hình nằm hoàn toàn trong RAM trên Mac mini M4 Pro).

Việc offload trọng số expert sang SSD giúp giảm dung lượng bộ nhớ, nhưng việc tải theo yêu cầu (on-demand) theo cách thông thường lại tạo ra điểm nghẽn tuần tự hóa cứng: việc lựa chọn expert cho tầng $N+1$ phụ thuộc vào đầu ra của tầng $N$. Các thao tác đọc SSD không thể bắt đầu đủ sớm để ẩn sau quá trình tính toán, khiến quá trình thực thi bị đình trệ do độ trễ đĩa một lần trên mỗi tầng cho mỗi token.

Ý tưởng cốt lõi

• SSD đóng vai trò là tầng lưu trữ trọng số (Weight Tier): Trọng số của các expert MoE được giữ trên SSD ở định dạng int4 và được ánh xạ bộ nhớ (mmap) vào một vùng đệm streaming nhỏ gọn. Bộ nhớ đỉnh được xác định bởi tập làm việc (working set) đang hoạt động thay vì tổng kích thước tham số của mô hình.
• Dự đoán làm định tuyến (Prediction as Routing): Thay vì chờ tầng $N$ hoàn thành, một prerouter trên mỗi tầng sẽ dự đoán định tuyến của tầng $N+1$ trước hẳn một token. Quan trọng nhất, mô hình xem dự đoán này chính là định tuyến thực tế tại thời điểm giải mã. Điều này loại bỏ hoàn toàn việc bỏ rơi token (dropped tokens) hoặc tải dự phòng (fallback loads) theo thiết kế, vì tập hợp được chuẩn bị sẵn hoàn toàn trùng khớp với tập hợp được định tuyến.
• Recovery LoRA chưa hợp nhất (Unmerged Recovery LoRA): Lượng tử hóa 4-bit và định tuyến xấp xỉ gây ra suy giảm chất lượng. Edge0 huấn luyện một adapter low-rank trên đường dẫn định tuyến của student thông qua kỹ thuật chưng cất (distillation). Vì việc hợp nhất LoRA vào mô hình gốc int4 rồi lượng tử hóa lại sẽ xóa sạch hầu hết phần delta (chỉ còn giữ lại 18% ở cấp độ logits), Edge0 phục vụ adapter dưới dạng một nhánh delta song song chưa hợp nhất với chi phí tính toán không đáng kể.

Phương pháp

Kiến trúc Prerouter Head

Mỗi prerouter head thuộc sở hữu của tầng $N$ chạy tại token $t$ trên đầu ra layer norm sau attention và dự đoán định tuyến cho tầng $N+1$ tại token $t+1$. Các đặc trưng đầu vào nối ghép trạng thái ẩn với hai vector định tuyến one-hot top-$k$ (của token hiện tại và token trước đó):

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

Đối với phiên bản 35B, kích thước này là $2048+2\times 256=2560$. Head tính toán:

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

trong đó $\ell (x)$ là đường dẫn phần dư tuyến tính (linear residual path) được khởi động từ trọng số router của tầng tiếp theo.

Thực thi Adapter chưa hợp nhất

Tầng MoE được phục vụ vẫn giữ nguyên các byte gốc đã lượng tử hóa và chạy delta LoRA song song:

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

Huấn luyện ba giai đoạn

1. Chưng cất Head (Head Distillation): Các prerouter head được huấn luyện để bắt chước quyết định của router thực tế ở tầng tiếp theo của teacher trong khi vẫn đóng băng phần còn lại của mạng.
2. SFT trên đường dẫn Student (Student-Path SFT): Các recovery LoRA adapter (được gắn vào các phép chiếu attention, linear-attention và shared-expert, nhưng không gắn vào các streaming expert) được huấn luyện bằng văn bản do teacher tạo ra, đồng thời buộc định tuyến phải đi qua prerouter.
3. Chưng cất On-Policy (On-Policy Distillation): Student tạo văn bản dưới sự thực thi của prerouter, teacher gốc fp16 chấm điểm các văn bản được tạo, và các tham số LoRA được tối ưu hóa thông qua phân kỳ Kullback–Leibler ngược với mục tiêu xác suất top-$k$ có xét đến phần đuôi (tail-aware).

Kết quả

• Thông lượng & Bộ nhớ (Mac mini M4 Pro, 24 GB):
  • edge0-35b ($K=4$): Đạt tốc độ giải mã 20,4 tok/s chỉ với 2,9 GiB bộ nhớ hoạt động đỉnh từ checkpoint int4 19,5 GB (so với 3,9 tok/s và 18,2 GiB khi nạp toàn bộ vào RAM).
  • edge0-8b ($K=8$): Đạt tốc độ giải mã 28,0 tok/s trong mức bộ nhớ đỉnh 1,5 GiB.
• Gia tốc nạp trước (MacBook M2, 16 GB):
  • So với streaming theo yêu cầu thuần túy phản ứng (reactive), cơ chế prefetch bằng prerouter trước 1 token tăng tốc độ giải mã thêm +80% đến +84% trên các độ rộng định tuyến ($K=2$: từ 4,8 lên 8,6 tok/s; $K=4$: từ 3,5 lên 6,4 tok/s; $K=8$: từ 1,8 lên 3,3 tok/s).
  • Thời gian luồng chính bị nghẽn do nạp (blocked load time) giảm đáng kể (ví dụ: từ 244,0 ms xuống 101,9 ms mỗi bước ở $K=4$).
• Chất lượng mô hình (OpenCompass):
  • Trên năm benchmark (AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench), edge0-35b đạt điểm trung bình 79,2 so với 83,2 của mô hình gốc Qwen3.6 fp16 (khoảng cách 3,9 điểm).
  • edge0-8b đạt 69,9 so với 72,7 của mô hình gốc Ling 3.0 tiny fp16 (khoảng cách 2,8 điểm).

Hạn chế và câu hỏi mở

• Đồng thời khi phục vụ (Serving Concurrency): Edge0 hoạt động tuần tự hóa FIFO cho các yêu cầu đơn lẻ; việc xử lý theo lô (batching) sẽ làm thay đổi tập làm việc của expert và động lực học caching theo những cách mà các hồ sơ luồng đơn chưa giải quyết được.
• Chi phí xây dựng đồ thị của CPU: Tốc độ giải mã bị giới hạn bởi chi phí phụ trợ của framework phía CPU (~44 ms/bước khi xây dựng đồ thị MLX trên 40 tầng), điều mà các tối ưu hóa SSD không thể giảm bớt.
• Giảm chất lượng ở tác vụ suy luận: Sự suy giảm hiệu năng tập trung chủ yếu vào các benchmark suy luận phức tạp (ví dụ: giảm 6,1 điểm trên AIME đối với 35B và 10,0 điểm đối với 8B).
• Phụ thuộc vào phần cứng: Mức tăng thông lượng tương đối của prerouter phụ thuộc vào độ trễ của đĩa và sẽ giảm đi khi bộ nhớ đệm trọng số đã nóng (warm cache) hoặc khi chạy trên các tầng lưu trữ cực nhanh.