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 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.
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
TorchNeuron runs PyTorch code asynchronously. Each PyTorch operation is placed in a queue, and the user program moves on without waiting. In the background, the Neuron compiler turns each queued operation into an executable file called a NEFF, for Neuron executable file format, which then runs on the device in program order. In eager mode each operation gets its own NEFF; with torch.compile, a whole captured region becomes one NEFF.
Compiling is expansive, so TorchNeuron works hard to never do it twice. Every NEFF it produces is cached at several levels. An in-memory cache keeps recently used NEFFs instantly available, and a persistent cache on disk means a NEFF compiled once can be reused across restarts and future training runs.
In a distributed training setup, caching also leverages all processes participating in the run. Processes that run the same part of the model need the same NEFFs, so on a single machine, dozens of them may ask for identical compiled code. TorchNeuron runs a shared compilation service so that each operation is compiled only once per machine, and the other processes reuse the result, without overwhelming the host CPUs. In tests with a large language model spread across 64 processes on a single machine, the service 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.
