Amazon launches TorchNeuron toolkit for Trainium PyTorch workloads

Amazon has introduced TorchNeuron, an open-source toolkit that lets developers run PyTorch workloads on AWS Trainium chips with minimal code changes. Developers need only configure a setting to direct workloads to the correct chip. TorchNeuron supports Trainium2 and Trainium3, includes pre-packaged NKI operators for model optimization, and works with tools like TorchTitan for distributed training. Early adopters include Cartesia AI and the Lieber Institute for Brain Development. Amazon plans to open-source TorchNeuron at the PyTorch Conference in San Jose.