Google เปิดตัว MaxText ฝึก AI Long-Context ได้ถึง 14 ล้าน Token

9 พ.ย. 2026

การฝึกอบรมแบบ Long-context training ทั้งในช่วง pretraining และ post-training ได้กลายเป็นรากฐานสำคัญของโมเดล AI ระดับโลก ความยาวลำดับ (native sequence lengths) ที่เพิ่มขึ้นช่วยให้โมเดลภาษาขนาดใหญ่ (LLMs) มีหน่วยความจำสำหรับประมวลผลคลังโค้ดทั้งหมดหรือวิดีโอที่มีความยาวหลายชั่วโมงได้ รวมถึงเพิ่มความสามารถด้านเหตุผลในงานแบบ agentic ที่มีหลายขั้นตอน

ในช่วง 3 ปีที่ผ่านมา context window ได้ขยายจาก 128K - 200K token สู่ระดับ 1M - 2M และโมเดลที่รองรับถึง 10M token กำลังจะมาถึง อย่างไรก็ตาม การฝึกอบรมที่ความยาวระดับหลายล้าน token มักจะพบกับคอขวดสำคัญสองประการ:

1. Attention Compute: ค่า attention FLOPs จะเพิ่มขึ้นแบบยกกำลังสองตามความยาวลำดับ ทำให้การคำนวณส่วนนี้กลายเป็นภาระหลักของระบบ
2. Memory and Communication: ปริมาณ Activations จะเพิ่มขึ้นตามความยาวลำดับ การทำซ้ำข้อมูลบนทุกชิปแบบเดิมจะทำให้หน่วยความจำ (HBM) เต็มจนเกิดข้อผิดพลาด out-of-memory ทางแก้คือการใช้ context parallelism (CP) เพื่อสตรีมข้อมูลแทนการเก็บข้อมูลทั้งหมดไว้ในคราวเดียว

วันนี้ Google มีความยินดีที่จะประกาศรองรับ long-context training ใน MaxText ซึ่งเป็น open-source JAX training framework ของ Google Cloud โดยรองรับกลยุทธ์ context-parallel 4 รูปแบบบน Cloud TPU ได้แก่ All-gather, A2A (DeepSpeed Ulysses), P2P (Ring Attention) และ A2A + P2P hybrid (Unified Sequence Parallelism หรือ USP)

ทุกรูปแบบใช้งาน Tokamax Splash Attention kernel ร่วมกับระบบสื่อสาร 3D inter-chip interconnect (ICI) บน TPU รุ่นใหม่อย่าง TPU 5p และ TPU 7 (Ironwood) จากการทดสอบกับ Llama 3.1 8B พบว่าทำประสิทธิภาพได้ถึง 70% Model FLOPs Utilization (MFU) บน TPU 5p ที่ความยาว 4 ล้าน token และขยายได้สูงสุดถึง 14 ล้าน token โดย TPU 7 ยังคงรักษา 50% MFU ได้อย่างสม่ำเสมอ ฟีเจอร์นี้พร้อมใช้งานแล้วใน MaxText v0.2.4(doc)

กลยุทธ์ Context parallelism ทำหน้าที่แบ่งลำดับข้อมูลที่ยาวออกไปยังชิปต่างๆ เพื่อไม่ให้ชิปตัวใดแบกรับภาระหน่วยความจำเพียงลำพัง โดยแต่ละกลยุทธ์มีจุดเด่นที่แตกต่างกัน:

• All-gather: เรียบง่ายที่สุด แต่ละชิปจะรวบรวม keys และ values จากชิปอื่นทั้งหมด แต่ต้องมีการปรับสมดุลภาระงาน (load balancing) เนื่องจากลักษณะของ causal masking

Figure 1. All-gather สำหรับหนึ่งลำดับที่แบ่งระหว่างอุปกรณ์สี่เครื่อง (ชิป TPU) ทุกอุปกรณ์รวบรวม keys และ values ทั้งหมด จากนั้นทำการ attend เฉพาะส่วน บล็อกที่มีเส้นขอบแสดงถึง attention ที่กำลังถูกคำนวณ บล็อกเส้นขอบขนาดเล็กคือข้อมูลที่กำลังรับส่ง

