AWS TorchNeuron: Can PyTorch on Trainium Crack Nvidia’s CUDA Lock-In?
Key takeaways
- AWS TorchNeuron, in public beta since Oct. 20, 2026, runs PyTorch on Trainium2 and Trainium3 via a neuron device.
- Deepgram reports about 60% more audio per dollar. 89% of top Hugging Face text-generation architectures run as is.
- Beta limits: no dynamic shapes in torch.compile, limited operators and about 22,000 skipped upstream tests.
Nvidia GPUs dominate AI development for more than raw performance. The CUDA software ecosystem Nvidia has built over the years matters just as much. Moving to another accelerator means rewriting code and retesting compatibility. So even a faster or cheaper chip is hard to switch to.
AWS TorchNeuron, now open source, aims to lower that barrier. It lets existing PyTorch code run almost unchanged on Amazon’s own AI accelerators, Trainium2 and Trainium3. This article covers how it works and whether it can dent Nvidia’s CUDA ecosystem.
How AWS TorchNeuron works
TorchNeuron is the software layer that makes Trainium usable from PyTorch. It is built on PrivateUse1, PyTorch’s official extension point for custom hardware, which keeps it compatible with existing development setups.
Using another vendor’s chip used to mean installing a separate framework or converting code. TorchNeuron registers Trainium as a device named neuron, so in many cases you switch by changing .to('cuda') to .to('neuron').
Code written against standard APIs such as torch.accelerator needs even fewer changes. Not every operation is supported, though, so some models need extra work. Both Trainium2 and Trainium3 are supported, though FP8 and MXFP8 data types need Trainium3, according to AWS’s TorchNeuron overview.
During development you check behavior in eager mode, which runs operations one at a time. For production, torch.compile optimizes performance.
In that step, TorchDynamo turns the model into an FX graph. Then AOTAutograd, Torch-MLIR and StableHLO carry it down to a NEFF, the executable format Trainium runs. When a specific operation has to run faster, you can write your own kernel with the Neuron Kernel Interface (NKI).
Why developers don’t leave Nvidia CUDA easily
PyTorch and CUDA are often used together, but they do different jobs. PyTorch is the framework for building and training AI models. CUDA is the software platform that makes those computations run efficiently on Nvidia GPUs.
Nvidia has spent more than a decade building libraries and developer tools around CUDA, and many PyTorch projects are tuned for it.
Even if another chip performs better, switching is hard when changing and retesting code costs time and money. This dependence on one platform, which makes moving to other products difficult, is called lock-in. Nvidia’s strength lies not only in GPU performance but also in the software environment developers have relied on for years.
Nvidia CUDA vs. AWS Trainium vs. Google Cloud TPU
Google is moving the same way. Together with Meta, it released TorchTPU, a native PyTorch backend for Google TPUs, so PyTorch runs on Cloud TPUs.
Rather than asking developers to learn a new programming model, AWS and Google let them reuse as much existing PyTorch code as possible. The goal is to cut the code changes and compatibility testing that a move to another accelerator requires.
| Item | Nvidia GPU | AWS Trainium | Google Cloud TPU |
|---|---|---|---|
| Main frameworks | PyTorch, Hugging Face | PyTorch, Hugging Face, TorchTitan | PyTorch, Hugging Face, TorchTitan |
| Device hookup | CUDA | PrivateUse1 (neuron) | PrivateUse1 (tpu) |
| Compilation | TorchDynamo, Inductor | TorchDynamo, StableHLO, NEFF | TorchDynamo, StableHLO, XLA |
| Custom kernels | CUDA C++, Triton | NKI | Pallas, Helion |
| License | Proprietary CUDA platform | Apache 2.0 | Apache 2.0 |
Early adopters and cost savings
According to Amazon, more than 30 companies and startups in speech AI, robotics, large language models and drug discovery have used TorchNeuron for early validation.
Speech AI company Deepgram moved its Nova-3 model to Trainium. CTO Adam Sypniewski said that in Deepgram’s own benchmarks, Trainium processes about 60% more input audio per dollar than its existing GPU baselines. I’d treat that 60% as a pointer until it holds up on your own workload.
That means more speech data handled for the same money. Because these are Deepgram’s own tests, whether other AI models see similar efficiency has to be checked separately. Compute cost is the same pressure behind the AI token price war between DeepSeek, OpenAI and Anthropic.
Real-time voice AI company Cartesia moved most of its existing inference pipeline to Trainium without major changes. It applied custom NKI kernels only to the repeated operations that drive response time. It cut latency by optimizing the parts that mattered instead of rewriting all of its code.
Hugging Face compatibility improved as well. Among the top 100 text-generation model architectures, the share that runs without conversion rose from 46% with Optimum Neuron to 89%.
That makes existing open-source models that much easier to move to Trainium. Even when a model runs, you still need to check processing speed and memory use to see whether it performs as required.
Compatibility problems TorchNeuron still has to solve
TorchNeuron is still in public beta. Compatibility with existing PyTorch code is much better, but a few limits deserve a check before production. First, torch.compile does not support dynamic shapes, where input sizes change, so you have to set dynamic=False.
In services such as chatbots, where prompt length keeps changing, recompiling for each new input size can take time. To reduce that, batch inputs of similar length together (bucketing) or pad them to a fixed length.
Memory handling has limits too. Non-contiguous tensors, whose data is not laid out in one continuous block, can trigger extra copies during processing. More data movement uses more memory bandwidth and can slow processing down.
Operator support is still limited. Some features, including sparse tensors and NestedTensor, were left out of the first release, and about 22,000 PyTorch upstream tests are being skipped.
If an unsupported operation falls back to the CPU, overall speed can drop sharply. Before adopting TorchNeuron, check each model’s operator support, compile time and memory use.
What it would take to rival Nvidia
TorchNeuron is unlikely to replace CUDA right away. Nvidia has spent years building development tools and optimized libraries, and AWS still has compatibility problems to fix. Even so, because existing PyTorch code can be brought over and tested almost as is, evaluating another accelerator is less of a burden than before. I don’t expect it to move most teams off Nvidia soon, given the beta limits above and CUDA’s long head start.
For teams on PyTorch, the practical step is a trial run. Move a model with .to('neuron'), then check operator coverage, compile time with fixed input shapes and memory use against the beta limits above. Whether that turns into a real alternative depends on how stable Trainium2 and Trainium3 prove across different AI models and how much they actually cut costs.
