Google เปิดซอร์ส TorchTPU: torch.device("tpu") ใช้ได้จริง พร้อม vLLM/SGLang บน TPU และประกาศเลิก PyTorch/XLA (archive 15 พ.ย.) — แต่ matmul fp32 ค่าเริ่มต้นคำนวณแบบ bf16 ไม่ใช่ fp32 เต็มแบบ CUDA

Google ปล่อยงานชุดใหญ่สำหรับ PyTorch บน TPU เมื่อ 20 ต.ค. 2026 ซึ่งตรงกับช่วง PyTorch Conference ข่าวหลักคือ repo TorchTPU เปิดเป็นสาธารณะภายใต้ Apache-2.0 และ torch_tpu 0.2.0 ขึ้น PyPI แล้ว (อัปโหลดเมื่อ 19 ต.ค.) TorchTPU เป็น backend แบบ native ที่ทำให้ TPU เป็น device ปกติของ PyTorch เขียน torch.device("tpu") หรือ .to("tpu") ได้เลย ในวันเดียวกัน Google ประกาศ inference engine ใหม่ vLLM-TorchTPU (พัฒนาร่วมกับ Inferact) และ SGLang-TorchTPU (ร่วมกับ RadixArk) เผยแพร่ TorchTitan เวอร์ชัน TPU ที่ทำร่วมกับ Meta และ Hugging Face ประกาศว่า Transformers, Diffusers และ Kernels รันบน TPU ได้โดยตรงแล้ว

ข่าวที่สำคัญไม่แพ้กันซ่อนอยู่ใน README ของ pytorch/xla ซึ่งถูกแก้เมื่อ 19 ต.ค. ให้แสดงว่า PyTorch/XLA เลิกพัฒนาแล้ว และ repo จะถูก archive ในวันที่ 15 พฤศจิกายน 2026 ใครที่ยังมีงาน torch_xla อยู่จึงมีเวลาย้ายไม่ถึงหนึ่งเดือนก่อนที่ repo จะหยุดรับการแก้ไข ข่าวนี้มาในวันเดียวกับที่ AWS เปิดซอร์ส TorchNeuron สำหรับ Trainium ผู้ให้บริการคลาวด์รายใหญ่สองรายจึงเลิกใช้แนวทาง XLA แบบ lazy tensor พร้อมกัน แล้วหันมาทำ device แบบ native ของ PyTorch เหมือนกัน

จาก torch_xla มาเป็น device ปกติ: อะไรเปลี่ยน

ตามเอกสาร PyTorch on Cloud TPU ของ Google ปัญหาของ PyTorch/XLA คือใช้ lazy tensor ที่บันทึกกราฟไว้ก่อนแล้วค่อยรัน โค้ดจึงไม่ได้ทำงานในจังหวะที่ผู้เขียนคาดไว้ ต้องใส่โค้ดกำหนดขอบเขตที่จะคอมไพล์ และการแก้เล็กน้อยก็ทำให้ต้องคอมไพล์กราฟใหม่ได้ TorchTPU เปลี่ยนมาทำงานแบบ eager เป็นค่าเริ่มต้น โดยใช้หนึ่ง process ต่อหนึ่ง device ตามโมเดลการทำงานปกติของ PyTorch และไม่ต้องพึ่ง runtime ของ JAX

• โหมด eager มี 3 แบบ: Strict (ค่าเริ่มต้น ส่ง op ทีละตัวแบบ async), Debug (รันแบบ sync พร้อมตรวจ bounds) และ DeferAndFuse ที่รวม op หลายตัวเข้าด้วยกันเพื่อลดการอ่านเขียน HBM
• Graph mode: ใช้ torch.compile(backend="tpu") ซึ่งแปลงผ่าน TorchDynamo → StableHLO → OpenXLA และแคชผลคอมไพล์ข้าม process, host และการรีสตาร์ทงานได้
• Distributed: ใช้ backend tpu_dist ที่รองรับ DDP, FSDP2 และ DTensor ส่วน topology_aware_mesh ช่วยวางมิติ parallelism ให้ตรงกับการเชื่อมต่อจริงระหว่างชิป
• Kernel: เรียก Pallas kernel เป็น custom op ของ PyTorch ได้ และใช้ kernel สำเร็จรูปจาก Tokamax ได้ ส่วน Helion ยังอยู่ในขั้นทดลอง
• ชิปที่รองรับ: v5e, v5p, v6e (Trillium) และ v7x (Ironwood) ส่วน TPU 8t/8i อยู่ใน roadmap ระยะใกล้

