Amazon releases TorchNeuron to let PyTorch developers run AI on its own chips

Amazon Web Services released a tool called TorchNeuron on Monday that lets AI startups move their apps from other hardware onto Amazon’s own AI chip, Trainium. The move is a direct challenge to the software edge that has made Nvidia the default hardware choice for AI development.

TorchNeuron works with PyTorch, the free, open-source framework developers use to build, train and run AI models. Amazon says about 85% of AI apps are written in PyTorch. Facebook created the framework, and the Linux Foundation, a nonprofit that maintains open-source software, has overseen it since 2022. Nvidia has spent years making sure its CUDA software, the layer that lets code run on Nvidia chips, works smoothly with PyTorch. That work has made Nvidia hardware the default option for most AI developers.

TorchNeuron is built on PyTorch, so developers who already know the framework do not need to learn a new one. Amazon says a single line of configuration change, not a code rewrite, is enough to run an app on Trainium instead. Amazon is not the first to offer this kind of shortcut; PyTorch users can already shift workloads to AMD chips or Apple hardware with a similar change.

AWS argues its version offers something more: access to Amazon’s full cloud, not just a chip. “You’re building on AWS. You’re not betting on a single-purpose compute provider and scrambling to build everything else around it,” said Colin Brace, vice president of technology software at AWS Annapurna Labs, the unit behind Amazon’s custom chips, in comments to TechCrunch. He contrasted that with neoclouds such as CoreWeave, Nebius and Firmus, which rely heavily on Nvidia hardware.

AWS says Trainium’s advantage comes from handling both training, the process of teaching a model, and inference, the process of putting a finished model to work, on the same chip within infrastructure startups already use. Brace pointed to Deepgram, a voice-AI company that competes with ElevenLabs, saying it moved to Trainium using native PyTorch support and now processes 70% more audio per dollar at production-grade accuracy, according to AWS.

Rewriting PyTorch’s ties to Nvidia

Amazon’s ambitions extend past its own hardware. The company says it is contributing code changes directly to the PyTorch project to strip out parts of the framework tied to specific chip makers, with the goal of making PyTorch work equally well across hardware vendors, not just Trainium. “The real story is removing the lock-in that kept developers from having a real choice at all,” Brace said. Whether that succeeds is unclear. PyTorch’s appeal has always rested partly on shielding developers from the complexity of matching code to specific hardware. What happens next depends on whether the PyTorch community adopts Amazon’s upstream changes, and whether developers find the switch to Trainium as simple as Amazon describes once they try it outside AWS’s own examples.