Overview
- Amazon announced TorchNeuron on Monday, Oct. 19, and said it will publish the software under an Apache 2.0 license at the PyTorch Conference on Oct. 20.
- TorchNeuron lets developers run the same PyTorch training and inference code on AWS Trainium2 and Trainium3 with a single setting instead of rewriting model code for new hardware.
- The tool plugs into PyTorch’s device registration and works with distributed training tools such as TorchTitan and Fully Sharded Data Parallelism so models can scale across Trainium instances.
- Amazon has integrated Trainium support into community tooling like Hugging Face and vLLM, open-sourced performance kernels, and tests TorchNeuron against hundreds of thousands of models each build to catch issues early.
- More than 30 early customers are using TorchNeuron in robotics, voice, world models and drug discovery, a move that aims to lower costs and challenge Nvidia’s long-standing software lead tied to CUDA.