The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin; Yiming Wang; Runyuan Cai; Hanze Liu; Xiaodong Zeng

งานวิจัย · arXiv 2609.18063

อีกฟากหนึ่งของ Memory Wall: การรัน MoE ขนาด 35B จาก SSD ด้วยการทำนายเส้นทางที่ผ่านการเทรน

The Other Half of the Memory Wall: Serving 35B MoEs from SSD with Trained Routing Prediction

Yu Lin, Yiming Wang, Runyuan Cai, Hanze Liu, Xiaodong Zeng

เผยแพร่ 2026-09-16cs.AI arXiv PDF

อ่านพร้อมคำแปลและ AI

สรุปสามบรรทัด

🎯 MoE ระดับ 35B ต้องใช้พื้นที่ราว 19.5 GB แม้จะทำ int4 แล้วก็ตาม ซึ่งเกินกว่า RAM ระดับผู้บริโภค ในขณะที่การทำ offload ไปยัง SSD แบบทั่วไปต้องสะดุดกับความหน่วงของดิสก์เพราะการเลือกเส้นทางขึ้นกับเอาต์พุตของชั้นก่อนหน้า
🔧 Edge0 สตรีม expert ของ MoE จาก SSD โดยใช้ prerouter ประจำแต่ละชั้นเพื่อทำนายเส้นทางล่วงหน้า 1 โทเค็น โดยนำหลักการ 'ใช้การทำนายเป็นเส้นทางจริง' ควบคู่กับ unmerged parallel recovery LoRA ที่ผ่านการกลั่นกรองบนเส้นทาง student
📈 รัน MoE ขนาด 35B ได้ที่ความเร็ว ~20.4 tok/s ภายใต้การใช้หน่วยความจำสูงสุด 2.9 GiB บน Mac mini M4 Pro (24 GB) โดยมีคะแนนเฉลี่ยเบนช์มาร์กห่างจาก fp16 teacher เพียง 3.9 จุด

บทคัดย่อ

การทำ inference ของ Mixture-of-experts (MoE) บนฮาร์ดแวร์ระดับผู้บริโภคถูกจำกัดด้วยหน่วยความจำสำหรับเก็บน้ำหนัก: โมเดลระดับ 35B มีขนาด 19.5GB ที่ int4 และความกระจัดกระจาย (sparsity) ช่วยลดเฉพาะการคำนวณต่อโทเค็น แต่ไม่ได้ลดจำนวนไบต์ที่ต้องเก็บไว้ การทำ offloading ไปยัง SSD แบบทั่วไปไม่สามารถช่วยได้ด้วยตัวเอง เนื่องจาก expert ของชั้น $N + 1$ จำเป็นต้องถูกเลือกก่อนที่เอาต์พุตของชั้น $N$ จะเกิดขึ้น ทำให้การอ่านข้อมูลไม่สามารถเริ่มได้เร็วพอที่จะซ่อนอยู่เบื้องหลังการคำนวณ เรานำเสนอ Edge0 ซึ่งเป็นเอนจินการทำ inference ของ MoE แบบสตรีมมิ่งที่เชื่อมช่องว่างนี้ด้วย prerouter: ส่วนหัวระดับชั้นจะทำนายการเลือกเส้นทางของชั้นถัดไปล่วงหน้า 1 โทเค็น และการทำนายนี้จะถูกนำมาใช้เป็นการเลือกเส้นทางโดยตรง ส่งผลให้ชุด expert ที่ถูกเตรียมไว้ตรงกับชุดที่ถูกเลือกและไม่มีสิ่งใดตกหล่น นอกจากนี้ unmerged recovery LoRA ที่เทรนบนเส้นทางของ student จะช่วยชดเชยคุณภาพที่สูญเสียไปจากการทำควอนไทเซชันแบบ int4 และการแทนที่เส้นทาง บนเครื่องขนาด 24GB เพียงเครื่องเดียว Edge0 สามารถรัน MoE ขนาด 35B ได้ที่ความเร็ว 20 tok/s โดยใช้หน่วยความจำสูงสุดขณะทำงานไม่เกิน 3GiB และมีคะแนนเฉลี่ยบน 5 เบนช์มาร์กสาธารณะห่างจาก teacher แบบ fp16 เพียงไม่กี่จุด ทั้งนี้ โมเดลระดับ 8B ก็สามารถทำงานบนเฟรมเวิร์กเดียวกันได้ โดยเฟรมเวิร์ก เช็กพอยต์ และอะแดปเตอร์ทั้งหมดถูกเปิดให้ใช้งานแบบโอเพนซอร์ส

