Amazon opens Trainium AI chips to all PyTorch developers

Amazon Amazon.com, Inc. ↘ 0.95% has introduced TorchNeuron, software that lets developers run code written in PyTorch on its Trainium AI chips with a single configuration change and no rewriting, Reuters reported. The tool comes from Annapurna Labs, Amazon's chip unit. It targets the software advantage behind Nvidia's ↘ 0.59% lead in AI computing.reuters.com

TorchNeuron had been in private beta since late last year. On Monday, October 19, Amazon opened it to all developers as a public beta, according to GeekWire. Amazon said it will release the software as open source under the Apache 2.0 license on Tuesday, October 20, at PyTorch Conference North America. The tool supports Trainium2 and Trainium3.geekwire.com+1

How it works

Until now, developers who wanted to run PyTorch code on Trainium had to change it to call programming interfaces built specifically for Amazon's chips, Annapurna Labs executive Colin Brace told GeekWire. "If you know PyTorch, then you know Trainium," he said. TorchNeuron uses PyTorch's built-in system for registering new hardware, so the tools developers already use for testing, debugging and scaling work on Trainium unchanged, Amazon said. It also works with TorchTitan, PyTorch's distributed training framework.aboutamazon.com+1

Amazon tests the tool every night against more than 100,000 models from Hugging Face, Brace said. The company has also submitted changes to PyTorch itself to remove code that tied the framework to particular chips, GeekWire reported.geekwire.com

"This has been about meeting customers where they're at," Brace told Reuters. He said more than 30 early customers in robotics, world models and voice AI are using the software. Speech AI startup Deepgram said its benchmarks showed Trainium processing about 60% more audio per dollar than its existing GPU setup, according to GeekWire.reuters.com+1

Taking on Cuda

PyTorch is overseen by Meta ↘ 0.96% and used by most AI developers. Nvidia has spent more than a decade building Cuda, its software platform that makes PyTorch code run well on its own chips, Reuters noted. Moving that code to other hardware can take months of engineering work. That has made many companies hesitant to switch, GeekWire reported.reuters.com+1

Rivals are making similar moves. Google, part of Alphabet ↘ 1.08%, is building TorchTPU to run PyTorch natively on its chips. Microsoft ↗ 0.66% said last week it plans to run most of its AI services on its own Maia chips, GeekWire reported.geekwire.com

Trainium's reach

So far, Trainium's biggest users are AI labs with financial ties to Amazon, including Anthropic and OpenAI. In April, CEO Andy Jassy wrote that Amazon's chip business had an annual revenue run rate of more than $20 billion, according to GeekWire. Brace called TorchNeuron "an accelerant" toward wider use of the chips. Amazon's stated goal for the hardware, he said, has "always been about choice."geekwire.com