SkewAdam ช่วยลดการใช้หน่วยความจำในการฝึกแบบ Mixture-of-Experts ลงมากกว่า 60% และเพิ่มความแม่นยำขึ้นอย่างเห็นได้ชัด ช่วยให้นักพัฒนาสามารถรันโมเดล MoE ขนาด 6.78 พันล้านพารามิเตอร์บน GPU ขนาด 40 GB เพียงตัวเดียวได้

ทำไมการฝึก MoE ถึงติดปัญหาเรื่องขีดจำกัดของหน่วยความจำ (memory wall)

สถาปัตยกรรม Mixture-of-Experts จะส่งข้อมูลนำเข้าแต่ละชุดผ่านเครือข่ายย่อยที่เป็น "expert" เพียงบางส่วน ในขณะที่ยังคงรักษา dense backbone เอาไว้ ข้อดีคือโมเดลสามารถขยายขนาดไปถึงระดับพันล้านพารามิเตอร์ได้โดยที่การใช้พลังในการคำนวณ (compute) ไม่เพิ่มขึ้นตามสัดส่วน ในทางปฏิบัติ ตัว optimizer โดยเฉพาะ AdamW ที่ทีมส่วนใหญ่ใช้งาน คือตัวการหลักที่กินพื้นที่ GPU RAM สำหรับโมเดลขนาด 6.78B พารามิเตอร์ ลำพังแค่ momentum และ variance tensors ของ optimizer ก็ต้องการพื้นที่ถึงประมาณ 50 GB แล้ว เมื่อรวมกับ activations และ model weights หน่วยความจำสูงสุดจะพุ่งไปถึง 81.4 GB ซึ่งบีบให้ต้องใช้การตั้งค่าแบบหลาย GPU หรือต้องใช้ตารางการฝึกที่ช้าลง

สิ่งที่ SkewAdam ทำแตกต่างออกไป

SkewAdam ไม่ได้คิดค้นกฎการเรียนรู้ใหม่ แต่มันใช้วิธีจัดสรรตำแหน่งที่เก็บค่า moments ของ Adam ใหม่:

  • dense backbone จะเก็บค่า momentum แบบความละเอียดเต็ม (full-precision) เพื่อรักษาความราบรื่นในการอัปเดตที่เลเยอร์แบบ dense ต้องการ
  • expert bank จะเก็บค่าประมาณการของ variance แบบ factored ซึ่งช่วยลดปริมาณข้อมูลที่ต้องเก็บสำหรับแต่ละ expert
  • router ซึ่งทำหน้าที่ตัดสินใจว่าจะเปิดใช้งาน expert ตัวไหน จะยังคงเก็บค่าประมาณการ second-moment ที่แม่นยำเอาไว้

ด้วยการจัดการองค์ประกอบทั้งสามนี้แยกเป็น "ระดับ" (tiers) ที่ต่างกัน ตัว optimizer จึงสามารถลดความละเอียดที่เกินความจำเป็นในจุดที่สำคัญน้อยกว่า และรักษาความละเอียดไว้ในจุดที่สำคัญที่สุด

ผลลัพธ์ที่วัดผลได้

การรันโมเดล MoE ขนาด 6.78B พารามิเตอร์ตัวเดิมด้วย SkewAdam ให้ผลลัพธ์ดังนี้:

  • หน่วยความจำ GPU สูงสุด: 31.3 GB (ลดลงจาก 81.4 GB)
  • พื้นที่ที่ใช้สำหรับ optimizer-state: 1.29 GB (ลดลงจาก 50 GB)

สถานะ (state) ที่ลดลงนี้สามารถบรรจุลงในตัวเร่งความเร็ว (accelerator) ขนาด 40 GB เพียงตัวเดียวได้อย่างสบายๆ

ความแม่นยำที่เพิ่มขึ้น ไม่ใช่แค่การประหยัดหน่วยความจำ

การลดหน่วยความจำมักจะส่งผลเสียต่อคุณภาพของโมเดล แต่ SkewAdam กลับช่วยปรับปรุงค่า perplexity ซึ่งเป็นตัวชี้วัดมาตรฐานของโมเดลภาษา จาก 126.8 (AdamW) เป็น 108.4 นอกจากนี้ยังทำผลงานได้ดีกว่า Muon (120.2) และ Lion (393.7) ในงานเดียวกัน ผู้เขียนระบุว่าการพัฒนาที่เพิ่มขึ้นนี้มาจากการรักษา momentum ไว้ในทั้งสามระดับ ในขณะที่วิธีการที่ตัด momentum ออกไปทั้งหมดอย่าง Adafactor นั้นให้ผลลัพธ์ที่ด้อยกว่า

คำถามที่ยังไม่มีคำตอบ

ผลลัพธ์ของ SkewAdam ได้รับการสาธิตบนโมเดลขนาด 6.78B พารามิเตอร์ ซึ่งยังไม่เป็นที่แน่ชัดว่าสัดส่วนของหน่วยความจำ (memory-state ratios) แบบเดียวกันนี้จะยังคงใช้ได้กับโมเดลที่มีขนาดใหญ่กว่าเดิมเป็นสิบเท่า หรือกับงานอื่นๆ ที่นอกเหนือจากการสร้างโมเดลภาษาหรือไม่

สิ่งที่ควรจับตามอง

  • การทดสอบประสิทธิภาพ (benchmarks) บนการตั้งค่า MoE ที่ใหญ่ขึ้น (ระดับหลายหมื่นล้านพารามิเตอร์)
  • การนำไปใช้ในเฟรมเวิร์กโอเพนซอร์สและการรวมเข้ากับสคริปต์การฝึกที่เป็นที่นิยม
  • ความคิดเห็นจากชุมชนเกี่ยวกับข้อแลกเปลี่ยนที่ซ่อนอยู่ (hidden trade-offs) เช่น ความเร็วในการลู่เข้า (convergence speed) หรือความเสถียรภายใต้ตารางการเรียนรู้ (learning-rate schedules) ที่แตกต่างกัน

ที่มา: โพสต์จากนักพัฒนาที่ให้รายละเอียดเกี่ยวกับการออกแบบ optimizer และผลลัพธ์เชิงประจักษ์