การทำ Quantization ของ KV Cache สำหรับ LLM บนอุปกรณ์ (On-Device)
การรัน Llama 3.2 3B บนอุปกรณ์ Android ที่มี RAM เพียง 2 GB นั้นเป็นเรื่องยาก นักพัฒนาส่วนใหญ่มักมุ่งเน้นไปที่น้ำหนักของโมเดล (model weights) แต่นี่คือความผิดพลาด เพราะตัวการที่กินหน่วยความจำที่แท้จริงคือ KV cache
KV cache จะขยายขนาดขึ้นเรื่อยๆ ในขณะที่คุณแชท หากคุณใช้ความละเอียดแบบ FP16 มาตรฐาน ตัว cache จะกินพื้นที่หลายร้อยเมกะไบต์ ซึ่งส่งผลให้แอปของคุณค้างหรือเด้งออก (crash) หลังจากคุยไปได้เพียงไม่กี่ประโยค
คุณสามารถแก้ไขปัญหานี้ได้ด้วย 3 ขั้นตอนเฉพาะเจาะจง ดังนี้:
- ใช้ Mixed-Precision Quantization Keys และ values ไม่จำเป็นต้องมีความละเอียดเท่ากัน Key caches สามารถรองรับความละเอียดต่ำได้ดี แต่ Value caches ทำไม่ได้
- ใช้ INT4 สำหรับ keys
- ใช้ INT8 สำหรับ values
วิธีนี้จะช่วยลดขนาด cache ลงได้ถึง 62% สำหรับ context ขนาด 2048 token คุณจะลดการใช้หน่วยความจำจาก 224 MB เหลือเพียง 84 MB โดยที่ไม่ต้องไปเปลี่ยนน้ำหนักของโมเดลเลย
ใช้การทำ Sliding Window Eviction คุณไม่สามารถเก็บทุก token ไว้ในหน่วยความจำหลัก (active memory) ได้ตลอดเวลา ให้ใช้ sliding window เพื่อเก็บเฉพาะ 1536 token ล่าสุดเท่านั้น และให้เก็บ 64 token แรกไว้เป็น anchor เพื่อรักษา system prompt เอาไว้
ใช้ Flash Spilling เมื่อ token หลุดออกจาก sliding window ให้ย้ายพวกมันไปยัง flash storage โดยใช้ memory-mapped files บน Android ซึ่งหน่วยความจำ UFS 4.0 ในปัจจุบันมีความเร็วเพียงพอที่จะดึงข้อมูลเหล่านี้กลับเข้าสู่หน่วยความจำ (page back) ได้โดยไม่เกิดอาการหน่วง
ผลลัพธ์ที่ได้นั้นน่าทึ่งมาก เมื่อทดสอบบน Snapdragon 8 Gen 3:
- การใช้หน่วยความจำสูงสุด (Peak memory) ลดลงจนต่ำกว่าขีดจำกัด 2 GB
- จำนวนรอบการสนทนาสูงสุดเพิ่มขึ้นจาก 4 รอบ เป็นมากกว่า 12 รอบ
- ความเร็วในการประมวลผล token เพิ่มขึ้น เนื่องจาก cache ที่มีขนาดเล็กลงช่วยให้ใช้ bandwidth ของหน่วยความจำได้ดีขึ้น
- คุณภาพของโมเดลยังคงใกล้เคียงเดิมเกือบทั้งหมด
หลีกเลี่ยงข้อผิดพลาดเหล่านี้:
- อย่าทำ quantization ให้กับ keys และ values ในระดับเดียวกัน เพราะจะทำให้คุณภาพลดลง
- อย่าละเลยเรื่อง thermal throttling เนื่องจากการประมวลผล (inference) อย่างต่อเนื่องจะทำให้เครื่องร้อน ควรตรวจสอบ Android Thermal HAL เพื่อจัดการประสิทธิภาพการทำงาน
- อย่าลืมเรื่องวงจรชีวิต (lifecycle) ของ cache ควรผูก mapped buffers เข้ากับ scope ที่เหมาะสมเสมอเพื่อป้องกันปัญหา memory leaks
วางแผนงบประมาณหน่วยความจำ (memory budget) ของคุณให้เรียบร้อย ก่อนที่จะเริ่มสร้างฟีเจอร์ต่างๆ
Source: https://dev.to/software_mvp-factory/kv-cache-quantization-for-on-device-llms-kf