ทำให้ Knowledge Distillation มีประสิทธิภาพ คุ้มค่า และใช้งานได้จริงในระดับ Scale ใหญ่

ทำให้ Knowledge Distillation มีประสิทธิภาพ คุ้มค่า และใช้งานได้จริงในระดับ Scale ใหญ่ในยุคที่โมเดลภาษาขนาดใหญ่ (Large Language Models - LLMs) ได้รับความนิยมอย่างแพร่หลาย การนำโมเดลเหล่านี้ไปใช้งานจริงมักมาพร้อมกับค่าใช้จ่ายที่สูงลิ่ว เนื่องจากขนาดที่ใหญ่และความต้องการทรัพยากรคอมพิวเตอร์มหาศาล เช่น โมเดล Kimi-K3 ที่มีพารามิเตอร์สูงถึง 2.8 ล้านล้านตัว ต้องการ VRAM ถึง 3TB เพียงเพื่อโหลดโมเดล เทคนิค Knowledge Distillation จึงกลายเป็นทางออกสำคัญในการย่อขนาดโมเดลให้เล็กลง โดยยังคงประสิทธิภาพใกล้เคียงของเดิมไว้อย่างไรก็ตาม กระบวนการ Knowledge Distillation นี้เองที่เป็นส่วนสำคัญในการกำหนดคุณภาพสุดท้ายของโมเดล แต่ก็มักจะเป็นขั้นตอนที่ใช้ทรัพยากรมากที่สุดและมีค่าใช้จ่ายสูงที่สุดเช่นกัน การต้องโหลดทั้งโมเดลครู (Teacher Model) และโมเดลนักเรียน (Student Model) ไปพร้อมกัน และการคำนวณการกระจายความน่าจะเป็น (Probability Distribution) สำหรับทุก Token ในชุดคำศัพท์ (Vocabulary) ต้องการ VRAM จำนวนมหาศาล ซึ่งโดยทั่วไปแล้วต้องใช้ GPU หลายร้อยตัวและกลยุทธ์ Tensor Parallelism ที่ซับซ้อนบทความวิจัยล่าสุดจาก Multiverse Computing ได้นำเสนอแนวทางใหม่ "Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss" เพื่อแก้ปัญหานี้ ด้วยการปรับปรุงระบบ 2 ส่วนหลัก คือ การแคช (Cache) Top-K Logits ของโมเดลครู เพื่อให้ไม่ต้องโหลดโมเดลครูค้างไว้พร้อมกับโมเดลนักเรียนตลอดเวลา และการพัฒนากลไก Loss ที่มีประสิทธิภาพด้านหน่วยความจำ (Memory-efficient KL-divergence loss) ซึ่งช่วยลดการใช้ VRAM ลงอย่างมากเมื่อเทียบกับการใช้งานทั่วไปในไลบรารีอย่าง PyTorch หรือ NVIDIA Megatron-Bridgeการปรับปรุงทั้งสองส่วนนี้ ทำให้ต้นทุนการเทรนลดลงอย่างมีนัยสำคัญ จนสามารถทำการ "รักษาบริบทให้ยาว (Long-context healing)" บน GPU เพียงตัวเดียว และทำให้การทดลองในระดับ Scale ใหญ่เป็นไปได้อย่างคุ้มค่าทำไมการทำ Knowledge Distillation ถึงมีค่าใช้จ่ายสูง?การทำ Knowledge Distillation แบบมาตรฐานที่ใช้ Kullback-Leibler (KL) divergence loss แบบออนไลน์ (Online Distillation) จะต้องโหลดทั้งโมเดลครูและโมเดลนักเรียนไว้พร้อมกัน ในทุกขั้นตอนการเทรน โมเดลครูจะต้องทำการ Forward Pass เต็มรูปแบบเพื่อสร้างการกระจายความน่าจะเป็นออกมา จากนั้นจึงนำไปฝึกโมเดลนักเรียนให้เลียนแบบ แม้ว่าการตั้งค่าแบบนี้จะให้ผลลัพธ์ที่ละเอียดอ่อนที่สุด เนื่องจากสามารถเข้าถึงการกระจายเต็มรูปแบบของโมเดลครูได้ แต่ก็เป็นวิธีที่ใช้หน่วยความจำและพลังประมวลผลมากที่สุดเช่นกันลองพิจารณาตัวอย่าง: โมเดล gpt-oss-120b มีชุดคำศัพท์ 201,088 Token หากใช้ความยาวบริบท (Sequence Length) 32K และ Batch Size 4 เฉพาะ Tensor ความน่าจะเป็นของโมเดลครูจะมีขนาดประมาณ 4 × 201,088 × 32,768 ซึ่งหากเป็นรูปแบบ bfloat16 ก็จะใช้ VRAM สูงถึง 50GB สำหรับ Tensor เดียว เมื่อรวมกับ Gradients, Activations, Model Weights และ Optimizer States การเทรนเพียงหนึ่งรอบอาจใช้ VRAM สูงสุดถึง 250GB ซึ่งเกินกว่าความสามารถของ GPU อย่าง H200 หรือ B200 เพียงตัวเดียวการทำ Knowledge Distillation แบบ Offline และ Fused Chunked KL Lossการทำ Distillation แบบ Offlineแทนที่จะคำนวณผลลัพธ์ของโมเดลครูใหม่ทุกครั้ง เราสามารถคำนวณผลลัพธ์เพียงครั้งเดียว จากนั้นแคช (Cache) Top-100 Token ที่มีความน่าจะเป็นสูงสุดในแต่ละตำแหน่ง (Position) แล้วนำไปฝึกโมเดลนักเรียนจาก Cache นี้ วิธีนี้ทำให้ไม่ต้องโหลดโมเดลครูค้างไว้ในหน่วยความจำระหว่างการเทรน และไม่ต้องคำนวณซ้ำเมื่อ Cache ถูกสร้างขึ้นแล้ว ทำให้สามารถนำ Cache เดียวกันไปใช้กับการทดลองอื่นๆ ได้อีกด้วยกลไก Fused Chunked KL Lossเพื่อทำความเข้าใจว่า Loss Function มีค่าใช้จ่ายสูงอย่างไร ลองจินตนาการว่า Loss Function ทำงานอย่างไร: สำหรับทุก Token ในบริบท และทุกคำในชุดคำศัพท์ Loss Function ต้องการค่าที่บอกว่าการทำนายของโมเดลนักเรียนแตกต่างจากโมเดลครูมากน้อยเพียงใด หากจัดเรียงเป็นตาราง โดยมีแถวเป็นรายการคำในชุดคำศัพท์ และคอลัมน์เป็นตำแหน่งในบริบท สำหรับชุดคำศัพท์ขนาดใหญ่ (100K+ คำ) และบริบทที่ยาวมากๆ ตารางนี้จะมีขนาดมหึมา และการคำนวณ KL Loss แบบดั้งเดิมจะสร้างตารางทั้งหมดขึ้นมาก่อนที่จะสามารถให้ค่าออกมาได้มีการเปรียบเทียบ 3 วิธีในการคำนวณ Loss เดียวกันนี้ ซึ่งทั้งหมดเทียบเท่ากันทางคณิตศาสตร์:Dense KL: เป็นวิธีการแบบดั้งเดิม สร้างตารางการกระจายความน่าจะเป็นเต็มรูปแบบของโมเดลครูจาก Top-100 Logits ที่แคชไว้ แล้วนำไปเปรียบเทียบกับ Logits ของโมเดลนักเรียน วิธีนี้ใกล้เคียงกับการทำ Online Distillation มากที่สุด จึงใช้เป็น Baseline ในการตรวจสอบความถูกต้อง แต่มีข้อเสียคือต้องเก็บตารางเต็มรูปแบบของชุดคำศัพท์ × บริบท ไว้ในหน่วยความจำถึงสองชุดForward-chunked KL: ใช้วิธีการเก็บข้อมูลของโมเดลครูแบบ Sparse (เฉพาะ Top-100 Logits ที่แคชไว้ โดยไม่ขยายเป็นตารางเต็มรูปแบบ) และคำนวณ Loss ทีละส่วน โดยประมวลผลทีละส่วนของบริบท (Slice of sequence positions) วิธีนี้ช่วยลดการใช้หน่วยความจำของโมเดลครูและส่วนเปรียบเทียบ และจากการทดสอบพบว่าเป็นวิธีที่เร็วที่สุด อย่างไรก็ตาม โมเดลนักเรียนเองก็ยังต้องคำนวณ Logits เต็มรูปแบบและเก็บไว้สำหรับการทำ Backward Pass ทำให้หน่วยความจำยังคงเพิ่มขึ้นอย่างรวดเร็วตามความยาวบริบทFused chunked KL (การพัฒนาหลัก): ก้าวไปอีกขั้นโดยการรวมการคำนวณ Output Projection ของโมเดลนักเรียนเข้ากับ Loss Function โดยตรง วิธีนี้ไม่เคยสร้าง Logits เต็มรูปแบบของโมเดลนักเรียนเลย แต่จะประมวลผลทีละส่วนของบริบทแบบ End-to-End โดยแปลง Hidden States ไปเป็น Logits สำหรับส่วนนั้นๆ แล้วนำผลลัพธ์ไปรวมกับ Loss ที่กำลังคำนวณ จากนั้นจึงทิ้งส่วนนั้นไปก่อนจะไปยังส่วนถัดไป ส่วน Backward Pass จะคำนวณแต่ละส่วนใหม่แบบ On-the-fly แทนที่จะเก็บไว้ ข้อดีคือ VRAM ที่ใช้จะเพิ่มขึ้นเป็นสัดส่วนเชิงเส้น (Linearly) กับความยาวบริบท แทนที่จะพุ่งสูงขึ้นตามขนาดของชุดคำศัพท์ × บริบทผลลัพธ์ที่ได้ในทางปฏิบัติการเปรียบเทียบทั้ง 4 วิธี (Online Distillation, Dense KL, Forward-chunked KL, และ Fused chunked KL) บน GPU H200 ตัวเดียว โดยใช้ Llama 3.1 8B Instruct เป็นโมเดลครู และโมเดล Llama ขนาด 3.2B เป็นโมเดลนักเรียน ที่ความยาวบริบท 8K Token พบว่าทั้ง 4 วิธีให้ Training Loss ใกล้เคียงกันมาก แม้ว่าการทำ Offline Distillation จะใช้เพียง Top-100 Logits ที่แคชไว้ก็ตามกราฟ Loss Curve เกือบจะทับซ้อนกันในทุกวิธี แสดงให้เห็นว่า Offline Distillation ด้วย Top-100 Logits ที่แคชไว้ ให้ผลลัพธ์ที่ไม่สูญเสียคุณภาพเมื่อเทียบกับ Online Distillationการขยายผลสู่บริบทที่ยาวขึ้นเมื่อทำการทดสอบเฉพาะส่วนของ Output-projection network (ไม่มีส่วน Transformer body) ที่ความยาวบริบท 32K Token พบว่า VRAM ที่ใช้ลดลงจาก 85.2 GiB (Dense KL) เหลือเพียง 5.45 GiB (Fused chunked KL) ซึ่งลดลงถึง 15.6 เท่า! และที่ความยาวบริบท 64K Token วิธี Dense KL ไม่สามารถทำงานได้อีกต่อไปที่ความยาวบริบท 256K Token, Fused chunked KL ใช้ VRAM เพียง 11.6 GiB เทียบกับ 134.2 GiB ของวิธีถัดไปที่ดีที่สุด และมีความเร็วต่อรอบการเทรนประมาณ 3.3 เท่าในการทำ Distilling โมเดล GPT-OSS 20B ที่ความยาวบริบท 32,768 Token การใช้ Fused Loss ช่วยลดจำนวน GPU ที่ต้องใช้จาก 4 โหนด เหลือเพียง 1 โหนด ความเร็วต่อรอบการเทรนลดลงจาก 57.0 เหลือ 12.23 วินาที (เร็วขึ้นประมาณ 5 เท่า) และ Throughput ต่อ GPU เพิ่มขึ้นจาก 74.2 เป็น 345.7 TFLOP/sโมเดลนักเรียนที่ได้การตั้งค่า Offline แบบมีประสิทธิภาพนี้เอง ที่ทำให้การทำ Knowledge Distillation ขนาดใหญ่เป็นไปได้ในราคาที่คุ้มค่า โมเดลนักเรียนที่ได้ ซึ่งถูกย่อส่วนจาก Llama 3.1 8B Instruct ลงมาเหลือประมาณ 3.2B พารามิเตอร์ ยังคงรักษาความแม่นยำส่วนใหญ่ของโมเดลครูไว้ได้ในชุดทดสอบ BoolQ และ HellaSwag โดยมีคะแนน MMLU ลดลงไปประมาณ 9 จุด แต่มีขนาดพารามิเตอร์น้อยกว่าครึ่งงานวิจัยนี้เป็นส่วนหนึ่งของความพยายามอย่างต่อเนื่องของ Multiverse Computing ในการทำให้ Knowledge Distillation และ Long-context healing มีความคุ้มค่าและสามารถนำไปใช้งานได้จริงในระดับ Scale ใหญ่ ไม่ใช่แค่เป็นสูตรสำเร็จที่ทำครั้งเดียว แต่เป็นสิ่งที่ทีมต่างๆ สามารถนำไปปรับปรุงและทดลองได้อย่างต่อเนื่องในต้นทุนที่ต่ำหากต้องการรายละเอียดทางเทคนิคทั้งหมด รวมถึง Gradient แบบ Closed-form ของ Fused Chunked KL Loss และการตั้งค่าการเทรนฉบับสมบูรณ์ สามารถอ่านบทความฉบับเต็ม หรือติดต่อทีมงานเพื่อพูดคุยเกี่ยวกับการนำไปปรับใช้กับกระบวนการ Distillation ของคุณได้เราได้เปิด Source Code ของการพัฒนา Chunked-loss ไว้แล้วที่: [github.com/CompactifAI/Full-Chunked-KL-Loss](https://github.com/CompactifAI/Full-Chunked-KL-Loss)#KnowledgeDistillation #LLMs #MachineLearning #AI #EfficientAIhttps://huggingface.co/blog/MultiverseComputingCAI/efficient-knowledge-distillation