Amazon adds native PyTorch support for Trainium

TorchNeuron is an open-source native backend that uses PyTorch’s official extension point for custom hardware, allowing standard PyTorch APIs to operate on Trainium’s neuron device.[1] Amazon says developers can move code onto the device using `.to('neuron')`, and that an existing PyTorch project needs only a configuration change rather than a rewrite.[2][1] That promise targets a key friction in trying another chip: GeekWire reports that adapting code built around Nvidia’s CUDA platform to other hardware can take months of engineering work.[3] The announcement does not establish that every PyTorch workload will run unchanged or perform equally well.

The backend supports PyTorch eager execution for research and debugging, `torch.compile` for production performance, and the Neuron Kernel Interface for custom kernels.[1] It also integrates with TorchTitan, PyTorch’s distributed-training framework, and supports scaling training across Trainium instances with Fully Sharded Data Parallelism.[2] Amazon says developers can run code on Trainium using `.to('neuron')`.[1] Adaptive eager mode is a separate experimental feature; Amazon says it aims to reduce launch costs and create more opportunities to fuse operations without requiring code changes.[1] Teams should distinguish that experimental feature from the listed execution and scaling support.

TorchNeuron was scheduled to become a public beta on October 20, with Amazon saying it would release the tool under an open-source license at PyTorch Conference North America on that date.[1][2] GeekWire reported that the tool had been in private beta since late last year.[3] Amazon identifies early adopters across areas including robotics, video, self-driving cars, speech, drug discovery, large language models and finance.[2] Amazon says Cartesia ran its standard inference pipeline on Trainium with minimal changes to its PyTorch code, while the Lieber Institute for Brain Development is using TorchNeuron to explore strategies on Trainium.[2]

GeekWire frames TorchNeuron as part of technology companies’ efforts to reduce Nvidia’s hold on AI computing, even as those companies remain among Nvidia’s biggest customers.[3] The report says Nvidia’s software ecosystem, including CUDA, is one reason for its AI-chip dominance.[3] Other moves cited by GeekWire include Google’s goal of moving existing PyTorch code to TPUs with minimal changes, and Microsoft’s plan to run most of its AI services on Maia chips rather than Nvidia hardware.[3] Amazon says 85% of AI researchers and developers write code in PyTorch, and says extending PyTorch work to new chips previously required months of rewriting and revalidation.[2]