TorchNeuron explained: AWS opens Trainium chips to PyTorch

Most AI models today are written in PyTorch, and most of them are trained on Nvidia GPUs. Amazon wants to change the second part without touching the first. Its new TorchNeuron backend lets PyTorch code run natively on AWS Trainium, Amazon's in-house AI chips, with what Amazon calls "a single configuration change and no rewrite", the company says. The code goes open source on 20 October at PyTorch Conference North America.

What TorchNeuron is

PyTorch has an official way for chipmakers to plug in new hardware, known as the PrivateUse1 device backend and the torch.accelerator APIs. TorchNeuron uses that mechanism to register Trainium as an ordinary PyTorch device called neuron. According to the PyTorch Foundation blog and AWS documentation, it supports:

• eager execution for research and debugging, including a new "adaptive eager" mode;
• torch.compile for production speed;
• standard distributed training APIs such as FSDP, DTensor, DDP and tensor parallelism, plus PyTorch's TorchTitan framework;
• the Neuron Kernel Interface (NKI) for hand-written custom kernels;
• profiling with the familiar torch.profiler, linked to AWS's Neuron Explorer tool.

The package is called torch-neuronx. It is published under the Apache 2.0 licence on GitHub and shipped as a public beta with the Neuron 2.33 release.

What actually changes for developers

In practice, the main change is one word. Where code would send a model to an Nvidia GPU with "cuda", it sends it to "neuron" instead. Amazon says its engineers ran the same training script on a conventional chip and on Trainium, changing only that setting, and the model "learned the same things at the same rate". Amazon also says it tests TorchNeuron against hundreds of thousands of models on every build. Hugging Face has added native Trainium support to its tools, so developers don't need Amazon-specific code there either.

Early users named by Amazon include voice-AI company Cartesia and the Lieber Institute for Brain Development, which uses AI models in drug discovery research.

What it doesn't do

• It is not for your laptop. Trainium chips exist only in AWS data centres. TorchNeuron supports Trainium2 and Trainium3 instances. If you want to run models at home, see our guide to laptops for local AI.
• It is a beta. Beta software can have gaps, and speed claims are Amazon's own. Independent benchmarks will matter.
• It doesn't make Trainium cheaper by itself. The pitch is lower switching costs, not a price cut.

Why it matters

Nvidia's lead isn't only about chips. Its CUDA software stack keeps teams on its hardware, because moving code elsewhere takes engineering time. Amazon is attacking that lock-in directly and says it is also sending changes upstream to make PyTorch itself less tied to specific chips. It isn't alone: start-ups such as Tenstorrent, which just raised $1.5 billion, are chasing the same opening. More competition among AI chips could ease pressure on the compute supply, which is also limited by power, as Microsoft's exit from a planned UK data centre showed this week.