AWS TorchNeuron: PyTorch on Trainium with one config change

AWS released TorchNeuron today, a native PyTorch backend for Trainium chips. Amazon says existing PyTorch code moves to Trainium with "a single configuration change and no rewrite," and the code goes open source under Apache 2.0 at PyTorch Conference North America on October 20. If your AI product runs PyTorch on rented Nvidia GPUs, you just got a cheap way to find out whether you are locked in.

What actually happened

Per Amazon's announcement and the torch-neuronx repo:

• What it is: A PyTorch device backend registered through PyTorch's built-in hardware mechanism. You call model.to("neuron") the same way you call .to("cuda"). Eager mode is the default.
• Chips: Trainium2 and Trainium3 instances. Python 3.12–3.14, pinned to torch 2.14 and torch-neuronx 2.14.*.
• Compile: torch.compile(backend="neuron") works, but only with static shapes. Every new input shape recompiles.
• License: Apache 2.0.
• Upstream: Amazon says it will contribute changes to PyTorch to remove code that ties the framework to specific chips.
• Early users: Amazon names Cartesia and the Lieber Institute. In a separate AWS case study, Deepgram reports 70% more audio processed per dollar on Trainium.

TechCrunch quotes AWS's Colin Brace framing it as "removing the lock-in." Amazon did not publish Trainium prices or a cost comparison with this launch.

Why TorchNeuron matters for your business

Most small teams never chose Nvidia. They chose PyTorch, and Nvidia came with it. That is fine until GPU prices jump or your cloud runs short of capacity. Then "we could switch" turns out to mean a three-month rewrite.

TorchNeuron makes the switch a test you can run in an afternoon. We would do it this way:

1. Grep for cuda. Every hard-coded .cuda() or device="cuda" is lock-in you wrote yourself. Move the device into config.
2. Watch for dynamic shapes. Variable-length batches recompile on every new shape under torch.compile. Pad to fixed buckets or the speed numbers lie.
3. Price per job, not per hour. Run the same workload on both, then compare cost per finished inference or per training step. Deepgram's 70% is one company's audio workload. Yours is a different number.

You may never move. But a backend you have tested is a price you can negotiate with. One you haven't is a guess.

Sources: Amazon, aws-neuron/torch-neuronx on GitHub, TechCrunch, AWS Machine Learning Blog.