การพอร์ตโมเดล Google’s Gemma-4 31B ไปยัง AWS Inferentia2 inf2.24xlarge ให้ผลลัพธ์ที่ตรงกันแบบ token-for-token กับตัวอ้างอิงบน CPU อย่างสมบูรณ์แบบ—แต่ประโยคที่สร้างออกมากลับเป็นภาษาที่อ่านไม่รู้เรื่องทั้งหมด ช่องว่างระหว่างคำว่า “matching” (ตรงกัน) กับ “working” (ใช้งานได้) กลายเป็นคำเตือนสำหรับใครก็ตามที่พยายามจะยัดโมเดล LLM ขนาดมหึมาลงในชิปประมวลผล (inference chips) แบบปรับแต่งพิเศษของ Amazon

ทำไมการตรงกันแบบ token-by-token ถึงยังไม่เพียงพอ

นักพัฒนาได้เปรียบเทียบ output token แต่ละตัวจากอุปกรณ์ Inferentia กับ token ที่ได้จากการรันโมเดลบน CPU ผลลัพธ์ที่ออกมานั้นเหมือนกันทุกประการ ทำให้ดูเหมือนว่าฮาร์ดแวร์สามารถจำลองการทำงานตามตัวอ้างอิง (reference implementation) ได้อย่างแม่นยำ แต่ในความเป็นจริง ทั้งสองกระแสข้อมูลได้ป้อน prompt ที่ผิดรูปแบบเข้าไปในโมเดลที่ถูกตัด chat template ออกไป และมีการใส่ turn markers ที่ไม่ถูกต้อง การขาด template นี้ทำให้โมเดลติดอยู่ในลูปไม่สิ้นสุดและพ่นข้อความที่ไร้สาระออกมา ฮาร์ดแวร์ทำหน้าที่ของมันได้ดีแล้ว นั่นคือการจำลอง bug ที่มีอยู่ในโค้ดอ้างอิงออกมานั่นเอง

บทเรียนนี้เรียบง่ายมาก: SEQ_MATCH (ความเท่ากันของ token ตามลำดับ) ไม่ได้เท่ากับความถูกต้อง หากตัวอ้างอิง (reference implementation) เสีย การจำลองฮาร์ดแวร์ที่ซื่อสัตย์ก็จะรับเอาความล้มเหลวนั้นมาด้วย การตรวจสอบความถูกต้อง (Validation) ต้องไปไกลกว่าแค่ความเท่ากันในระดับ token แต่จำเป็นต้องมีการตรวจสอบการทำงานแบบ end-to-end ด้วย input ที่จัดรูปแบบอย่างถูกต้อง

Buffers ที่ปลอมตัวเป็น parameters

ในช่วงขั้นตอนการโหลด ตัวโหลดโมเดลได้ข้ามส่วนประกอบที่เรียกว่า layer_scalar ไป เนื่องจากในนิยามโมเดลของ PyTorch โค้ดได้ลงทะเบียนออบเจกต์นี้เป็น buffer แทนที่จะเป็น parameter โดย buffer คือ static tensors ที่การเทรนจะไม่ทำการอัปเดต และตัวโหลดโมเดลหลายตัวมักจะมองข้ามสิ่งนี้เมื่อทำการแปลงเป็นรูปแบบที่รองรับ Neuron การข้ามส่วนนี้ไปทำให้ค่า scaling factors ของหลายๆ layer ยังคงเป็นค่าเริ่มต้น (default) ซึ่งส่งผลให้การคำนวณทางคณิตศาสตร์ทั่วทั้งเครือข่ายผิดเพี้ยนไป ไม่มีข้อผิดพลาดใดๆ ถูกแจ้งเตือน โมเดลสามารถ compile ได้ และ inference pipeline ก็ทำงานได้ตามปกติ แต่ผลลัพธ์ทางตัวเลขนั้นผิดพลาด

สำหรับใครก็ตามที่กำลังย้ายโมเดลขนาดใหญ่ไปยัง Inferentia ควรตรวจสอบ (audit) ทุกๆ non-parameter tensor แม้ว่า tensor นั้นจะไม่ได้มีไว้เพื่อการเรียนรู้ (learned) แต่มันอาจยังจำเป็นต่อการคำนวณ forward-pass ที่ถูกต้อง การตรวจสอบการรวม buffer ด้วยตนเองสามารถช่วยป้องกันข้อผิดพลาดด้าน scale ที่เกิดขึ้นเงียบๆ (silent scale errors) ซึ่งยากต่อการวินิจฉัยในภายหลัง

