Horace He on Building ML Systems for Trillion-Trillion Operations

I work on the PyTorch compilers team at Meta. Today I want to explain how we build machine-learning systems for workloads involving trillions of trillions of floating-point operations, from a single GPU to very large distributed runs.

The infrastructure buildout is extraordinary: headlines describe enormous GPU clusters and new power arrangements, while leading AI companies raise billions of dollars. Much of that investment ultimately supports enormous numbers of floating-point operations, particularly matrix multiplications repeated over millions of iterations.

A chart from Epoch AI illustrates the growth in training compute. I estimate that leading models are approaching around 10^26 floating-point operations. Behind those numbers sit many layers of software, from a model API through PyTorch, Triton, CUDA and compiler toolchains.

Infrastructure work can be rewarding because a small improvement in a widely used layer benefits many people. But like a language model behind a simple interface, the infrastructure beneath it can be a vast system we do not fully understand, built by people we may never have met.

ML systems differ from many other systems because the model's logic is often simple, while expectations for its performance are extremely high. The Bitter Lesson argues that simple methods that scale with compute have repeatedly outperformed approaches built around more specialized cleverness.

Andrej Karpathy's llama2.c demonstrates the simplicity: a Llama 2 implementation can run in under a thousand lines of C without dependencies, including its loops and matrix multiplies. It is not fast in that form, but its compactness shows what the huge computation is fundamentally doing.

A common metric is model FLOP utilization: the fraction of the hardware's theoretical floating-point peak achieved by the model. Reaching 100% on a CPU would mean all cores issuing maximum-width SIMD instructions all the time. Typical CPU programs are nowhere near that, yet large ML training often reaches roughly 50%.

The field has also consolidated. Transformers have become dominant across language, vision and audio, and a relatively small number of organizations train frontier models. That changes the relationship between infrastructure providers and the users running the largest jobs.

There are two ways to improve a system: optimize the machinery for everyone, or give users a programming model and tools that let them express and tune their own work. The history of ML frameworks shows how important that choice can be.

Early frameworks such as Caffe used declarative descriptions, including Protobuf files, to specify networks. Writing those descriptions by hand was restrictive, so graph-builder APIs such as TensorFlow 1 let programmers construct the graph through a domain-specific language.

Graph builders still left a gap between what users wrote and what actually executed on the GPU. PyTorch became successful around 2016–17 with an imperative, eager model: call an operation, and conceptually it runs. That model is much easier to understand than a function passing through an intermediate graph and a sequence of transformations.

I sometimes joke that, given one day to train a model, Python would get you a faster result than C++. The reason is not that Python itself is a faster language, but that its simple, growable execution model allowed an extensive ecosystem of ML infrastructure to be built around it.

PyTorch occupied a similar design point: its simple execution model did not automatically optimize everything, but it made it easy for others to build frameworks on top. Early on, that simplicity did not necessarily cost much performance compared with graph-based competitors.

One reason was that matrix multiplications consumed over 90% of runtime and were already handled by highly optimized vendor libraries such as cuBLAS and cuDNN. There was relatively little other work for a framework compiler to optimize.

Another reason was asynchronous execution. The CPU schedules GPU work onto a queue while the GPU processes it. I picture Gromit laying track ahead of a moving train: as long as Python can schedule work faster than the GPU consumes it, its scheduling overhead can effectively be hidden.

That balance changed when NVIDIA introduced Tensor Cores in 2017. These GPU units execute matrix multiplication directly, and the gap between matmul throughput and other GPU operations became very large. Accelerating matmuls leaves the rest of the program responsible for a larger fraction of runtime.

ML compilers emerged to optimize that remaining work. The user still writes code with an imperative programming model, but the implementation captures a graph and optimizes it. torch.compile, for example, intercepts execution at the Python bytecode interpreter level, somewhat like a just-in-time compiler for PyTorch programs.

Historically, graph-builder systems exposed a graph both to the programmer and at execution; early PyTorch exposed and executed eager operations. Modern frameworks commonly preserve an eager user interface while capturing a graph internally, although PyTorch, JAX, TinyGrad and MLX use different capture approaches.

I divide GPU runtime into three broad categories: compute, memory movement within the GPU, and overhead such as time the GPU spends idle. That framing helps identify what kind of optimization can actually improve a given workload.

