Trainium3 Advances Native PyTorch Training, but CUDA Gap Remains
SemiAnalysis reports that Anthropic is running production inference and reinforcement-learning rollouts on AWS Trainium2 and Trainium3, while AWS Neuron is moving Trainium toward native PyTorch training with eager mode, torch.compile and standard APIs. Its engineers found the Trainium3 direction promising after testing the TorchTitan and TorchNeuron stack, but encountered unspecified bugs, long compilation times and a platform that remains far less mature than NVIDIA's CUDA ecosystem. AWS says upcoming TorchNeuron releases will use precompiled NKI kernels and a redesigned compile stack to target faster time to first step, though that target has not yet been demonstrated.
The move toward native PyTorch training addresses a limitation of Trainium for research workloads that require broad API coverage. However, compilation delays, bugs and the CUDA maturity gap remain practical constraints.
Trainium has already proven itself in the realm of inference. Anthropic is running production inference and RL rollouts on Trainium2 and Trainium3. Historically, Trainium has been great for inference but not the best fit for research training. In inference, the model architecture changes at a much slower pace compared to training, where researchers iterate on many different experiments and ablations of model architecture and need a stack that covers a broad API surface. In the past, that is where Trainium fell short.
The AWS Neuron team is aggressively working to change this pitfall with the biggest shift being moving towards a first-class, native PyTorch training experience. Gone are the days of sprinkling xm.step() through your training loop and wrestling with a bespoke XLA-style execution model. Trainium now runs PyTorch the way PyTorch users expect, in eager mode, with torch.compile and the standard APIs.
The SemiAnalysis engineering team has been trying out the Trainium3 racks for the past couple of weeks and has found the experience to be great – we love the overall direction of where it is heading. To be clear, it isn’t perfect yet, and we have personally run into a couple of bugs ourselves. It’s nowhere close to the maturity of the CUDA ecosystem, but the direction is right. Their team is making rapid progress with the help of agentic coding.
Some source content is omitted. Continue reading on the original site.
#AWS#Anthropic#AWS Trainium2#AWS Trainium3#AWS Neuron#PyTorch#TorchTitan#TorchNeuron#Neuron Explorer#NVIDIA CUDA#Meta PyTorch