Amazon targets Trainium gains with new TorchNeuron developer tool

Amazon unveiled TorchNeuron, a new developer tool designed to make it easier to deploy PyTorch applications on its Trainium chips.

PyTorch is a free, open source toolkit that developers use to build, train, and run AI models. According to Amazon, TorchNeuron will allow developers to run their existing AI code on AWS Trainium chips by changing just one line of configuration, making it easier to adopt new hardware.

Long-term, aim appears to be aimed at providing developers with an easy-to-use alternative to Nvidia’s CUDA ecosystem for AI training and inference.

"Today, 85% of AI researchers and developers write their code in PyTorch. Until now, the hardware that code could run on was far more limited than the code itself, and expanding to new chips meant rewriting and revalidating work that took months to perfect," said the firm.

"Amazon is now closing that gap. With TorchNeuron, the PyTorch code developers already write runs natively on Amazon's AWS Trainium chips, with a single configuration change and no rewrite."

The company said that when its engineers ran the same AI training script on a conventional chip and Trainium, the AI model learned the same things at the same rate; the only difference was a setting that told the code which chip to use.

This, Amazon noted, means that the tools developers already use to test, debug, and scale their AI models all work on Trainium without modification.

A boon for startups?

PyTorch said early adopters include startups, enterprises, machine learning industry partners, and academic institutions, with applications spanning autonomous vehicles, robotics, video generation, speech, telecommunications, healthcare, and geospatial modeling.

"With PyTorch native support on AWS Trainium, we ran our standard inference pipeline with minimal changes to our PyTorch code," said Arjun Desai, co-founder of voice AI firm Cartesia.

"Small, fast ops map efficiently through the native path, so we didn’t have to hand-write kernels for them. That lets our team concentrate on writing high-performance NKI kernels for the layers that are called repeatedly and drive end-to-end latency."

TorchNeuron will roll out as free and open source software under the Apache 2.0 license at PyTorch Conference North America.

Looking ahead, Amazon said that it's contributing changes upstream to PyTorch itself, working with the community to remove code that ties the framework to specific chips.

The aim is to make PyTorch more hardware-agnostic for the entire developer community, not just Trainium users. Google is already working on TorchTPU, a similar project to run PyTorch natively on its chips.

"By reducing the friction of adopting new hardware, TorchNeuron introduces something the AI industry has largely lacked: meaningful choice," said Amazon. "When developers can move between chips without penalty, chipmakers have to earn developers' business, and developers benefit from lower costs and faster iteration."

FOLLOW US ON SOCIAL MEDIA

Follow ITPro on Google News and add us as a preferred source to keep tabs on all our latest news, analysis, views, and reviews.

You can also follow ITPro on LinkedIn, X, Facebook, and BlueSky.