SemiAnalysis: Trainium's Shift to Native PyTorch Training Is the Right Move; Compilation Time Is the Real Bottleneck for Researchers
The SemiAnalysis engineering team conducted real-world tests of the Trainium training stack and concluded that AWS's strategic correction from XLA to native PyTorch (TorchNeuron) is a significant and correct move. However, compilation time for the first training step reaches 30 to 55 times the average step duration, creating substantial friction for researchers who frequently iterate on model architectures.
For AI chips, inference and training are two entirely different engineering problems. Inference deals with forward propagation of a few fixed architectures and has a limited number of hot-path kernels; training requires a software stack that covers backward propagation for thousands of operators, multiple precision modes, and model structures that researchers may modify at any time. In a paid article published on October 23, 2026, SemiAnalysis explicitly stated that while Trainium has already proven itself in the inference field—Anthropic is running production inference and reinforcement learning on Trainium2 and Trainium3—it has historically struggled in research training because the requirements for API coverage and flexibility in training are far higher than in inference.
Amazon's response strategy involves a major software route correction. Previously, Trainium accessed the PyTorch ecosystem via PyTorch/XLA, relying on lazy tensor graph capture without a first-class eager mode, and did not support native PyTorch distributed APIs such as torch.distributed, DTensor, or SimpleFSDP. For researchers switching to Trainium, this meant relearning how to write training loops. The new TorchNeuron backend abandons this path and instead provides a native PyTorch experience: users simply replace .to("cuda") with .to("neuron"), theoretically allowing most ATen operators to run out of the box, while supporting both eager mode and torch.compile. SemiAnalysis believes this shift will significantly broaden the adoption range of Trainium3.
The SemiAnalysis engineering team conducted several weeks of practical testing on Trainium3 racks and Trainium2 hardware. Training the Qwen3-32B model under the TorchTitan framework yielded an average throughput of 11,099 tokens/second with 23.89% MFU in eager mode, and 14,156 tokens/second with 30.47% MFU in compiled mode, representing a 1.28x improvement in throughput. Training continued to function normally after they modified the model architecture. In the FlexAttention benchmark test, Trainium3 passed 89 out of 96 test cases, achieving a pass rate of 93%; all seven failed cases were compilation timeouts under large batch sizes and long sequences, rather than numerical errors. Regarding matrix multiplication performance, Trainium3's MFU was superior to Nvidia GB300 at medium scales (above 1 GFLOP), lagging only above 100 GFLOP, and regaining the lead above 1 TFLOP.
However, what truly hinders researcher adoption is not peak performance, but compilation time. SemiAnalysis's practical tests revealed that for the Qwen3-32B training run, the first step took 184.6 seconds in eager mode and 255.7 seconds in compiled mode, which is 30 times and 55 times the subsequent average step time, respectively. The compilation time for gpt-oss 20B was even higher, reaching 27 minutes and 17 seconds. The article traces the root cause of this issue to the historical design of the Trainium compiler stack: in inference scenarios, compilation is completed offline and cached once per image, so it does not affect production inference latency; but in training scenarios, researchers frequently modify model architectures, triggering recompilation with every change, which directly interrupts the iteration rhythm.
SemiAnalysis proposes an important judgment framework regarding this issue: for ML research ablation experiments, the most critical factor is not the fastest stack, but decent performance combined with broad support for model architecture modifications, enabling researchers to iterate without fighting the stack itself. This means that MFU percentage is not the sole indicator; the breadth of API coverage, numerical correctness, and iteration efficiency jointly determine the practicality of a training stack. Trainium performs well in API coverage and numerical accuracy, but compilation time—a seemingly minor engineering issue—actually constitutes the core friction preventing researchers from migrating from CUDA.
AWS has already marked solutions for this problem in its roadmap. The upcoming version of TorchNeuron introduces a new compiler backend featuring extensive precompiled NKI kernels and a re-architected fast-compilation stack, aiming to reduce the time-to-first-step by up to two orders of magnitude while maintaining minimal or zero performance regression. SemiAnalysis believes that if this promise is fulfilled, compilation time will cease to be a topic of discussion. Additionally, Trainium is rapidly investing in the vLLM/SGLang ecosystem, providing hardware access and early access to SGLang based on native PyTorch, a strategy consistent with Google's approach of integrating TPUs into the vLLM/SGLang ecosystem via TorchTPU.
It should be noted that SemiAnalysis's tests are first-party results and have not been independently reproduced; the article also explicitly acknowledges that Trainium is far from matching the maturity of the CUDA ecosystem. TokenSwitch, a first-class API in PyTorch for MoE dispatch/combine, has not yet been implemented on Trainium and currently relies on a custom NeuronEPTokenDispatcher class. The promised two-order-of-magnitude improvement in compilation time is an AWS commitment rather than a delivered result. Nevertheless, SemiAnalysis's core argument—that native PyTorch interfaces solve the API coverage problem, while compilation efficiency is the next key variable determining adoption speed in the training market—provides a more operational dimension for evaluating AI accelerator training readiness than peak FLOPS alone.