Google โชว์ MaxText เทรนโมเดลต่อได้ทั้งที่ TPU พังกลางคัน กู้คืนใน ~110 วินาที

การเทรนโมเดลขนาดใหญ่บนเครื่องหลายร้อยหลายพันตัวมีศัตรูตัวฉกาจที่นักพัฒนาทุกคนเจอ นั่นคือ ฮาร์ดแวร์พังกลางคัน ปกติแล้วถ้า worker ตัวเดียวล้ม งานเทรนแบบ multi-node ทั้งงานจะ crash ตามไปด้วย ต้องรีสตาร์ทใหม่ตั้งแต่ checkpoint ล่าสุด เสียทั้งเวลา warm-up และค่า compute ที่จ่ายไปเปล่า ๆ เมื่อวันที่ 6 กรกฎาคม 2026 ทีม Google Cloud เผยบล็อกที่ตอบโจทย์นี้ตรง ๆ ด้วยชื่อยั่วให้อ่านว่า “เราจงใจฆ่า TPU กลางการเทรน แล้วมันกู้คืนได้ในไม่กี่วินาที”

โพสต์นี้แนะนำ elastic training ใน MaxText — เฟรมเวิร์กเทรน LLM แบบ open-source บน JAX ของ Google — โดยสาธิตแบบจับต้องได้ว่าเมื่อ TPU slice หนึ่งถูกดึงออกกลางการเทรน ระบบไม่ crash แต่ “รับ exception เป็น Python error ธรรมดา” แล้วซ่อมตัวเองต่อได้ ประเด็นสำคัญสำหรับนักพัฒนาคือ นี่ไม่ใช่แค่ลูกเล่น แต่เป็นแพตเทิร์นที่เปลี่ยนวิธีคิดเรื่องความทนทานของงานเทรนขนาดใหญ่

กลไกกู้คืน: สาม component ที่ทำงานร่วมกัน

หัวใจของ elastic training อยู่ที่การเปลี่ยน “hardware failure = job crash” ให้กลายเป็น “exception ที่ดักจับได้” ประกอบด้วยสามส่วน:

• Pathways ตรวจจับการล้ม — เมื่อ worker ล้ม Pathways จะรู้ได้จาก operation ที่ล้มเหลวหรือ heartbeat timeout แล้วโยน JaxRuntimeError ออกมา (ประมาณ 10 วินาทีสำหรับการตรวจจับผ่าน heartbeat) แทนที่จะให้ทั้ง process ตาย
• ตัวห่อ elastic_retry — decorator ที่ดักจับ exception นี้และเริ่มกระบวนการกู้คืน โดยไม่ทำให้ controller process ล่ม
• Orbax checkpoint validation — ตรวจว่า checkpoint ที่จะโหลดกลับมานั้น “สมบูรณ์” จริง ด้วยการเช็กไฟล์ marker ชื่อ commit_success กันการโหลด checkpoint ที่เขียนค้างไว้ครึ่ง ๆ กลาง ๆ

จุดที่ทำให้ทั้งหมดนี้เป็นไปได้คือสถาปัตยกรรม single-controller ของ Pathways — controller เป็น Python process เดียวที่มองเห็นชิป TPU ทั้งหมด ทำให้กู้คืนได้โดยไม่ต้องรื้อ worker ที่ยังดีอยู่และ controller ทิ้ง เก็บทั้ง Python state และ compilation cache ไว้ข้ามการล้มได้

ตัวเลขจากการสาธิตจริง

ทีมงานตั้งการทดลองด้วย TPU v5e-16 จำนวน 3 slice (รวม 48 ชิป) เทรนโมเดล Qwen3-0.6b แล้วจงใจฆ่า slice หนึ่งกลางคัน ผลลัพธ์:

• ดาวน์ไทม์รวม ~110 วินาที นับจากตอนฆ่าจนถึง training step ถัดไป
• เสียความคืบหน้าไปเพียง 88 training steps
• แจกแจงเวลา: heartbeat detection ~13 วินาที, จัดสรร pod/รีสตาร์ท ~50 วินาที, กู้ checkpoint จาก Cloud Storage ~5.4 วินาที และ warm-up ~12.7 วินาที

