PyTorch is the one framework that most AI developers use to build their models. It’s a Python library of operations performed on tensors, higher-dimensional analogues of a matrix, or 2-D grid of numbers. Almost all of the crucial features of contemporary AI models, such as input data, model weights, and attention mechanisms, are represented by tensors, and PyTorch drastically simplifies writing code to manipulate them.
This week, we announced the availability of TorchNeuron, an open-source backend that translates PyTorch code into operations executable on AWS Trainium systems. PyTorch code that already runs on other accelerators will run on Trainium with minimal or no code changes.
TorchNeuron bridges PyTorch and Trainium, and both sides gain from it. PyTorch developers can now run their code on Trainium, which was built from the ground up for large-tensor workloads. They can also tap the AWS Neuron SDK, which includes a profiler that traces each instruction with cycle-level accuracy and the Neuron Kernel Interface (NKI, pronounced "Nikki"), which enables control almost at the assembly level.
TorchNeuron includes several optimizations that enable PyTorch code to run efficiently on Trainium, including caching at multiple levels and an optional “adaptive eager” mode that doesn’t execute individual operations one by one, as is typical in plain eager mode, but instead groups and compiles operations for more efficient execution, approaching fully compiled performance.
Much of the innovation behind TorchNeuron, however, is the way its performance was evaluated. AI agents searched through hundreds of thousands of open-source models and categorized them according to their architectural characteristics. This allowed the Annapurna team around Khaled ElGalaind and Sarang Ozarde to identify 841 representative models covering the full spectrum of model architectures. From those models, they also distilled tens of thousands of practically important tensor dimensions, or “shapes”, which were then used to tune and benchmark TorchNeuron’s performance. The result is that PyTorch programmers can trust TorchNeuron to execute their code accurately, efficiently, and most importantly, right out of the box.
Runtime optimizations
Adaptive eager mode
Standard PyTorch has two native execution modes. One is called “eager”, meaning that PyTorch operations are executed immediately, in program order.
The other mode is compiled mode, where the programmer selects blocks of code that should be compiled in the traditional sense, meaning that a compiler builds a control-flow graph of the tensor operations and optimizes it. There are many optimization techniques in place today, such as operator fusion to reduce memory traffic and increase locality. Developers often use eager mode when they’re debugging models and compiled mode once they’re confident their code is correct and want it to execute as efficiently as possible.
TorchNeuron adds something in-between, an opt-in setting called adaptive eager mode, which offers some of the advantages of compilation but operates on smaller subgraphs. This mode approximates both eager mode’s responsiveness and compiled mode’s performance, but without the compile time overhead.
At run time, a Python interpreter running on a CPU accumulates instructions it gets from PyTorch ’s operator dispatch and builds small graphs. The graphs are sent to TorchNeuron, which compiles them, caches them, and puts them in a queue for execution. Multiple such graph collections and compilations are happening in parallel, further minimizing the compile time overhead.
A key innovation here is that the accumulated graphs, their individual operations, and their input tensor shapes correlate directly with compile time. Larger graphs on larger tensors take longer to compile. Hence, a smart, adaptive algorithm dynamically selects the right subgraphs to maximize the overlap of compile time with the execution of already compiled snippets, effectively hiding all compile time for a great variety of models.
PyTorch developers running their code in eager mode are accustomed to being able to correlate the operations executing on hardware with individual lines of code; that’s why eager mode is so useful for debugging. We didn’t want to disrupt this experience, so TorchNeuron’s adaptive-eager mode maintains accurate source correlation.
Caching
When a compiled graph comes out of the queue, TorchNeuron converts it into an executable file called a NEFF, for Neuron executable file format. A single operation in eager mode can yield a single NEFF, and the same mechanism works for full or partial graphs. In this sense, the Neuron software stack uses ahead-of-time (AOT) compilation as the default.
Converting operations to NEFFs is expensive, so TorchNeuron tries not to do it unnecessarily. Every NEFF it produces is cached at several levels. An in-memory cache keeps recently generated NEFFs, and a persistent cache on disk means NEFFs can be reused for fast restarts and future work. The AOT-by-default approach makes this straightforward.
An Amazon EC2 UltraServer is a collection of up to four individual Graviton servers with, based on configuration, somewhere between 16 and 144 Trainium accelerator chips. For very large models, even an UltraServer is not big enough to hold the full model, and models are sharded across many servers with thousands of accelerator chips. Since they all run identical parts of the same model, they all need the same NEFFs.
The TorchNeuron runtime coordinates those accelerators so that they don’t have to compile identical code; instead, one system compiles operations and graphs, and the others reuse the compiled code. A shared compilation service pools the compiler to work for all the accelerators on a server, keeping the host CPUs from being overwhelmed. In tests with a large language model spread across accelerators, these optimizations cut the time to the first training step by about a third.
Evaluation
We built AI agents that conducted a broad sweep of Hugging Face repositories, ruled out the models whose broken configurations won't run on any accelerator, and surfaced hundreds of thousands with working, executable code.
Evaluating every revision of TorchNeuron’s code on all those models would have been immensely time consuming. However, and besides, many of them perform similar tasks or process similar types of data.
So members of our team built another agent that classified models according to a defined set of architectural properties. For example, a large language model has three main components: the front end is typically an encoder, specialized for different data types, such as text, audio, images, video, or some combination thereof. The trunk houses the main Transformer blocks, the attention mechanisms and feed-forward layers, but the attention mechanisms might be self-attention, cross-attention, or both, and they come in flavors such as flash attention, ring attention, and delta attention. Finally, a backend converts the model’s inner representations into a usable output form, which might be any type of data, produced through any number of strategies.
Across all models, LLMs and otherwise, we found that just 841 representative models covered all the design characteristics that TorchNeuron was likely to confront in the wild.
But fully evaluating new code revisions on even 841 models can still take too long for rapid regression detection, so with every evaluation run, we kept a record of the different-sized tensor operations that Trainium was asked to process. By the end, we had a “shape registry” cataloguing tens of thousands of operations and their tensor dimensions that TorchNeuron is likely to encounter in practice.
From the registry, we generated unit tests that could test and validate a new software version in dramatically faster fashion. Any revision that passed the fast unit tests would then be tested against the 841 representative models. And if it passed those, we tested it again on the complete distribution of models. This staging greatly improved developer velocity, as problems could be found way earlier in the development cycle.
All this testing and the wide coverage of models give us confidence that TorchNeuron will execute PyTorch models accurately and efficiently on Trainium, giving our customers more options as they confront an increasingly crowded AI landscape.
