Google สร้างซ้ำ OLMo 3 7B บน TPUs สำเร็จด้วย MaxText

24 ก.ย. 2026

OLMo 3 พัฒนาโดย Allen Institute for AI (AI2) เป็นโมเดลภาษาแบบเปิดเต็มรูปแบบ (fully open) ระดับล้ำสมัยที่ฝึกฝนด้วยสถาปัตยกรรมสมัยใหม่และสูตรการฝึกอบรมแบบหลายขั้นตอน เพื่อประเมินขีดความสามารถของ MaxText บน Google Cloud TPUs ทีมงานของเราจึงตั้งเป้าที่จะสร้าง OLMo 3 7B ของ AI2 ซ้ำตั้งแต่ต้น

เราเลือก OLMo 3 เพราะเป็นโมเดล 7B ที่แข็งแกร่งและทันสมัยซึ่งฝึกฝนในระดับการใช้งานจริง (production scale) อีกทั้ง AI2 ยังเปิดเผยกระบวนการเกือบทั้งหมด รวมถึงข้อมูล โค้ด และการตั้งค่าต่างๆ ทำให้เรามีตัวอ้างอิง PyTorch และ GPU ที่เป็นอิสระสำหรับทดสอบ MaxText และ TPUs ได้อย่างแม่นยำ

การสร้างซ้ำในครั้งนี้ใช้ MaxText บน Google Cloud TPUs ทั้งในส่วน stage-1 pre-training และ stage-2 mid-training anneal โดยพิสูจน์ความสอดคล้องผ่านเมทริกซ์ข้อมูลชุดทดสอบ (held-out metrics) มากกว่าแค่ดูจากกราฟการสูญเสีย (loss curve)

ไฮไลต์สำคัญจากการทดสอบ

การแปลงโมเดลจาก PyTorch ไปเป็น JAX: สถาปัตยกรรมของ OLMo 3 ถูกพอร์ตไปยัง MaxText และตรวจสอบด้วย logit-parity โดยพบว่าที่ step-0 ตรงกับตัวอ้างอิง HuggingFace ที่ค่า KL ≈ 1.5e-3 ซึ่งเป็นระดับสัญญาณรบกวนพื้นฐาน (noise floor) ปกติของโมเดลเดียวกันแต่ต่างเฟรมเวิร์ก และให้ค่า top-1 token ตรงกันถึง 98.75% ที่บริบท 8192 token

ความน่าเชื่อถือในการรันต่อเนื่อง: ระบบ Checkpoint-and-resume สามารถเล่นซ้ำการรันได้อย่างแม่นยำ โดยการทดสอบ A/B แสดงค่า Δ = 0.000 ในทุกขั้นตอนหลังจาก resume แม้เกิดเหตุขัดข้องที่ทำให้ต้องหยุดกลางคัน ระบบก็สามารถกู้คืนและฝึกฝนต่อได้โดยไม่มีความคลาดเคลื่อนทั้งในส่วนของ loss และ perplexity

การปรับขนาดและเปลี่ยนฮาร์ดแวร์กลางคัน: เมื่อสูญเสียกำลังการผลิตไปสามในสี่ที่ขั้นตอน 1.05M ทีมงานได้รัน resume ต่อบน slice ที่เหลือเพียง 1 ใน 4 โดยไม่ต้องแก้สคริปต์ run_olmo3_7b_stage1.sh ซึ่ง throughput ต่ออุปกรณ์ยังคงเดิมภายในช่วง 1% นอกจากนี้ใน Stage 2 ยังสามารถเปลี่ยนจาก Ironwood ไปใช้ v5p ได้ทันทีโดยรักษาค่า MFU ไว้ได้ที่ 57.4%

การออกแบบร่วมเพื่อประสิทธิภาพสูงสุด: การปรับรูปร่าง attention จาก 32 heads x 128 head-dim เป็น 16 x 256 ช่วยให้ทำงานเร็วขึ้น 12.4% เนื่องจากใช้ประโยชน์จาก 256x256 MXU ของ Ironwood ได้เต็มที่ โดยที่พารามิเตอร์และ FLOPs ยังเท่าเดิม

บทเรียนจากการสร้างซ้ำ

ทีมงานพบว่ากราฟการสูญเสีย (Training loss) เพียงอย่างเดียวไม่ใช่ข้อพิสูจน์การลู่เข้าที่เพียงพอ ในช่วงขั้นตอน 0.9M ค่า loss ของ MaxText ต่ำกว่า AI2 อย่างชัดเจนจนดูเหมือนแซงหน้า แต่เมื่อตรวจสอบด้วย C4 ชุดทดสอบกลับพบว่าเกิดจากการจดจำ (memorization) เนื่องจากบั๊กใน data loader ที่ทำให้เกิดการอ่านข้อมูลซ้ำซ้อน

บทเรียนสำคัญคือการสร้างซ้ำที่ซื่อตรงต้องมีการประเมินผลด้วยชุดทดสอบในทุก landmark เพื่อดักจับบั๊กข้อมูลที่อาจปลอมตัวมาในรูปแบบของชัยชนะ รวมถึงความสำคัญของการรักษาความแม่นยำในการ resume และการแยกสูตรการฝึกออกจากโครงสร้างเครือข่ายและเจเนอเรชันของฮาร์ดแวร์

ปัจจุบัน Stage 1 และ 2 เสร็จสมบูรณ์แล้ว โดยขั้นตอนถัดไปคือ Stage 3 ที่เน้นการปรับตัวเข้ากับบริบทที่ยาว (long-context adaptation) และการทำ Post-training ผ่าน Tunix เพื่อสร้าง OLMo-3-Instruct และ OLMo-3-Think ต่อไป

ผู้ที่สนใจสามารถใช้งานสูตรการฝึกและเครื่องมือการแปลงได้ใน MaxText โดยมีตัวอย่างคำสั่งดังนี้:

# Convert AI2's step-0 weights to Orbax, then:
export OLMO_INDEX_PATH=/path/to/olmo_index_seq8192.json
export LOAD_PARAMETERS_PATH=gs://<your-bucket>/olmo/checkpoints/stage1-step0/0/items
bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh submit
# Drive the full run, auto-resubmitting on preemption:
STEPS_OVERRIDE=1414078 bash src/maxtext/trainers/pre_train/scripts/olmo/xpk_olmo3_7b_stage1.sh resume_until_done