Amazon says PyTorch code now runs on Trainium chips with one setting changed

Amazon says developers can now run existing PyTorch code on its AWS Trainium chips with a single configuration change. According to the official announcement, the software, called TorchNeuron, will be open-sourced under Apache 2.0 on October 20.

One device setting instead of a port

TorchNeuron is a native PyTorch backend for Trainium. Instead of asking developers to adopt a separate toolchain, it registers the chip through the mechanism PyTorch already offers for new hardware. Amazon says this means the usual tools for testing, debugging and scaling models work on Trainium unchanged.

In a side-by-side demonstration described by Amazon, engineers ran the same training script on a conventional chip and on Trainium. The only difference was the setting that names the target device. Amazon reports that the model learned the same things at the same rate. The demonstration is Amazon’s own, and the announcement does not publish loss curves, model sizes or run times.

The release supports Trainium2 and Trainium3. It also integrates with TorchTitan, PyTorch’s distributed training framework, so scaling across Trainium instances uses Fully Sharded Data Parallelism, the same pattern used on other devices. For teams that want more control, TorchNeuron includes tools for multi-chip scaling and for writing custom low-level kernels.

Who is already building on it, and what else ships

Amazon names two early users. Cartesia, a voice AI company, ran its standard inference pipeline on Trainium with minimal changes to its PyTorch code. The Lieber Institute for Brain Development is using TorchNeuron to explore AI replacements for molecular dynamics simulations in drug discovery. Amazon also lists robotics, video, world models, self-driving cars, speech, LLMs and finance among the fields where startups and enterprises are building on it. No customer is quoted, and no performance or cost results from them are given.

Beyond TorchNeuron itself, Amazon says it is contributing changes upstream to PyTorch to remove code that ties the framework to particular chips. Hugging Face has enabled native Trainium support in its tools, according to Amazon, so no Amazon-specific code is required there. The company has also made a major contribution to the vLLM inference engine and open-sourced its full set of performance kernels as a starting point for custom work. Amazon says it tests TorchNeuron against hundreds of thousands of models on every build.

The source does not say how much TorchNeuron costs beyond the chips (the code is Apache 2.0), which AWS regions offer Trainium2 and Trainium3 capacity, or whether every PyTorch operator is covered. Those details should become clearer once the code is public.

Engineering time is the cost Amazon is attacking

The pitch is economic. Amazon argues that for startups the biggest expense is often engineer-weeks spent adapting software to new hardware, and that being able to move between chips without penalty forces chipmakers to compete for developers. It also points to large fleets that could shift capacity between inference by day and training off-peak, using the same hardware for both.

The hardware is the other half of that equation. According to figures that Introl attributes to Amazon, a Trainium3 chip delivers 2.52 petaflops of FP8 compute with 144 GB of HBM3e memory, and a 144-chip UltraServer reaches 362 petaflops. Introl also reports Amazon’s claim of 4.4 times the performance of the previous generation at the UltraServer level. Those are vendor numbers, and the page does not state the workloads or conditions behind them.

That context shows why software friction matters. A chip with those specifications only lowers a customer’s bill if the code runs on it without months of revalidation. Amazon says its custom silicon business has passed a $25 billion annual revenue run rate, so the company has a clear interest in making that adoption path as short as possible.

The 85% of AI researchers and developers who, per Amazon, write in PyTorch are the audience. What the announcement does not establish is how close real-world results come to the demonstration. Claims of identical learning behavior and lower cost still need independent benchmarks, and the open-source release on October 20 is the first point at which outside developers can check for themselves.