ความผันผวนของ Spot-instance และการ compile ที่ใช้เวลา 39 นาที

การรันโมเดลขนาด 31 พันล้านพารามิเตอร์บน spot instance ดูเหมือนจะมีราคาถูก แต่ความประหยัดนั้นมาพร้อมกับเหตุการณ์การดึงทรัพยากรคืน (reclaim events) ที่คาดเดาไม่ได้ เวลาในการ compile ของนักพัฒนา—ซึ่งใช้เวลาประมาณ 39 นาทีในการแปลงโมเดลเป็นโค้ดที่รองรับ Neuron—หายวับไปทันทีเมื่อ AWS ดึง instance คืน เพื่อรับมือกับการขัดจังหวะดังกล่าว พวกเขาจึงสร้างระบบความปลอดภัยแบบสามประสาน:

  • ModelBuilder ช่วยควบคุมการใช้หน่วยความจำให้อยู่ภายในขีดจำกัด 384 GB ของ host เพื่อหลีกเลี่ยงการ crash ที่จะบังคับให้ต้องเริ่มใหม่
  • การทำ S3 mirroring ทันที ทั้งไฟล์ weight ดิบและไฟล์ “neffs” (Neuron executable files) ที่ compile แล้ว ช่วยให้ instance ใหม่สามารถทำงานต่อจากจุดที่ instance เดิมค้างไว้ได้ทันที
  • multi-region poller ทำหน้าที่สแกน AWS regions เพื่อหา spot capacity ที่ว่างอยู่ และเริ่มทำงาน instance ใหม่ทันทีที่มีทรัพยากรปรากฏขึ้น

ขั้นตอนเหล่านี้เปลี่ยนการ compile ที่เปราะบางและมีจุดล้มเหลวเพียงจุดเดียว (single-point) ให้กลายเป็น pipeline ที่ยืดหยุ่นและสามารถทนทานต่อความผันผวนของตลาด spot ได้

หลุมพรางของการทำ Sharding ด้วย attention layouts แบบผสม

Gemma-4 31B ใช้การกำหนดค่า attention สองรูปแบบ บาง layer ใช้ key-value (KV) heads จำนวน 4 หัว ในขณะที่ layer อื่นๆ ใช้จำนวนที่ต่างออกไป การแบ่งโมเดลให้เท่ากันใน 8 parallel ranks จะล้มเหลวเมื่อจำนวน KV head ของ layer นั้นๆ ไม่สามารถหารได้ลงตัว การพยายามทำ shard layer ที่มี 4 heads ลงใน 8 ranks จะบังคับให้แต่ละ rank ต้องจัดการกับ "ครึ่งหัว" ซึ่งเป็นสิ่งที่เป็นไปไม่ได้ในทางคณิตศาสตร์ และจะทำให้เกิด shape mismatches และ runtime errors

วิธีแก้ไขคือการ replicate layers ที่ถูก sharded แบบ global (layer ที่มีจำนวน head ที่เข้ากันได้) ไปยังทุก ranks และทำ shard เฉพาะ “sliding” layers ที่มีจำนวน head ที่สามารถแบ่งได้อย่างลงตัวเท่านั้น กลยุทธ์แบบไฮบริดนี้ช่วยรักษาประสิทธิภาพของ tensor-parallel ในขณะที่หลีกเลี่ยงการหาร KV heads ที่ไม่ถูกต้อง ซึ่งช่วยกำจัดข้อผิดพลาดด้าน tensor-parallelization ที่เคยเป็นปัญหาในการพยายามครั้งก่อนๆ

บทสรุป

การพอร์ต LLM ขนาดมหึมาไปยัง Inferentia เป็นมากกว่าแค่การ compile แล้วรัน แต่มันต้องการการทดสอบการทำงาน (functional testing) ที่เข้มงวดมากกว่าแค่ความเท่ากันของ token, การตรวจสอบอย่างละเอียดถี่ถ้วนว่าทุกๆ tensor ไม่ว่าจะเป็น parameter หรือ buffer ถูกจัดการอย่างถูกต้อง และกลยุทธ์การปรับใช้ (deployment strategy) ที่คาดการณ์ถึงการดึงทรัพยากร spot-instance คืน และท้ายที่สุด การทำ sharding ต้องคำนึงถึงโครงสร้าง attention geometry ภายในของโมเดล มิฉะนั้น ความสามารถในการประมวลผลแบบขนาน (parallelism) ที่สัญญาว่าจะให้ความเร็ว จะกลายเป็นต้นเหตุของความล้มเหลวที่เกิดขึ้นอย่างเงียบๆ แทน