AWS Introduces TorchNeuron Toolkit for PyTorch Workloads on Trainium Chips
Amazon's open-source toolkit eliminates code-heavy deployment barriers for PyTorch AI models on custom Trainium accelerators, with early exploration in medical research workflows.
Amazon Web Services has introduced TorchNeuron, an open-source toolkit that enables developers to run PyTorch workloads on AWS's custom Trainium chips with minimal code modifications. The company plans to open-source the toolkit on October 22, 2024, coinciding with the first day of the two-day PyTorch Conference in San Jose.
The toolkit addresses a longstanding technical barrier: until now, running PyTorch projects on Trainium-based AWS instances required significant manual configuration work, increasing development costs and creating entry barriers for users. According to Amazon, TorchNeuron requires developers only to configure a setting that informs workloads which chip to use.
Technical Architecture and Capabilities
TorchNeuron supports both Trainium2 and Trainium3 accelerators for training and inference workloads. Trainium3, the more capable processor, features eight computing cores, with each core comprising four modules optimized for different subsets of AI model calculations. The chip maintains model data in a 154-gigabyte pool of onboard memory.
AWS runs Trainium3 in custom servers that can each accommodate up to 144 accelerators. According to the company, a single system can perform up to 362 quadrillion calculations per second when processing data in the MXFP8 format, which is widely used by AI models.
The Trainium series features NKI, a custom programming language with its own compiler that enables developers to make low-level optimizations to their AI models. Such customizations reduce AI models' hardware usage and boost their speed. TorchNeuron includes a set of pre-packaged NKI operators—software building blocks that customers can use to develop AI models. Some operators ease the task of implementing a mixture of experts architecture, while others are designed to reduce neural networks' memory usage and speed up related tasks such as quantization.
Integration with PyTorch Ecosystem
TorchNeuron supports not only PyTorch but also open-source tools that software teams often use alongside the framework. The list of supported tools includes TorchTitan, a popular AI training engine that automates some of the work involved in distributing a training run across accelerators.
TorchTitan splits workloads across chips using a method called fully sharded data parallelism. It divides not only the model being trained but also optimizer states and gradients. An optimizer is a program that hones a model's capabilities, while gradients are files that play an important role in that process. According to Amazon, TorchNeuron includes features that enable developers to further optimize the training workflow using Trainium chips.
Early Adoption
The toolkit has already been adopted by several AWS customers. Cartesia AI Inc., a venture-backed voice agent startup, is using TorchNeuron to run inference workloads on Trainium chips. The Lieber Institute for Brain Development is exploring ways of optimizing its AI-powered medical research workflows with the toolkit.
Sources
Reporting & accountability
How we maintain accuracy
Health AI Central separates reported facts, analysis, and public-source data. Readers can review our standards, request a correction, or consult the relevant research methodology before relying on a claim.