Amazon introduced TorchNeuron on 19 October, saying it lets developers run existing PyTorch code on its AWS Trainium AI chips by changing one configuration setting rather than rewriting the programme. The software is designed to remove a practical obstacle to trying Amazon’s hardware: code written for one kind of chip has often needed adaptation before it will run on another.
Key points
- Amazon says existing PyTorch code can run on Trainium after a single configuration change.
- In Amazon’s demonstration, a training script learned at the same rate on Trainium and a conventional chip.
- More than 30 early customers are using TorchNeuron, an Annapurna Labs executive told Reuters.
TorchNeuron connects Trainium to PyTorch
PyTorch is the framework in which Amazon says 85% of AI researchers and developers write code. It already provides a way for hardware to register as a device, and Amazon says TorchNeuron uses that connection to make Trainium available to existing programmes. Think of the connection as a standard socket: the programme still needs to be pointed towards the chip it should use, but its underlying instructions do not have to be rebuilt simply to make the connection. Amazon says familiar tools for testing, debugging and scaling PyTorch models can then work on Trainium without modification.
For a demonstration, Amazon engineers ran the same training script on Trainium and on what the company described as a conventional chip. Amazon says the scripts differed only in the device setting, and that the model learned the same things at the same rate in those runs. That is a result from a side-by-side test of one script, rather than a measurement of how every model or workload would perform.
Amazon also says Cartesia ran its usual voice-AI inference pipeline on Trainium with minimal changes to its PyTorch code. Inference is the work a trained model does when it produces a response, rather than when it is learning from training data. The Cartesia example concerns that response-making work, while Amazon’s side-by-side demonstration concerned training.
Having a spoken conversation could therefore involve a model running on a different chip without its existing programme being rewritten for that chip. The change would be in the computing behind the response, rather than in how the conversation begins.
Annapurna Labs reports more than 30 early customers
Colin Brace, vice president of software for Amazon’s Annapurna Labs chip unit, said more than 30 early customers in fields including robotics, world models and voice AI were using TorchNeuron, Reuters reported. Brace said moving to Amazon’s chips could bring steep cost savings. Those are his accounts of adoption and potential savings, not figures from the training demonstration.
Reuters reported that Nvidia has spent more than a decade developing Cuda, software that helps PyTorch code run efficiently on Nvidia chips. That history puts Amazon’s emphasis on existing code in context: attracting developers to another chip depends on the software around it as well as the chip itself. Google, which also makes custom AI chips, is working on PyTorch compatibility tools to attract developers.
Amazon says TorchNeuron supports both Trainium2 and Trainium3. The company also says it checks each build of the software against hundreds of thousands of AI models, with the aim of catching problems before customers run their workloads. That testing is a different kind of evidence from the reported early use: it describes how Amazon checks compatibility as the software changes, while Brace’s customer count describes uptake before its public release.
Hugging Face and TorchTitan extend Trainium support
Amazon says Hugging Face has added native Trainium support to its AI development tools, allowing work on the chips without Amazon-specific code. Amazon is also contributing changes to PyTorch itself to remove code tied to particular hardware. Those changes are intended to make the framework less dependent on any one chip, rather than confining the compatibility work to TorchNeuron.
For larger training jobs, Amazon says TorchNeuron works with TorchTitan, PyTorch’s distributed training framework, through Fully Sharded Data Parallelism. This gives developers a route to spread training across Trainium instances using an existing PyTorch approach. Amazon also offers tools for writing lower-level code, which gives developers working directly on performance a path beyond the single-setting approach.
Amazon says it acquired Annapurna Labs in 2015 and has since built four families of custom chips for different cloud workloads. TorchNeuron addresses the software side of using that hardware, with support for the two Trainium generations named by the company.
Amazon plans to release TorchNeuron as open-source software under the Apache 2.0 licence at PyTorch Conference North America on 20 October.