On modern hardware, matrix multiplication can be dramatically faster than other arithmetic. On the GPU in my example, I compare about 67 TFLOPS of FP32 compute with roughly 1,000 TFLOPS of matrix-multiply-oriented TF32 compute. Even fully occupied non-matmul arithmetic can therefore look poor against the matmul peak.

An unexpected detail is that the values inside matrices can affect performance even when the memory access pattern and control flow are identical. Benchmarks using zeros, Gaussian samples and uniformly distributed data can produce different results.

The explanation is power. Transistor switching consumes dynamic power, and different bit patterns cause different amounts of switching. Hardware can throttle when it reaches power or thermal limits, so a matrix full of zeros may run faster than data that causes more switching.

That can make benchmarks misleading. Someone described a model producing NaNs and suddenly appearing faster: once many tensor values became NaN, the hardware switched differently and throttled less. Better throughput in that situation was a symptom of a broken computation, not a genuine optimization.

Memory movement is the second major cost. In one BERT-like workload, tensor contractions account for 99.8% of floating-point operations but only 61% of runtime. Operations doing comparatively little arithmetic can still occupy a substantial share of time.

Even when data already resides in GPU VRAM, it must move to the compute units and back. Think of a factory with limited space and a larger warehouse farther away. Moving supplies between them consumes bandwidth and time even though it does no useful arithmetic itself.

Suppose we apply an addition, an activation and another elementwise operation separately. In eager execution, intermediate data may travel from memory to compute and back for each operation. Sending the result back to the warehouse only to fetch it immediately for the next step is wasteful.

Operator fusion combines operations in one GPU kernel: bring the data to compute once, perform several operations, and write the result back. I consider this the most important optimization in an ML compiler. It requires looking across operations, which the simplest eager execution model does not do.

Fusion also involves tradeoffs between saving and recomputing values. Recomputing from data already nearby can avoid memory accesses. It may reduce peak memory use and improve runtime, despite performing additional arithmetic.

Backpropagation makes this especially relevant. Training first runs a forward pass through the layers and then a backward pass in reverse order. Intermediates, or activations, may need to be retained between the two passes; their long lifetimes can cause out-of-memory errors.

The third category is overhead. If the CPU cannot enqueue work quickly enough, the GPU waits idle, like a train catching up with the person laying its tracks. CUDA Graphs and lower-overhead generated wrappers are among the ways to reduce that scheduling overhead.

With only a few dominant architectures and billions spent on a training run, one might ask why a compiler is needed at all. Why not assign specialists to optimize a model by hand? In practice some teams do avoid compilers, and their reasons reveal a real usability problem.

Compilers offer leverage: implement an optimization once and many users can benefit. But imagine a library that does not always work, does not document when it works, and may change its behavior after an update. That is how compiler optimization can feel when 'works' means achieving the performance you need.

An influential discussion of ISPC puts it this way: auto-vectorization is not a programming model. If a vectorizer can fail, programmers who depend on SIMD performance must understand why it failed and adjust their code or coax the compiler into applying the optimization.

An explicit programming model gives a user something more dependable. A SIMD intrinsic specifies an operation that maps to a SIMD instruction; users can reason from that contract rather than reverse-engineering an optimizer. An optimization that always works can become part of the programming model.

A compiler normally hides its complex implementation behind a simple interface. When it misses an optimization, that implementation becomes visible to users struggling to understand why code was not inlined or vectorized. In ML, numerical precision makes the situation even more delicate.

Hardware keeps pushing toward lower-precision arithmetic, from 16-bit operations through 8-bit and, in newer hardware, 4-bit formats. Small numerical differences can have large effects, and an optimization that looks better in isolation can disrupt an algorithm's assumptions.

I encountered NaNs in a FlashAttention implementation involving fused multiply-add, or FMA. FMA computes a multiply and addition with one final rounding, often giving a more accurate result than separately rounded operations. Yet that change can break code relying on two computations producing exactly the same rounded value.

The calculation subtracts a maximum before exponentiation to avoid large exponents. If the maximum was computed from a rounded product but a later expression uses FMA on only one branch, the expressions can differ in a way the algorithm did not expect. In this case FMA produced NaNs while disabling it avoided them.

Compilers also struggle with algebraic rewrites. FlashAttention relies on an online softmax reformulation that removes a global synchronization. The mathematics is manageable for a person, but a compiler generally cannot be relied upon to invent the rewrite from a sequence of ordinary operations.

