Amazon’s Trainium Takes a Strike at the Training King

What it claims

SemiAnalysis tested AWS's TorchNeuron, a native PyTorch backend that replaces the PyTorch/XLA route to Trainium. Their view is that it makes Trainium usable for research training, not just inference. The authors used Trainium3 racks for a few weeks. They report a positive direction with bugs and rough edges, and say the stack is "nowhere close to the maturity of the CUDA ecosystem." Anthropic is named as running production inference and RL rollouts on Trainium2 and Trainium3.

How it works

TorchNeuron supports eager mode, torch.compile, torch.distributed.*, DTensor, SimpleFSDP and DDP. Under PyTorch/XLA, users had lazy tensor capture, xm.step() calls and out-of-tree SPMD APIs such as torch_xla.experimental.spmd_fsdp. The stated porting change is replacing .to("cuda") with .to("neuron"). Code that uses the torch.accelerator namespace, such as bf16 autocast, needs no change. Code that calls torch.cuda.* directly still does.

Ecosystem support named on the page includes Hugging Face Accelerate and TorchTitan. Neuron also ships an official TorchTitan image. Its gpt-oss 20B recipe runs through GraphTrainer, TorchTitan's experimental trainer. The recipe adds Trainium configurations of SimpleFSDP, a graph-capture compile backend, and custom NKI kernels for attention and MoE.

For quality assurance, AWS tests top 100k+ Hugging Face models across normalization, MLP, MoE and attention variants. The tests cover crashes, performance, CPU-fallback ops, graph breaks, and numerics against CPU and Nvidia baselines, including tensor-level parity and loss-curve divergence. AWS uses up to 80,000 physical cores for this. The page says pass rates are not yet 100%.

Evidence

• Qwen3-32B on TorchTitan. The run used 16 Trainium2 chips, which is 64 LNC=2 logical cores. A Trainium3 UltraServer with 64 chips was also tried, but no number is given for it. Small changes to the model architecture worked.
• Time to first step. The first step took 184.6 s in eager mode and 255.7 s in compiled mode. That is 30-55x the average step time. The gpt-oss compile took 27 min 17 s.
• Matmul. The test used one Trainium3 logical core with nn.functional.linear on projection shapes from open-weight models, at M = 1, 128 and 2048. Trainium3 had higher MFU than GB300 from about 1 GFLOP per matmul. GB300 and B200 led between about 100 GFLOP and 1 TFLOP. Trainium3 led again above 1 TFLOP.
• FlexAttention. All 7 failures out of 96 cases were at batch 4 and sequence length 8192. They were compile-time budget overruns, not numerical errors. Custom attention variants outside the suite matched Nvidia numerics.

How to use it

The profiler workflow that worked for the authors has three steps:

1. Launch a persistent Neuron Explorer server.
2. Capture with neuron-explorer capture.
3. Upload the trace to the server as a web request.

Kernel profiling and PyTorch profiling setups differ slightly. Both need several environment variables and produce several file types, and the authors found this poorly documented. Once set up, the timeline shows each hardware unit on its own row and highlights flow, order and output dependencies. The authors contrast this with Nsight Compute, which gives aggregate metrics rather than an event timeline inside the kernel.

Limits

• No MXFP8. TorchAO MXFP8 is not supported at launch. AWS is prioritizing BF16.
• No TokenSwitch. PyTorch's MoE dispatch/combine API is not supported. Trainium uses a custom NeuronEPTokenDispatcher instead, and AWS says it will work with Meta to adapt the TokenSwitch ABC.
• Slow compiles. Compile time is the main friction point. Upcoming releases target up to two orders of magnitude lower time to first step, which the authors have not verified.
• Only Trainium2 numbers for training. The training throughput figures are from Trainium2. No Trainium3 training throughput is reported, and no tuned speed-of-light numbers are given.
• Source relationship. The authors are a hardware analysis firm and describe their own advocacy role with vendors. The article includes a hiring section.