บทคัดย่อต้นฉบับ (ภาษาอังกฤษ)

Mixture-of-experts (MoE) inference on consumer hardware is bounded by weight memory: a 35B-class model is 19.5GB at 4-bit, and sparsity shrinks the compute per token, not the bytes that must be held. Naive offloading to SSD does not help on its own, because layer N+1's experts must be chosen before layer N's output exists, so the reads cannot start early enough to hide behind compute. We present Edge0, a streaming MoE inference engine that closes the gap with a prerouter: a per-layer head predicts the next layer's routing one token ahead, and the prediction is consumed as the routing itself, so the staged expert set equals the routed set and nothing is dropped. An unmerged recovery LoRA, trained on the student path, pays back the quality lost to int4 quantization and routing replacement. On a single 24GB machine, Edge0 serves a 35B MoE at 20tok/s inside 3GiB of peak active memory, within a few points of its fp16 teacher on average across five public benchmarks. An 8B tier runs on the same framework, and the framework, checkpoints, and adapters are open source.

คู่มือการอ่าน

ปัญหาและแรงจูงใจ

โมเดล Mixture-of-Experts (MoE) สามารถเพิ่มพารามิเตอร์รวมได้โดยยังคงจำนวนพารามิเตอร์ที่ทำงานต่อโทเค็นไว้ในระดับต่ำ (เช่น มีการเปิดใช้พารามิเตอร์ ~3B ต่อโทเค็นในโมเดลขนาด 35B) อย่างไรก็ตาม การทำ inference บนฮาร์ดแวร์ระดับผู้บริโภคยังคงถูกจำกัดด้วยหน่วยความจำสแตติกสำหรับน้ำหนัก: โมเดล MoE 35B ต้องใช้พื้นที่ราว 19.5 GB แม้จะผ่านการทำควอนไทเซชันเป็น int4 แล้วก็ตาม ในคอมพิวเตอร์ที่มี RAM 24 GB สิ่งนี้จะแย่งพื้นที่ของระบบปฏิบัติการและ KV cache จนหมด ส่งผลให้เกิด memory paging อย่างรุนแรงและทรูพุตในการดีโค้ดลดลงอย่างมาก (เช่น เหลือเพียง 3.9 tok/s สำหรับโมเดลที่โหลดไว้ในแรมทั้งหมดบน Mac mini M4 Pro)

การย้ายค่าน้ำหนักของ expert ไปไว้บน SSD ช่วยลดการใช้หน่วยความจำลงได้ แต่การโหลดตามความต้องการ (on-demand) แบบทั่วไปก่อให้เกิดคอขวดที่ต้องประมวลผลตามลำดับอย่างเลี่ยงไม่ได้: การเลือก expert สำหรับชั้น $N + 1$ จำเป็นต้องพึ่งพาเอาต์พุตจากชั้น $N$ ทำให้ไม่สามารถเริ่มอ่านข้อมูลจาก SSD ได้เร็วพอที่จะซ่อนอยู่เบื้องหลังเวลาการคำนวณ และส่งผลให้การประมวลผลต้องสะดุดรอความหน่วงของดิสก์ในทุกๆ ชั้นต่อหนึ่งโทเค็น

แนวคิดหลัก

SSD ในฐานะระดับชั้นของน้ำหนักโมเดล: ค่าน้ำหนักของ MoE expert จะถูกเก็บไว้บน SSD ในรูปแบบ int4 และถูกทำ memory-mapped (mmap) เข้าสู่ streaming pool ขนาดกะทัดรัด ทำให้การใช้หน่วยความจำสูงสุดถูกกำหนดโดยชุดข้อมูลที่กำลังทำงานอยู่ (working set) ไม่ใช่ขนาดพารามิเตอร์รวมของโมเดล
ใช้การทำนายเป็นเส้นทางจริง (Prediction as Routing): แทนที่จะรอให้ชั้น $N$ คำนวณเสร็จสิ้น prerouter ประจำแต่ละชั้นจะทำนายการเลือกเส้นทางของชั้น $N + 1$ ล่วงหน้าหนึ่งโทเค็นเต็ม และจุดสำคัญคือ โมเดลจะนำผลการทำนายนี้มาใช้ เป็น เส้นทางจริงในขั้นตอนการดีโค้ด วิธีนี้ช่วยกำจัดปัญหาโทเค็นตกหล่นหรือการโหลดสำรองฉุกเฉินได้ตั้งแต่ระดับการออกแบบ เนื่องจากชุด expert ที่ถูกสเตจเตรียมไว้จะตรงกับชุดที่ต้องใช้คำนวณอย่างแท้จริง
Unmerged Recovery LoRA: การทำควอนไทเซชัน 4 บิตและการประมาณเส้นทางนำมาซึ่งการสูญเสียคุณภาพ Edge0 จึงเทรน low-rank adapter บนเส้นทางการเลือก expert ของ student ผ่านการกลั่นกรอง (distillation) และเนื่องจากการรวม LoRA เข้ากับเบส int4 แล้วทำควอนไทเซชันซ้ำจะทำลายค่าความเปลี่ยนแปลงส่วนใหญ่ไป (เหลือการคงไว้ระดับ logits เพียง 18%) Edge0 จึงรันอะแดปเตอร์นี้เป็น delta แบบขนานโดยไม่ผสานรวม (unmerged) ซึ่งมี overhead น้อยมาก