An attention computation can be written as a matrix multiply, softmax and another matrix multiply. A compiler might pattern-match that sequence and replace it with a fused kernel. But a user who writes a custom softmax can accidentally defeat the pattern and suddenly see much greater memory use and slower execution.

One response is to expose a single fused attention operator. That gives users a predictable fast operation, but a monolithic API cannot easily express every new idea. Researchers continue to introduce variants such as sliding-window, ALiBi, paged and neighborhood attention.

A fused operator accumulates flags and special cases as variants are added, and existing quirks cannot simply be removed without breaking users. Even with many options, someone will still need an attention variant the operator cannot express.

The choice need not be between an unpredictable optimizer and an inflexible monolith. A more useful programming model can combine a complicated hand-written FlashAttention kernel with simple user modifications that a compiler can generate mechanically.

FlexAttention is an example. It uses torch.compile but gives users a dependable fused attention kernel and its memory properties. Users can experiment within the model without needing to understand every implementation detail of the underlying compiler.

For a sliding-window causal mask, a user can check that the query position is at or after the key/value position and that their distance is within the window, then combine those conditions. The kernel machinery handles the implementation while the user supplies the variation.

Before this approach, support across attention masks, biases, training and inference was a patchwork of manually implemented kernels. A flexible abstraction fills in more combinations. One user even expressed a mask for molecular graphs of different sizes, a use I had not anticipated when designing the API.

I think about programming models through a mathematical analogy: instead of hitting a walnut harder, change how you approach the problem until it becomes easier to open. A different abstraction can turn a fragile optimization into a property that users can rely on.

Consider parallelism in CUDA. Asking when the compiler will automatically make a CUDA program parallel misses the point: parallel execution is built into its programming model. The program may still be wrong or deadlock, but parallelism is not a speculative auto-vectorization decision.

Scaling beyond one GPU introduces another set of problems. A traditional distributed service often scales many small, interchangeable queries and may accept extra hardware for availability. A large training run resembles one huge job with frequent global synchronization and very tight performance constraints.

One way to organize parallelism is to imagine a cube of work: separate tasks along a batch dimension, operations within a task, and time. Data, tensor and pipeline parallelism divide the work along different dimensions of that picture.

Data parallelism assigns different data or tasks to different GPUs, but gradients must be synchronized after each step. We cannot simply increase the batch indefinitely to avoid communication: excessively large batches can prevent training from converging as desired. Fully sharded data parallelism also avoids naively replicating all parameters.

Tensor parallelism splits the same task across GPUs. Communication and computation are often harder to overlap because the devices are cooperating on that task, yet the performance penalty of idle hardware makes overlap valuable.

Asynchronous tensor parallelism can create a small pipeline inside an operation. It divides communication and matrix multiplication into pieces, allowing computation on one piece while another is being transferred. That creates overlap even within a tensor-parallel operation.

Pipeline parallelism instead assigns successive portions of a task to different GPUs at different times. Global synchronization can leave gaps in the pipeline, and backpropagation adds scheduling choices because forward and backward microbatches must both be accommodated.

A training pipeline may have points where it can choose to run a forward or a backward microbatch. Different schedules exploit that choice to reduce idle periods. The Llama 3 training diagram combines data, tensor and pipeline parallelism, along with context parallelism.

Compilers have historically struggled to keep up with distributed optimization. Given a fixed set of choices, an automatic model may select a good configuration. But much of the progress comes from inventing a new dimension of parallelism, expanding the search space rather than simply searching the old one.

Fault tolerance becomes another constraint at large scale. With globally synchronized work across thousands of GPUs, a hardware, network or user-related failure can interrupt the entire run. A failure rate that seems manageable for a smaller fleet becomes a frequent obstacle as the fleet grows.

The Llama 3 fault-tolerance figures illustrate the issue: at around 16,000 GPUs an error occurs roughly every 1.8 hours in his example, while at 131,000 GPUs it is about every 15 minutes. At sufficient scale, failures can occur before a training step completes.

ML has become an exceptionally important computational workload, but its social and technical characteristics mean conventional systems and compiler approaches do not always transfer directly. Workloads keep evolving; a compiler team cannot optimize in isolation for years and ignore those changes.

The most interesting question is therefore not only which optimizations to implement, but what programming models give people reliable tools to express and develop their own optimizations. Those models are what can help users build the next generation of very large training systems.