Amazon open-sources TorchNeuron to run PyTorch on Trainium

Amazon is going after the switching cost that keeps AI teams on Nvidia, and it is giving the tool away.

Amazon says TorchNeuron, a native PyTorch backend for its Trainium2 and Trainium3 chips, lets existing PyTorch code run on Trainium with "a single configuration change and no rewrite." It is being released as open source under the Apache 2.0 license at PyTorch Conference North America on Tuesday, October 20.

What Amazon is releasing

TorchNeuron plugs into PyTorch's built-in mechanism for supporting new hardware, according to Amazon. That means the testing, debugging and scaling tools you already use with PyTorch keep working without changes, the company says.

Amazon says 85% of AI researchers and developers write their code in PyTorch, and that it tests TorchNeuron against hundreds of thousands of AI models on every build.

In a demo Amazon describes, its engineers ran the same training script on a conventional chip and on Trainium. The only difference was the setting that told the code which chip to use. Amazon says the model learned the same things at the same rate.

Amazon names three early users:

• Cartesia, a conversational AI company, ran its standard inference pipeline with minimal changes to its PyTorch code.
• The Lieber Institute for Brain Development is using Trainium to explore drug discovery and molecular dynamics.
• Hugging Face added native Trainium support to its developer tools.

The moat was never only the chip

Nvidia's hold on AI training has as much to do with software as silicon. Your code, your kernels and your team's habits were all built for one vendor's stack. Moving meant rewriting, and rewriting meant risk nobody wanted to sign off on.

TorchNeuron is Amazon's bet that if moving costs one config line, price will decide where you train and serve.

Fair warning: Amazon's post gives no performance numbers and no price comparison for TorchNeuron. "Runs" and "runs fast and cheap" are different claims, and Amazon has only made the first one so far.

What it means if you train or serve on AWS

If you already run PyTorch workloads on AWS, this lowers the cost of a test. You can point an existing job at a Trainium instance and compare speed and bill against your current GPUs without a port.

Start with something you can measure end to end: a fine-tuning run or a batch inference job, not your flagship training run. Track time to finish, total cost, and whether the outputs match.

If you depend on custom CUDA kernels, expect more work than one line. The promise covers standard PyTorch code. Anything written for Nvidia's hardware directly is still yours to port.

And because it is Apache 2.0, you can read what it does before you trust it.

What to watch

The repository and documentation that land at the conference, and how much of PyTorch they cover on day one.

Independent benchmarks of Trainium against Nvidia GPUs on the same PyTorch code, run by someone other than Amazon.

Run one real job on it this month. Then decide.