Amazon’s Trainium Takes a Strike at the Training King
Source
newsletter.semianalysis.com
Author
Kimbo Chen
Date
Why it matters
Trainium moving to native eager PyTorch with torch.compile removes a major barrier to using it for research training. SemiAnalysis's hands-on view, including bugs and the CUDA gap, helps teams judge a non-NVIDIA option.
Key takeaways · AI-distilled
AWS's TorchNeuron replaces the PyTorch/XLA path, which had no eager mode and no torch.distributed, DTensor or FSDP support. SemiAnalysis says porting is mostly swapping .to("cuda") for .to("neuron"), though code calling torch.cuda.* directly still needs edits.
SemiAnalysis trained Qwen3-32B with TorchTitan on 16 Trainium2 chips at 11,099 tokenThe chunk of text a model reads and writes in — roughly three-quarters of a word — and the unit AI usage is billed in.Full definition →/s (23.89% MFU) in eager mode and 14,156 tokens/s (30.47% MFU) compiled, a 1.28x gain. A slight change to the model architecture also worked.
Compile time is the main friction: the first step took 184.6s eager and 255.7s compiled, 30-55x a normal step, and the gpt-oss 20B recipe took 27m17s to compile. AWS says upcoming releases target up to two orders of magnitude faster time-to-first-step.
Trainium3 passed 89 of 96 FlexAttention benchmarkA standard public test set for comparing AI models — the shared scoreboards behind every "model X beats model Y" claim.Full definition → cases; all 7 failures were at batch size 4 and sequence length 8,192 and hit the compile time budget rather than numerical errors. TorchAO MXFP8 training is not supported on day 0.
AWS tests across 100k+ Hugging Face models, checking CPU fallbacks, graph breaks, tensor parity and loss-curve divergence against CPU and NVIDIA baselines. SemiAnalysis rates the Neuron Explorer profiler above NVIDIA's and AMD's but found its setup confusing.
Terms in this piece · Glossary
inference — Running a trained model to get answers — the phase where AI is actually used, as opposed to trained.
token — The chunk of text a model reads and writes in — roughly three-quarters of a word — and the unit AI usage is billed in.
benchmark — A standard public test set for comparing AI models — the shared scoreboards behind every "model X beats model Y" claim.
