Amazon introduces toolkit for running PyTorch workloads on AWS Trainium chips

Amazon introduces toolkit for running PyTorch workloads on AWS Trainium chips

Amazon.com Inc. today introduced TorchNeuron, an open-source toolkit that enables developers to run PyTorch workloads on its AWS Trainium chips.

PyTorch is the industry’s most popular artificial intelligence development framework. Until now, running AI projects powered by the tool on Trainium-based Amazon Web Services instances involved a significant amount of manual work. That increased development costs and thereby created an entry barrier for users.

TorchNeuron makes it possible to run PyTorch projects on Trainium chips without extensive code changes. According to Amazon, developers only have to configure a setting that informs workloads which chip they should use. TorchNeuron supports Trainium2 and Trainium3, the two latest additions to the accelerator series. The chips support both training and inference use cases. 

Trainium3, the more capable of the two processors, features eight computing cores. Each core comprises four modules that are each optimized for a different subset of the calculations with which AI models process data. The chip keeps 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 Amazon unit, a single system can perform up to 362 quadrillion calculations per second when processing data in the MXFP8 format. The format is widely used by AI models.

The Trainium series features a custom programming language called NKI. The syntax, which comes with its own compiler, 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. They also speed up related tasks such as quantization.

TorchNeuron supports not only PyTorch but also the open-source tools that software teams often use alongside the framework. The list of supported tools includes TorchTitan, a popular AI training engine. It 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 splits 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.

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. 

Amazon plans to open-source TorchNeuron Tuesday, the first day of a two-day PyTorch Conference in San Jose.

Image: Amazon