ฝั่ง Hugging Face บอกว่า ตอนนี้ไม่ต้องใช้ Optimum TPU แล้ว แค่ส่ง device_map="tpu" ใน from_pretrained หรือเรียก pipe.to("tpu") ใน Diffusers ก็พอ ทีม HF ใส่ตัวเลขจริงไว้ด้วย ได้แก่ ชุดทดสอบของ Transformers ที่รันบน TPU ผ่าน 96.5% ของเคสที่รันได้ ทีมแจ้ง issue ไปที่ TorchTPU 34 เรื่องและแก้แล้ว 19 เรื่อง ส่วน ragged-dot kernel จาก Tokamax ทำให้ decode เร็วขึ้น 21.5 เท่า บน Qwen3-30B-A3B และ 12.8 เท่าบน Gemma 4 26B-A4B ที่ batch 64 บน v6e-8 แต่ตัวเลขนี้เทียบกับ implementation เริ่มต้นของ Transformers ไม่ได้เทียบกับ GPU

Serving และ training: ตัวเลขที่ Google อ้าง

• vLLM-TorchTPU / SGLang-TorchTPU ใช้ model definition ของ PyTorch จาก upstream แล้วเสริม kernel เฉพาะ TPU ด้านล่าง รองรับ prefix caching, KV offload ไป host/Managed Lustre ผ่านไลบรารีโอเพนซอร์ส TPU Sync, speculative decoding (MTP, EAGLE-3, DFlash, DSpark, n-gram), P/D disaggregation และ checkpoint FP8/FP4/INT4 หัวหน้าทีม PyTorch-TPU ของ Meta ยืนยันในบล็อกว่า SGLang-TorchTPU “รันใน production ที่ Meta แล้ว” และ Pallas megakernel ของ Kimi K3 (2.8T) ให้ความเร็ว 700+ token/วินาที สำหรับผู้ใช้คนเดียว บน Ironwood ส่วน engine ที่ใช้ JAX เดิม (vLLM-JAX, SGLang-JAX) จะยังได้รับการดูแลต่อ
• TorchTitan บน TPU (บล็อกเชิงเทคนิค) เป็น fork ที่ google-pytorch/torchtitan ไม่ใช่ upstream ของ Meta ตัวเลขที่อ้างบน v7x คือ Llama 3 405B ได้ราว 57% MFU ที่ 128–1,024 ชิป, Qwen3 235B ได้ 32.42% MFU ที่ 1,024 ชิป และ FLUX.1 ทำ throughput ได้ราว 99% ของ MaxDiffusion (JAX) บนฮาร์ดแวร์เดียวกัน แต่ Google ก็เขียนเองว่าการ scale Qwen3 235B ไปถึง 8,192 ชิปเป็น “การทดลองด้าน scaling ไม่ใช่ recipe สำหรับเทรนจริง”

ตัวเลขเหล่านี้ต้องผ่านการปรับจูนตาม playbook หลายขั้น ไม่ใช่สิ่งที่ได้ทันทีหลังติดตั้ง ในกรณี FLUX.1 โหมด eager กับ FSDP2 แบบไม่ปรับอะไรใช้เวลา 11.49 วินาทีต่อ step ต้องเปลี่ยนไปใช้ SimpleFSDP, splash attention, per-block torch.compile และแก้โค้ดโมเดลอีก 3 จุดจึงลดลงเหลือ 2.78 วินาที เอกสารยังระบุตรง ๆ ว่า FSDP2 บน TPU “low perf” และยังไม่รองรับ torch.Stream

สิ่งที่ต้องรู้ก่อนย้ายโค้ด

1. ความแม่นยำของ matmul ค่าเริ่มต้นไม่เหมือน CUDA ตาม เอกสาร numerical precision ค่าเริ่มต้นของ float32_matmul_precision บน TPU คือ "medium" ซึ่งปัด input fp32 ลงเป็น bf16 (mantissa 7 บิต) แล้วสะสมผลเป็น fp32 ในขณะที่ PyTorch บน CUDA ใช้ค่าเริ่มต้น "highest" (fp32 เต็ม) ถ้าต้องการความแม่นยำใกล้ fp32 ต้องตั้ง "high" ซึ่งใช้ MXU ประมาณ 3 เท่า หรือ "highest" ราว 6 เท่า โค้ดเดิมจึงอาจรันได้แต่ผลตัวเลขต่างจากบน GPU โดยไม่มีคำเตือน (มีคำเตือนเฉพาะตอนเลือกโหมดที่สูงกว่า)
2. “โอเพนซอร์ส” ครอบคลุมแค่ตัว backend โค้ดใน repo เป็น Apache-2.0 แต่ wheel ต้องพึ่ง libtpu>=0.0.50 ซึ่งเป็น runtime ที่ PyPI ระบุ license ว่า “Google Cloud Platform Terms of Service” ส่วนที่คุยกับชิปโดยตรงจึงยังเป็นของ Google เหมือนเดิม แม้จะไม่ต่างจาก CUDA driver ในทางปฏิบัติ
3. ข้อมูลเวอร์ชันในเอกสารไม่ตรงกัน README ระบุว่ารองรับ Python 3.11+ และมี wheel cp311 บน PyPI แต่ metadata ของ wheel 0.2.0 เขียน Requires-Python: >=3.12 และรองรับ torch>=2.11,<=2.15 จึงควรล็อกเวอร์ชันไว้ และติดตั้งด้วย --extra-index-url https://download.pytorch.org/whl/cpu ไม่เช่นนั้น pip จะดึง torch รุ่น CUDA ขนาดใหญ่มาโดยไม่จำเป็น
4. สถานะยังไม่เท่ากันทุกส่วน ตาราง feature set ระบุว่า Qwen, Llama, Gemma, Muse Glimmer, WAN และ FLUX ผ่านการ validate แล้ว ส่วน GLM, Kimi, DeepSeek, GPT-OSS และ Mamba ยังอยู่ในหมวด experimental และ expert parallelism สำหรับ MoE ใน TorchTitan ยังไม่เสร็จ

