Amazon launches TorchNeuron to make moving AI apps from Nvidia to Trainium chips easier

Amazon Web Services (AWS), Amazon's cloud computing arm, on Monday, October 19, released a developer tool called TorchNeuron. It lets AI startups easily move their AI apps onto Amazon's own AI chip, Trainium. In doing so, Amazon is going after the software advantage that has helped Nvidia become the default hardware choice for AI.

How the tool works

TorchNeuron works with PyTorch, a free, open-source toolkit developers use to build, train and run AI models. According to Amazon, an app written with PyTorch can run on Trainium after a change to just one line of configuration, with no code rewrite. The company says about 85% of AI apps are written with PyTorch. Facebook originally developed it, but since 2022 it has been part of the Linux Foundation, which maintains open-source projects.

Nvidia has spent years making sure its CUDA software supports PyTorch seamlessly, so its hardware is now embedded in the toolkit essentially by default. PyTorch users can also port apps to AMD GPUs or Apple's chips with a simple configuration change.

AWS's argument

Amazon stresses that moving to Trainium also gives startups access to its entire cloud. It claims this is a better choice than neoclouds, many of which are tied to Nvidia, such as Coreweave, Nebius and Firmus. Colin Brace, VP of Technology Software at AWS Annapurna Labs, the unit that makes Amazon's custom chips, said the same chip handles both training and inference, which improves the economics of a full workload. As an example he cited Deepgram, an AI voice app provider, which he said processes 70% more audio per dollar after moving to Trainium.

Impact on PyTorch

Amazon also says it will contribute changes to the main PyTorch project to remove code that ties the framework to specific chips, aiming to make it more hardware-agnostic. Whether this gives developers more choice and cheaper options, as Amazon promises, or just more configuration headaches, remains to be seen.