AWS เปิดซอร์ส TorchNeuron: แค่ .to('neuron') โค้ด PyTorch ก็รันบน Trainium ได้ — public beta (Apache-2.0) แต่ torch.compile รับแค่ shape คงที่, op ที่ยังไม่รองรับแอบไปรันบน CPU และมี wheel ให้แค่ Ubuntu 26.04

AWS เผยแพร่บทความบน PyTorch Blog เมื่อ 19 ต.ค. 2026 เปิดตัว TorchNeuron ซึ่งเป็น backend แบบ native ของ PyTorch สำหรับชิป AWS Trainium และในวันถัดมา (20 ต.ค.) ซึ่งเป็นวันงาน PyTorch Conference North America ก็เปิดซอร์สโค้ดที่ github.com/aws-neuron/torch-neuronx ภายใต้ Apache-2.0 ในสถานะ public beta (release Neuron 2.33.0)

เดิมทีการใช้ PyTorch บน Trainium ต้องผ่าน PyTorch/XLA ที่มีวิธีคิดคนละแบบกับ PyTorch ปกติ ทั้งการ trace graph ที่ทำให้ debug ยาก และการต้องแก้โค้ดเฉพาะชิป TorchNeuron เปลี่ยนแนวทางมาเป็นการลงทะเบียน Trainium เป็น device ปกติผ่าน PrivateUse1 ซึ่งเป็นช่องทางที่ PyTorch เปิดไว้ให้ฮาร์ดแวร์ภายนอกโดยเฉพาะ ผลคือย้ายโมเดลไปรันได้ด้วยคำสั่ง .to("neuron") ส่วน torch.nn, torch.autocast, torch.distributed และ torch.profiler ใช้ได้ตามเดิม

ใช้งานยังไง

import torch  # torch-neuronx ถูกโหลดอัตโนมัติ

model = MyModel().to("neuron")
torch.compiler.set_default_backend("neuron")  # ถ้าจะใช้ torch.compile
model.compile()

• Eager เป็นค่าเริ่มต้น — รันทีละ op เหมือน GPU ใช้ print หรือ breakpoint ได้ตามปกติ แต่ทุก op จะถูกคอมไพล์เป็นไฟล์ NEFF ในครั้งแรกที่ถูกเรียก จึงช้าในรอบแรก และมี adaptive eager (ยังเป็น experimental) ที่รวมหลาย op เป็นหน่วยเดียว เปิดได้ด้วย TORCH_NEURONX_ENABLE_CONCATENATION=1
• torch.compile — ใช้ TorchDynamo → AOTAutograd → Torch-MLIR → StableHLO → Neuron compiler
• Distributed — มี c10d backend ชื่อ neuron รองรับ DDP, FSDP2, tensor/sequence parallel ผ่าน DTensor และ pipeline parallel โดย collective วิ่งบน collective core ภายในเครื่อง และใช้ EFA ข้ามเครื่อง
• Custom kernel — เขียนด้วย NKI แล้วลงทะเบียนเป็น custom op ด้วย torch_neuronx.nki_op ใช้ได้ทั้ง eager และใน compiled graph
• ระบบนิเวศ — Hugging Face รองรับ Trainium เป็น device ใน Transformers v5, Diffusers, TRL และ kernels แล้ว (device_map="neuron") ส่วน TorchTitan มี fork torchtitan-neuron ให้ใช้ และฝั่ง inference ใช้ plugin vLLM Neuron / vLLM Omni Neuron

ความต้องการระบบ: ต้องใช้อินสแตนซ์ Trn2 หรือ Trn3, Python 3.12–3.14, และ PyTorch 2.13 หรือ 2.14 โดย README ระบุให้ pin torch-neuronx==2.14.* เพราะ ชื่อแพ็กเกจซ้ำกับของเก่า — torch-neuronx เวอร์ชันก่อนหน้านี้คือตัวที่ใช้ XLA

ข้อจำกัดของ beta ที่ต้องรู้ก่อนย้ายงาน

ข้อจำกัดเหล่านี้ไม่ได้อยู่ในบทความเปิดตัวของ Amazon แต่อยู่ใน README และ release notes ของ repo:

• torch.compile รองรับแค่ dynamic=False — ไม่รองรับ dynamic shape ทุก shape ใหม่ต้องคอมไพล์ใหม่ จึงต้อง pad หรือแบ่ง bucket เอง ส่วน mode รับเฉพาะ "default" และ option ของ Inductor จะไม่มีผล
• op ที่ไม่รองรับใน eager จะไปรันบน CPU — งานไม่ล้ม แต่ช้าลงโดยไม่มีอะไรเตือนชัดๆ ต้องเรียก torch_neuronx.get_fallback_ops() เพื่อดูว่ามี op ไหนตกไปที่ CPU บ้าง
• float64 ถูกรันเป็น float32 (มี warning) และการรองรับ int64 ขึ้นกับแต่ละ op
• FlexAttention ยังรองรับได้จำกัด ส่วน TorchAO / MXFP8 / FP8 ยังอยู่ใน private beta
• wheel สำเร็จรูปมีเฉพาะ Ubuntu 26.04 ถ้าใช้ 22.04 หรือ 24.04 ต้อง build จากซอร์สไปก่อน (release ถัดไปจะเพิ่มให้) หรือใช้ Neuron DLAMI / container ที่ติดตั้งมาให้แล้ว
• ฟีเจอร์ที่ตั้งใจยังไม่ทำในรุ่นแรก: sparse tensor, NestedTensor, forward-mode autograd, attention mask แบบกำหนดเอง (รองรับแค่ไม่มี mask กับ causal) และ per-channel quantization