มุมมอง AI Curator

วันนี้ผู้ให้บริการคลาวด์รายใหญ่สองรายประกาศแนวทางเดียวกันห่างกันไม่กี่ชั่วโมง AWS มี TorchNeuron ส่วน Google มี TorchTPU และเลิกใช้ PyTorch/XLA ไปพร้อมกัน ข้อสรุปคือ lazy tensor ของ XLA แพ้แล้วในฝั่ง PyTorch มาตรฐานใหม่สำหรับชิปที่ไม่ใช่ NVIDIA คือ “เป็น device ปกติ, eager เป็นค่าเริ่มต้น, ใช้ torch.compile เมื่อต้องการความเร็ว” ซึ่งเป็นแนวทางเดียวกับที่ Apple MPS และ Intel XPU ใช้มาก่อน ความได้เปรียบของ CUDA ที่เหลืออยู่จึงไม่ใช่เรื่อง “พอร์ตโค้ดได้หรือไม่” อีกต่อไป แต่เป็นเรื่อง kernel, ความครบของ op และความแม่นยำเชิงตัวเลขที่ใช้แทนกันได้จริง ซึ่งข้อ 1 ข้างบนแสดงให้เห็นว่าเรื่องนี้ยังไม่จบ

จุดที่น่าสนใจที่สุดคือการที่ Google ร่วมมือกับทีมผู้สร้าง vLLM (Inferact) และ SGLang (RadixArk) โดยตรง ไม่ได้ทำ fork ของตัวเอง สำหรับคนที่ทำ serving ข้อดีคือ API server, scheduler และ config ที่ใช้อยู่กับ GPU ใช้ต่อบน TPU ได้เกือบทั้งหมด การเปลี่ยนฮาร์ดแวร์จึงกลายเป็นการตัดสินใจเรื่องราคามากกว่าเรื่องวิศวกรรม และคนที่ได้ประโยชน์มากที่สุดคือลูกค้าที่ต่อรองราคา GPU กับผู้ให้บริการได้

สำหรับทีมในไทยและอาเซียน ข้อจำกัดแรกคือที่ตั้ง ตามหน้า regions/zones ของ Cloud TPU ยังไม่มีโซน TPU ในอาเซียน โซนที่ใกล้ที่สุดคือไต้หวัน (asia-east1) และโตเกียว (asia-northeast1) งานเทรนหรือ batch ไม่ได้รับผลกระทบมาก แต่งาน serving ที่ต้องการ latency ต่ำหรือมีข้อกำหนดให้เก็บข้อมูลในประเทศต้องคิดเรื่องนี้ก่อน ทีมที่ fine-tune โมเดลภาษาไทยบน GPU ที่เช่าอยู่จึงทดลองได้ทันทีโดยไม่ต้องเรียน JAX ก่อน ขั้นแรกที่แนะนำคือเอา eval เดิมมารันเทียบผลบน GPU กับ TPU โดยตั้ง float32_matmul_precision ให้ตรงกันทั้งสองฝั่ง ก่อนเทียบต้นทุนต่อ token ส่วนทีมที่มีโค้ด torch_xla ใน production ควรเริ่มวางแผนย้ายตั้งแต่สัปดาห์นี้ เพราะ archive วันที่ 15 พ.ย. หมายความว่าจะไม่มีการแก้บั๊กอีกแล้ว

แหล่งข่าว

• Announcing New vLLM and SGLang Inference Engines on TPU — Google Developers Blog
• Integrating TorchTPU into TorchTitan and Optimizing Models for TPU — Google Developers Blog
• Hugging Face ❤️ TorchTPU: TPUs, Natively — Hugging Face Blog
• google-pytorch/torch_tpu (LICENSE, feature set, numerical precision) — GitHub