ระเบียบวิธี

สถาปัตยกรรม Prerouter Head

Prerouter head แต่ละตัวของชั้น $N$ จะทำงานที่โทเค็น $t$ บนเอาต์พุตของ post-attention layer norm และทำนายการเลือกเส้นทางสำหรับชั้น $N + 1$ ที่โทเค็น $t + 1$ โดยฟีเจอร์อินพุตจะนำ hidden state มาต่อกับเวกเตอร์ one-hot ของ top- $k$ routing จำนวน 2 ชุด (โทเค็นปัจจุบันและโทเค็นก่อนหน้า):

$dim = d_{model} + 2 E$

สำหรับระดับ 35B ค่านี้คือ $2048 + 2 \times 256 = 2560$ โดยที่ head จะทำการคำนวณดังนี้:

$head (x) = fc_{2} (erf - gelu (fc_{1} (x))) + ℓ (x)$

โดยที่ $ℓ (x)$ คือ linear residual path ที่ warm-start มาจากค่าน้ำหนัก router ของชั้นถัดไป

การประมวลผล Unmerged Adapter

ชั้น MoE ที่ให้บริการจะคงสภาพไบต์ของเบสที่ถูกทำควอนไทเซชันไว้ และประมวลผล LoRA delta ควบคู่กันไปแบบขนาน:

$y = W_{int4} (x) + \frac{α}{r} B A x$

กระบวนการฝึกอบรม 3 ขั้นตอน

Head Distillation: Prerouter heads ถูกเทรนให้เลียนแบบการตัดสินใจของ router จริงในชั้นถัดไปของ teacher โดยที่เครือข่ายส่วนที่เหลือจะถูกตรึงค่าไว้
Student-Path SFT: Recovery LoRA adapters (เชื่อมเข้ากับ attention, linear-attention และ shared-expert projections แต่ไม่เชื่อมกับ streaming experts) ถูกเทรนโดยใช้ข้อความที่ teacher สร้างขึ้น พร้อมทั้งบังคับให้การเลือกเส้นทางต้องผ่าน prerouter
On-Policy Distillation: ตัว student สร้างข้อความขึ้นมาภายใต้การทำงานแบบ prerouted จากนั้น teacher ฐาน fp16 จะให้คะแนนข้อความที่สร้างขึ้น และพารามิเตอร์ของ LoRA จะถูกปรับปรุงผ่าน reverse Kullback–Leibler divergence ด้วยฟังก์ชันวัตถุประสงค์ความน่าจะเป็น top- $k$ ที่คำนึงถึงส่วนหาง (tail-aware)

ผลการทดลอง

ทรูพุตและการใช้หน่วยความจำ (Mac mini M4 Pro, 24 GB):
- edge0-35b ( $K = 4$ ): ทำความเร็วดีโค้ดได้ 20.4 tok/s โดยใช้หน่วยความจำสูงสุดขณะทำงานเพียง 2.9 GiB จากเช็กพอยต์ int4 ขนาด 19.5 GB (เปรียบเทียบกับ 3.9 tok/s และ 18.2 GiB เมื่อโหลดไว้ในแรมทั้งหมด)
- edge0-8b ( $K = 8$ ): ทำความเร็วดีโค้ดได้ 28.0 tok/s ภายใต้หน่วยความจำสูงสุด 1.5 GiB
ความเร็วที่เพิ่มขึ้นจากการดึงข้อมูลล่วงหน้า (MacBook M2, 16 GB):
- เมื่อเทียบกับการสตรีมแบบ on-demand ปกติ การทำ prefetch ล่วงหน้า 1 โทเค็นด้วย prerouter ช่วยเพิ่มความเร็วการดีโค้ดขึ้น +80% ถึง +84% ในทุกความกว้างของการเราต์ ( $K = 2$ : 4.8 เป็น 8.6 tok/s; $K = 4$ : 3.5 เป็น 6.4 tok/s; $K = 8$ : 1.8 เป็น 3.3 tok/s)
- เวลาที่เธรดหลักต้องบล็อกรอการโหลดลดลงอย่างเห็นได้ชัด (เช่น จาก 244.0 ms เหลือ 101.9 ms ต่อสเต็ปที่ $K = 4$ )
คุณภาพของโมเดล (OpenCompass):
- ในเบนช์มาร์กทั้ง 5 รายการ (AIME 2026, HumanEval, GPQA-Diamond, MMLU-Pro, IFBench) edge0-35b ได้คะแนนเฉลี่ย 79.2 เทียบกับ 83.2 ของเบส Qwen3.6 fp16 (ต่างกันเพียง 3.9 จุด)
- edge0-8b ทำได้ 69.9 เทียบกับ 72.7 ของเบส Ling 3.0 tiny fp16 (ต่างกันเพียง 2.8 จุด)

