AWS launches TorchNeuron to move PyTorch workloads onto Trainium

Amazon Web Services on Monday released TorchNeuron, a developer tool that lets AI startups move PyTorch applications onto Trainium, Amazon's own AI chip, by changing a single line of configuration instead of rewriting code.

According to AWS, TorchNeuron works with PyTorch, the free open-source toolkit used to build, train and run most AI models. The tool is adapted from PyTorch itself, so developers who already know the framework can use it without learning a new stack.

The tool targets Nvidia's deepest software advantage. Amazon says about 85% of AI applications are written with PyTorch, and Nvidia has spent years making CUDA, its software layer for running code on its GPUs, support those applications seamlessly, which helped make its hardware the default choice.

Amazon is not the first company to try. PyTorch users can already port applications to AMD GPUs or even Apple chips with a simple configuration change. AWS argues that moving to Trainium is not only easy but also plugs startups into the rest of its cloud.

"You're building on AWS. You're not betting on a single-purpose compute provider and scrambling to build everything else around it," Colin Brace, VP of technology software at AWS Annapurna Labs, told TechCrunch in an email. On AWS, he added, the infrastructure to go from prototype to production is already there.

AWS also says the same Trainium chip handles both training and inference, and that tight integration with its cloud saves startups money even when its per-token price is not the cheapest. Brace cited voice AI company Deepgram, which he said moved to Trainium using native PyTorch and now processes 70% more audio per dollar at production-grade accuracy.

Perhaps more consequential, Amazon says it is contributing changes upstream to PyTorch so the framework relies less on code written for specific chips. The stated goal is to make PyTorch hardware-agnostic for the entire developer community, not just for Trainium users.

That would benefit every vendor selling hardware that competes with Nvidia, but it also raises a question: part of PyTorch's appeal is that it hides hardware complexity from developers. Whether decoupling gives them more choice or just more configuration headaches is still unresolved.

If Amazon succeeds in loosening PyTorch's ties to Nvidia, the impact could go well beyond chip competition. As Brace put it, "The real story is removing the lock-in that kept developers from having a real choice at all."