Amazon’s TorchNeuron lets PyTorch code run on AWS Trainium chips

Amazon has launched TorchNeuron, an open-source backend that lets PyTorch code run on its Trainium AI chips with a one-line change. It goes into public beta on 20 October under the Apache 2.0 licence and supports Trainium2 and Trainium3. Some PyTorch features are not yet supported.

Amazon has launched TorchNeuron, a tool that lets developers run their existing PyTorch code on its Trainium AI chips without rewriting it, the company said on Monday. PyTorch is the most widely used software framework for building AI models.

TorchNeuron will go into public beta on 20 October at the PyTorch Conference in San Jose, according to a PyTorch blog post by Amazon engineers. It will be open source under the Apache 2.0 licence and supports Trainium2 and Trainium3.

Until now, moving AI code to a new chip often meant months of rewriting and testing. With TorchNeuron, developers point their code at Trainium by changing a single device setting, Amazon said.

In a demo, Amazon ran the same training script on a standard chip and on Trainium, and said the model learned at the same rate on both. The tool plugs into PyTorch’s built-in system for adding new hardware, so standard testing and scaling tools still work.

However, the first release has gaps. It does not yet support some PyTorch features, such as dynamic input shapes, and about 22,000 PyTorch tests are skipped, according to the PyTorch blog post.

Early users include voice AI startups Cartesia and Deepgram, as well as the Lieber Institute for Brain Development, which uses it in drug research. Hugging Face has also added Trainium support to its AI tools. Uber recently signed up to use the chips in an AWS expansion deal.

The launch is part of Amazon’s push to compete with Nvidia, whose chips and CUDA software dominate AI computing. In April, chief executive Andy Jassy said Amazon’s custom chip business could be worth $50B, and Amazon now says it has passed a $25B annual revenue run rate.

One of TorchNeuron’s launch partners is Decart, an AI chip software startup that Amazon was reported to be in talks to buy.