ข้อจำกัดและประเด็นที่ยังเปิดกว้าง

การรองรับคำขอพร้อมกัน (Serving Concurrency): Edge0 ทำงานแบบลำดับ FIFO สำหรับคำขอเดี่ยว การทำ batching จะเปลี่ยนชุด expert ที่ต้องใช้งานและพลวัตของการทำแคชไปในรูปแบบที่ยังไม่ครอบคลุมในโปรไฟล์สตรีมเดี่ยวนี้
Overhead การสร้างกราฟของ CPU: ความเร็วในการดีโค้ดถูกจำกัดโดย overhead ระดับเฟรมเวิร์กฝั่ง CPU (~44 ms/สเต็ป ในการสร้างกราฟ MLX ข้าม 40 ชั้น) ซึ่งการปรับแต่งบน SSD ไม่สามารถบรรเทาได้
คุณภาพที่ลดลงในด้านการใช้เหตุผล: การลดทอนประสิทธิภาพจะกระจุกตัวอยู่ในเบนช์มาร์กการใช้เหตุผลที่ซับซ้อน (เช่น คะแนนบน AIME ลดลง 6.1 จุดสำหรับ 35B และ 10.0 จุดสำหรับ 8B)
การพึ่งพาฮาร์ดแวร์: ทรูพุตสัมพัทธ์ที่ได้เพิ่มขึ้นจาก prerouter นั้นขึ้นอยู่กับความหน่วงของดิสก์ และผลลัพธ์นี้จะลดลงเมื่อแคชของน้ำหนักโมเดลมีสถานะอุ่น (warm) หรือเมื่อรันบนอุปกรณ์จัดเก็บข้อมูลความเร็วสูงพิเศษ

คำศัพท์สำคัญ

Edge0: เอนจินการทำ inference โมเดล MoE แบบสตรีมมิ่งที่รันโมเดลจาก SSD โดยซ้อนทับการอ่านข้อมูลไปกับการคำนวณผ่านการทำนายเส้นทาง และชดเชยคุณภาพที่สูญเสียไปด้วย unmerged recovery LoRA
prerouter: ส่วนหัวโครงข่ายประสาทขนาดเล็กประจำแต่ละชั้นที่ทำนายการตัดสินใจเลือกเส้นทางของชั้น N+1 ล่วงหน้า 1 โทเค็น โดยใช้ hidden state ของชั้น N และประวัติการเลือกเส้นทางล่าสุด
prediction-as-routing: หลักการออกแบบที่นำชุด expert ที่ทำนายได้จาก prerouter ไปใช้เป็นการเลือกเส้นทางจริงของชั้นนั้นทันทีในขั้นตอนการดีโค้ด ซึ่งช่วยขจัดปัญหาการทิ้งโทเค็นหรือบทลงโทษจากการทำนายพลาดระหว่างรันไทม์
recovery LoRA: low-rank adapter ที่ฝึกฝนผ่านการกลั่นกรองบนเส้นทาง prerouted ของ student และทำงานเป็นกิ่งขนานแบบไม่ผสานรวม เพื่อชดเชยข้อผิดพลาดจาก int4 quantization และการประมาณเส้นทาง
staged decode: โหมดการประมวลผล expert ที่ใช้ double buffering แบบช่องคงที่ โดยดัชนีของ expert จะถูกแมปผ่านตารางสล็อตบน GPU เพื่อหลีกเลี่ยงการจัดสรรหน่วยความจำใหม่ในทุกๆ ขั้นตอน
incremental sticky-slot stack (incr_stack): กลไกที่อัปเดต stacked expert tensors แบบถาวร ณ ตำแหน่งเดิม โดยเขียนทับเฉพาะช่อง expert ที่เปลี่ยนแปลงในแต่ละสเต็ป แทนที่จะประกอบเทนเซอร์ของ expert ใหม่ทั้งหมด
unmerged delta: การประมวลผล LoRA adapter ควบคู่ไปกับน้ำหนักเบสที่ถูกทำควอนไทเซชันแทนที่จะรวมเข้าเป็นเนื้อเดียวกัน ซึ่งช่วยป้องกันการสูญเสียรายละเอียดการปรับแต่งระหว่างการทำควอนไทเซชันซ้ำ
memory wall: ความเหลื่อมล้ำระหว่างอัตราการประมวลผลของหน่วยประมวลผลและแบนด์วิดท์/ความจุของ DRAM ซึ่งในการรัน LLM ภายในเครื่องจะถูกแบ่งระหว่างสถานะ KV-cache แบบไดนามิกและพื้นที่จัดเก็บน้ำหนักโมเดลแบบสแตติก