MaxText ให้เลือก recovery path ได้สองแบบ: pause and resume (รอให้ slice ที่ล้มกลับมาครบก่อนค่อยเทรนต่อ) หรือ replica resize (เทรนต่อบน slice ที่เหลือทันทีด้วย throughput ที่ลดลง) แบบหลังเหมาะกับงานที่ห้ามหยุดยาว ยอมแลกความเร็วชั่วคราวเพื่อไม่ให้งานสะดุด สำหรับโมเดลใหญ่ Google แนะนำให้ตั้ง ENABLE_PATHWAYS_PERSISTENCE=1 เพื่อให้ checkpoint I/O วิ่งตรงไปที่ storage แทนที่จะผ่าน proxy

ทำไมนักพัฒนาต้องสนใจ

ยิ่งงานเทรนใหญ่และใช้เครื่องเยอะ โอกาสที่จะมีเครื่องสักตัวล้มระหว่างทางยิ่งสูงขึ้นแบบทวีคูณ — ในสเกลหลายพันชิป การล้มไม่ใช่ “ถ้า” แต่เป็น “เมื่อไหร่” ต้นทุนจริงของการล้มไม่ใช่แค่ compute ที่เสียไป แต่รวมถึงเวลารื้อ-ประกอบ cluster ใหม่ทั้งหมด elastic training ตัดขั้นตอนนั้นทิ้งด้วยการเก็บ worker ที่ยังดีและ controller ไว้เหมือนเดิม

มุมมอง AI Curator

สิ่งที่น่าสนใจที่สุดของโพสต์นี้ไม่ใช่ตัวเลข 110 วินาที แต่คือการที่ Google เลือก “ป่าวประกาศจุดแข็งเชิงสถาปัตยกรรม” ของ Pathways อย่างเปิดเผย — โมเดล single-controller คือสิ่งที่ทำให้ทั้งหมดนี้เป็นไปได้ และมันต่างจากแนวทาง multi-controller (เช่น SPMD ล้วน ๆ ที่หลายเฟรมเวิร์กใช้) ที่การล้มของ worker ตัวเดียวมักลามเป็นการล้มทั้งงาน นี่เป็นการวางหมากเชิงกลยุทธ์: JAX + Pathways + TPU ขายเป็นชุดที่ความทนทานเป็นคุณสมบัติในตัว ไม่ใช่สิ่งที่ต้องต่อเพิ่มเอง

สำหรับทีมไทยและ SEA ที่เทรนหรือ fine-tune โมเดลบนคลาวด์ บทเรียนที่นำไปใช้ได้ทันทีมีสองข้อ แม้ไม่ได้อยู่บน TPU: หนึ่ง ออกแบบ training loop ให้มองการล้มของฮาร์ดแวร์เป็น “เหตุการณ์ที่ต้องจัดการ” ไม่ใช่ “จุดจบของงาน” — การเขียน checkpoint บ่อยและตรวจความสมบูรณ์ด้วย marker แบบ commit_success เป็นแพตเทิร์นที่ยกไปใช้กับ PyTorch/DeepSpeed ได้เลย สอง ต้นทุนแฝงของงานเทรนที่เปราะบางคือค่าเช่า GPU/TPU ที่จ่ายไปตอนงานล้มแล้วต้องเริ่มใหม่ ในบริบทที่งบ compute จำกัด ความทนทานคือการประหยัดเงินโดยตรง ไม่ใช่แค่ความสะดวก — และเป็นเรื่องที่ควรถามผู้ให้บริการคลาวด์ก่อนจองเครื่องล็อตใหญ่

คำเตือนตามความเป็นจริง: ตัวเลขในบล็อกมาจากการสาธิตบนโมเดลจิ๋ว (Qwen3-0.6b) และ 48 ชิป การกู้คืนในสเกลจริงหลายพันชิปกับโมเดลหลายร้อยพันล้านพารามิเตอร์จะมีตัวแปรมากกว่านี้มาก โดยเฉพาะเวลากู้ checkpoint ที่โตตามขนาดโมเดล ถึงอย่างนั้น แนวคิด “เปลี่ยนการล้มให้เป็น exception ที่ดักจับได้” ก็เป็นทิศทางที่ถูกต้องและควรจับตา

แหล่งข่าว:

• We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText — Google Developers Blog (6 ก.ค. 2026)
• MaxText — GitHub (AI-Hypercomputer)