Amazon Unveils TorchNeuron Open-Source Toolkit for PyTorch on AWS Trainium Chips

Key summary

Amazon has announced TorchNeuron, an open-source toolkit enabling developers to run PyTorch workloads on its AWS Trainium chips with minimal code changes. The toolkit supports the latest Trainium2 and Trainium3 accelerators, which handle both training and inference tasks. Trainium3 features eight processing cores, each containing four modules optimized for different subsets of data used in AI model processing, and includes 154 gigabytes of on-board model memory. AWS custom servers can host up to 144 Trainium3 accelerators per machine, achieving a maximum throughput of 362 trillion operations per second when processing MXFP8 data formats widely used by AI models. The toolkit also supports open-source tools like TorchTitan for distributed training across accelerators using Fully Sharded Data Parallelism (FSDP). Early adopters include Cartesia AI and the Lieber Institute for Brain Development. TorchNeuron is set to be released as open source on Tuesday, the first day of the PyTorch Conference in San Jose, California.