คำถามที่พบบ่อย

เหตุใด Edge0 จึงเลือกใช้การทำนายเป็นเส้นทางจริง แทนที่จะมองการทำนายเป็นเพียงคำใบ้?

การมองผลทำนายเป็นเพียงคำใบ้ทำให้ต้องมีระบบสำรองเมื่อรันไทม์หรือต้องทิ้งโทเค็นเมื่อทำนายพลาด การนำผลทำนายมาใช้เป็นเส้นทางจริงโดยตรงช่วยรับประกันว่าชุด expert ที่เตรียมไว้จะตรงกับชุดที่ถูกเลือกเสมอ ซึ่งเป็นการผลักภาระต้นทุนของการประมาณค่าไปไว้ในการเทรนแบบออฟไลน์ที่การทำ distillation สามารถกู้คืนคุณภาพกลับมาได้

ทำไม recovery LoRA จึงถูกรันแบบไม่ผสานรวม (unmerged) แทนที่จะรวมเข้ากับน้ำหนัก 4-bit?

เดลตาของน้ำหนัก LoRA มีค่า RMS อยู่ที่ประมาณ 10^-3 ซึ่งเล็กกว่าขนาดขั้นของ int4 quantization การรวมแล้วทำควอนไทเซชันใหม่จะทำลายผลของอะแดปเตอร์ส่วนใหญ่ไป โดยคงความแม่นยำในระดับ logits ได้เพียง 18% ขณะที่การรันเดลตาคู่ขนานแบบไม่ผสานรวมจะรักษาความแม่นยำได้เต็ม 100%

เหตุใด Edge0 จึงทำนายล่วงหน้าหนึ่งโทเค็น แทนที่จะทำนายภายในโทเค็นเดียวกัน?

การทำนายภายในโทเค็นเดียวกันต้องสั่ง flush GPU pipeline ในทุกๆ ชั้น ซึ่งกินเวลา 30 ถึง 100 ms ต่อสเต็ป ซึ่งนานกว่าเวลาโหลดข้อมูลที่พยายามจะซ่อน การทำนายนำหน้าหนึ่งโทเค็นเต็มเปิดโอกาสให้ทำนายทุกชั้นที่เตรียมไว้ได้ใน batch flush เดียว ทำให้มีเวลาคำนวณตลอดทั้งขั้นตอนดีโค้ดเพื่อซ่อน SSD I/O

Edge0 สามารถดีโค้ดได้เร็วกว่าระบบที่โหลดลงหน่วยความจำทั้งหมดบนเครื่อง 24 GB ได้อย่างไร?

โมเดลขนาด 19.5 GB เมื่อโหลดลงแรมทั้งหมดจะกินพื้นที่ถึง 18.2 GiB ทำให้แทบไม่เหลือที่ว่างสำหรับระบบปฏิบัติการและ KV cache ส่งผลให้เกิด OS memory paging ขั้นรุนแรงและทำให้การดีโค้ดตกไปอยู่ที่ 3.9 tok/s ส่วน Edge0 ใช้วิธีสตรีม expert จาก SSD โดยใช้หน่วยความจำสูงสุดเพียง ~2.9 GiB ทำให้ไม่เกิดแรงกดดันต่อหน่วยความจำของระบบ และสามารถทำงานได้ที่ 20.4 tok/s

คำอธิบายในหน้านี้เขียนโดย AI จากงานวิจัยและอาจมีข้อผิดพลาด โปรดอ่านต้นฉบับเพื่อดูรายละเอียด

อ่านพร้อมคำแปลและ AI