Amazon’s TorchNeuron brings PyTorch code to Trainium chips

Amazon has announced TorchNeuron, a native PyTorch backend that it says lets developers run the AI code they already write on its AWS Trainium chips. Nafea Bshara and 15 co-authors introduced it Monday in a post on the PyTorch blog.

How it works: TorchNeuron registers Trainium as a standard PyTorch device through PrivateUse1, PyTorch’s official extension point for custom hardware. Developers move a model onto the chip with .to('neuron'), and compiled models need the Neuron backend named once, according to the post. The post says code written against torch.accelerator, PyTorch’s device-agnostic API, runs unmodified.

What it supports: The backend handles eager execution for research and debugging, including an experimental adaptive eager mode, plus torch.compile for production and the Neuron Kernel Interface for custom kernels. It runs on Trainium2 and Trainium3.

When: TorchNeuron will be released under the Apache 2.0 license at github.com/aws-neuron/torch-neuronx and go live as a public beta Tuesday at PyTorch Conference North America, the post says.

Why Amazon is pushing it: Amazon says 85% of AI researchers and developers write their code in PyTorch, and that rewriting work for new chips has long been a barrier, according to its announcement. The company says its custom silicon business has passed a $25 billion annual revenue run rate.