ตัวเลขทดสอบ — อ่านให้ครบก่อนเชื่อ

• Hugging Face 100k model sweep (ข้อมูล ก.ย. 2026): eager รันผ่าน 96,618 จาก 101,076 โมเดล (95.6%) และ compile ผ่าน 95,097 จาก 98,420 (96.6%) โดยใช้คอนฟิกที่ย่อขนาดลง บทความบอกว่าได้บันทึกด้วยว่า op ไหนรันบนชิปและ op ไหนตกไป CPU แต่กราฟไม่ได้แสดงตัวเลขนี้ คำว่า “ผ่าน” ในกราฟจึงไม่ได้แปลว่ารันบน Trainium ทั้งหมด
• Upstream test suite ของ PyTorch (ข้อมูล ส.ค.): ATen ผ่าน 94.7%, Dynamo 96.6%, distributed 94.7%, profiler 100% แต่เปอร์เซ็นต์นี้คิดจากเฉพาะเทสต์ที่เปิดรัน บทความระบุว่ามีประมาณ 39K เทสต์ที่ข้ามโดยเจตนา ในจำนวนนี้ราว 22K เป็นเทสต์ของ op ที่ยังไม่มี lowering ลง Neuron และตัวเลขในกราฟ (รันรวม 96,789 เทสต์) ก็ไม่ตรงกับตัวเลขในเนื้อความ (~154K ผลลัพธ์ ข้ามไป ~69K ซึ่งจะเหลือราว 85K)
• ตัวเลขด้านราคามีเพียงคำพูดของ Deepgram ว่า Trainium ประมวลผล “input audio seconds per dollar” ได้มากกว่า GPU ที่ใช้อยู่ราว 60% ซึ่งได้มาจากโปรเจกต์พอร์ตที่ทำร่วมกับทีม AWS นานหลายเดือน ไม่ใช่การรันโค้ดเดิมโดยแค่เปลี่ยนบรรทัดเดียว

มุมมอง AI Curator

ประเด็นสำคัญของข่าวนี้ไม่ใช่ตัว Trainium แต่เป็นการที่ PrivateUse1 กลายเป็นช่องทางมาตรฐานสำหรับชิปที่ไม่ใช่ NVIDIA ก่อนหน้านี้ ข้อได้เปรียบของ CUDA ส่วนหนึ่งมาจากการที่ชิปค่ายอื่นต้องให้นักพัฒนาใช้ stack แยก (XLA, SDK เฉพาะ) ซึ่งทำให้การย้ายโค้ดมีต้นทุนสูง AWS ถอยออกจากแนวทาง XLA ของตัวเองแล้วมาใช้แนวทาง device แบบ native เหมือนที่ Apple MPS และ Intel XPU ทำมาก่อน แถมยังส่ง PR เข้า upstream เพื่อลบโค้ดที่ผูกกับ CUDA ด้วย ผลระยะยาวคือ “เปลี่ยน device string แล้วรันได้” จะกลายเป็นมาตรฐานขั้นต่ำที่ชิปทุกค่ายต้องทำได้

แต่ “รันได้” กับ “รันได้เร็ว” ไม่ใช่เรื่องเดียวกัน beta นี้มีเงื่อนไขที่ส่งผลต่องานจริงชัดเจน อย่างแรก ไม่รองรับ dynamic shape งาน LLM inference ที่ความยาว sequence เปลี่ยนตลอดจึงต้องจัด bucket เอง อย่างที่สอง CPU fallback ที่ไม่มีการเตือน อาจทำให้ benchmark ภายในทีมออกมาแย่โดยไม่รู้สาเหตุ ทีมที่ทดลองควรเรียก get_fallback_ops() ทุกครั้งก่อนสรุปผล และควรดูตัวเลข 95–96% ด้วยความระมัดระวัง เพราะไม่ได้บอกว่ามีงานกี่ส่วนที่รันบนชิปจริง

สำหรับทีมในไทยและอาเซียน ข้อจำกัดหลักคือการเข้าถึงอินสแตนซ์ Trn2/Trn3 ซึ่งยังมีเฉพาะบาง region ทีมที่ fine-tune โมเดลภาษาไทยบน GPU ที่เช่าอยู่ควรมองว่านี่คือทางเลือกด้านต้นทุนที่ทดลองได้ง่ายขึ้นมาก โดยไม่ต้องเขียน training loop ใหม่ ลองรัน python -m torch_neuronx.collect_env แล้วเทียบต้นทุนต่อ token กับงานจริงของทีมเอง อย่าใช้ตัวเลข 60% ของ Deepgram เป็นเกณฑ์ เพราะเป็นตัวเลขของงานเสียงที่ผ่านการพอร์ตร่วมกับ AWS มาแล้ว

แหล่งข่าว

• Introducing TorchNeuron: Native PyTorch backend for AWS Trainium — PyTorch Blog
• aws-neuron/torch-neuronx (README, ข้อจำกัด, release v2.33.0) — GitHub
• Amazon opens its AWS Trainium AI chips to every PyTorch developer — About Amazon