• All-to-All (Ulysses): ใช้การสลับแกนการทำงาน (sharding) ระหว่าง token และ attention heads ทำให้แต่ละชิปประมวลผลลำดับเต็มสำหรับบาง head แทน วิธีนี้ไม่ต้องทำ load balancing แต่จะถูกจำกัดด้วยจำนวน attention heads ของโมเดล

Figure 2. All-to-All (Ulysses) การทำ All-to-all แลกเปลี่ยนชิ้นส่วน token กับ attention heads (แถว) แต่ละอุปกรณ์ attend บนลำดับเต็มสำหรับ head ของตนเอง และเอาต์พุตจะส่งกลับด้วย All-to-all ครั้งที่สอง

• P2P (ring attention): ทลายข้อจำกัดเรื่องจำนวน head โดยการส่งต่อข้อมูลแบบวงแหวนระหว่างชิปทีละก้าว การสื่อสารจะซ้อนทับไปพร้อมกับการคำนวณ (overlap) ช่วยให้ขยายขอบเขตบริบทได้ไกลขึ้นมาก

Figure 3. P2P (ring attention) ชิ้นส่วน KV หมุนทีละก้าว แต่ละอุปกรณ์ attend ไปยังชิ้นส่วนที่ถือไว้ในขณะที่ชิ้นส่วนถัดไปมาถึง และถือชิ้นส่วนที่มาเยือนเพียงครั้งละหนึ่งชิ้น

• A2A + P2P (USP): การรวมจุดเด่นของทั้งสองระบบเข้าด้วยกันเพื่อการขยายขนาดที่มีประสิทธิภาพสูงสุด ลดจำนวนก้าวในการสื่อสารและช่วยซ่อนการถ่ายโอนข้อมูลไว้เบื้องหลังการคำนวณได้อย่างดี

Figure 4. A2A + P2P (USP) การทำ All-to-all ทำงานภายในแต่ละคู่ของอุปกรณ์และ P2P ring ทำงานข้ามคู่ ดังนั้นแต่ละอุปกรณ์จึง attend บนชิ้นส่วนที่ยาวขึ้นด้วยจำนวนก้าวน้อยลง

การเลือกใช้กลยุทธ์ขึ้นอยู่กับขนาดของหน่วยความจำ HBM และจำนวน KV heads ของโมเดล หากหน่วยความจำเพียงพอ All-gather จะเป็นตัวเลือกที่ง่ายและเร็วที่สุด แต่หากต้องการฝึกโมเดลที่มีบริบทระดับ 14 ล้าน token กลยุทธ์แบบ P2P หรือ Ring Attention จะเป็นตัวหลักที่ให้ประสิทธิภาพสูงสุด โดยเฉพาะบน TPU 7 ที่มีหน่วยความจำเพิ่มขึ้นเป็นสองเท่า

นอกจากนี้ MaxText ยังมีระบบ Offloading เพื่อช่วยบริหารจัดการหน่วยความจำ โดยการย้ายข้อมูลบางส่วนไปยัง host memory เมื่อหน่วยความจำชิปเริ่มเต็ม ซึ่งแทบไม่มีค่าใช้จ่ายด้านเวลาเพิ่มเติมเพราะสามารถทำไปพร้อมกับการคำนวณได้

ในอนาคต Google วางแผนขยายการรองรับไปยังโมเดลที่หลากหลายขึ้น เช่น Gemma ที่ผสมผสาน sliding-window และ global layers รวมถึงกลไก linear attention อย่าง Gated Delta Net (GDN) เพื่อผลักดันขีดจำกัดของการฝึกอบรม AI ให้ทรงพลังยิ่งขึ้น

ผู้ที่สนใจสามารถใช้งาน Context parallelism ได้แล้ววันนี้ผ่าน MaxText v0.2.4 บน PyPI และเข้าถึงสูตรตัวอย่างได้ใน GitHub ของ TPU 7 และ TPU 5p