AI NewsInfrastructureAnnouncement
AWS releases TorchNeuron, a PyTorch backend for Trainium chips
AWS says a one-line change (.to('neuron')) is enough to move a PyTorch training or inference script to Trainium chips, with no lazy-tensor rewrite.
Image: Hugging Face
Why it mattersA team pricing an AI product against Nvidia hardware can now test the same PyTorch script on Amazon's chips without a port, so the comparison stops being a research project.
A team that wanted to try Amazon's Trainium chip for a PyTorch model used to have to rewrite the training loop for lazy tensors, static shapes, and a graph that recompiled whenever it broke. From 20 October the rewrite is no longer required.
AWS announced TorchNeuron, a native PyTorch backend for its Trainium chips, open sourced under Apache 2.0 at github.com/aws-neuron/torch-neuronx. The code goes live as a public beta on 20 October, timed to the PyTorch Conference in San Jose. The authors are AWS Annapurna Labs engineers; the announcement is co-signed by Jana van Greunen, Director of PyTorch Engineering at Meta.
What changes in your code
AWS says a model moves to Trainium with one line, .to('neuron'), and the rest of the training script stays as written. The standard PyTorch primitives work on the new device: torch.nn, torch.autograd, torch.autocast, torch.distributed, and torch.profiler. Both execution modes work natively: eager for research and debugging, and torch.compile(model, backend="neuron") for production. The company also lists an experimental "adaptive eager mode" that it says fuses operations to approach compiled performance without the trace step.
For distributed work, AWS says DDP, FSDP2, tensor parallelism and pipeline parallelism all run on the standard PyTorch APIs, with collectives sent over Trainium's dedicated collective cores inside a node and over Elastic Fabric Adapter between nodes.
The number the vendor is leading with
Hugging Face published the companion piece the same day. Hugging Face reports that 89 percent of the top 100 most-used text-generation model architectures run on Trainium through TorchNeuron out of the box, against 46 percent through its previous library, Optimum Neuron v0.4.5. Hugging Face recommends that Neuron users work with Transformers, Diffusers, Accelerate, PEFT and TRL directly from now on, instead of the Optimum Neuron wrapper that used to hide the lazy-tensor shape of PyTorch XLA.
TechCrunch, writing the same day, carries a vendor-sourced figure from Colin Brace, VP of Technology Software at AWS Annapurna Labs: Deepgram, an AI voice provider, "moved to Trainium using native PyTorch support and they're processing 70 percent more audio per dollar at production-grade accuracy." Treat this as Deepgram's own figure routed through AWS; TechCrunch did not verify it, and the piece carries no third-party benchmark.
The piece Amazon is adding upstream
Amazon is also contributing changes to the PyTorch project itself to remove code it says ties the framework to specific chips. The company's stated goal is to make PyTorch hardware-agnostic for every developer who uses it. TorchNeuron is built on PyTorch's own PrivateUse1 device backend and torch.accelerator APIs, extension points the PyTorch community created for exactly this.
Amazon says 85 percent of AI apps are written with PyTorch today. CUDA, Nvidia's own layer, has been the default for most of them. Alternative chips from AMD and Apple already run PyTorch with a configuration change, so the mechanism is not new; what is new is that an AWS Trainium instance now joins that list without a research project between the laptop and the compile step.
For teams reviewing cloud bills on a workload that already runs on Nvidia GPUs, the practical move is to clone the training script, change .to('cuda') to .to('neuron'), and read what the same script costs per hour on each.
Source
This item was written by an AI system from the linked source. Reveneau is responsible for what it publishes.
Get AI News in your inbox
New developer tools, model and agent releases, and how teams are actually using them to release software. Short, and only when there is something worth reading.
