Conference Presentation, Lecture, Keynote
Horace He: Building Machine Learning Systems for a Trillion Trillion Floating Point Operations
- The speaker, Horace He, is a compiler engineer on the PyTorch team at Meta, responsible for developing
torch.compileandFlexAttention. - AI infrastructure build-out has accelerated dramatically, with every major AI company now securing its own nuclear power plants or massive GPU clusters (e.g., XAI's 300k GPU cluster).
- Startup fundraising benchmarks have shifted; raising $1 billion in 2024 is now a prerequisite to simply enter the market, compared to being a "unicorn" ($1B valuation) in 2016.
- Leading-edge models are currently trained using approximately $10^{26}$ floating-point operations (100 trillion trillion FLOPs), or roughly 1 trillion teraflops.
- Modern ML stacks are multi-layered abstractions (LM API, PyTorch, Triton, CUDA, etc.) that often obscure the underlying complexity, similar to the "benevolent dictator" metaphor of infrastructure work.
- Despite simple model logic (e.g., Llama-2 can be implemented in ~973 lines of C), performance expectations are extremely high, with deep learning systems targeting ~50% Model FLOP Utilization (MFU).
- The field has consolidated significantly over the last decade; Transformers are now the dominant architecture for vision, language, and audio, while SOTA training is dominated by a few companies.
- Performance optimization strategies have shifted from pure "optimization" (improving the compiler) to defining robust "programming models" that allow users to specify optimizations directly.
- Early frameworks like Caffe (2010-2012) used declarative Protobuf files, while TensorFlow 1 used graph-building DSLs that were difficult for users to reason about.
- PyTorch's success in 2016-2017 was driven by "eager execution," a simple imperative model where code runs line-by-line, matching Python's native execution behavior.
- PyTorch's early performance was competitive because 90% of runtime was spent in matrix multiplications (matmuls) provided by highly optimized vendor libraries like cuBLAS/cuDNN.
- The introduction of NVIDIA Tensor Cores (2017) created a 10x speed gap for matmuls versus other operations, necessitating ML compilers to optimize non-matmul paths.
- Modern ML compilers capture eager Python code into graphs to perform optimizations while maintaining an imperative programming model for the user.
- GPU runtime is fundamentally split between compute (floating-point operations) and memory movement (data shuffling), with data movement often dominating runtime despite contributing <1% to total FLOPs.
- Memory bandwidth costs arise because GPU compute units are in SRAM (limited capacity) while data resides in VRAM (large capacity), requiring constant data transfer.
- Operator fusion is the most critical compiler optimization, combining multiple operations (e.g., add, relu, sin) into a single kernel to minimize data movement between SRAM and VRAM.
- Recomputation vs. reuse trade-offs are essential in ML due to the unique program structure of forward and backward passes, which require saving large intermediates (activations).
- Overhead in ML systems occurs when the CPU (scheduling) cannot keep the GPU busy, often visualized as the GPU waiting for the CPU to lay down "train tracks."
- Compilers often struggle with numerical precision; for example, Fused Multiply-Add (FMA) operations can introduce NaNs if they change rounding behavior between different branches of a calculation.
- Flash Attention optimization relies on algebraic rewrites (online softmax) that compilers generally fail to discover automatically, requiring explicit API exposure.
- Monolithic fused operators (e.g., early Flash Attention) struggle to evolve as new variants (sliding window, ALiBi, Page Attention) emerge weekly.
- The "FlexAttention" API addresses this by allowing users to define attention masks via Python code, which
torch.compileguarantees will fuse into an efficient kernel. - Distributed ML training differs from traditional distributed systems: it involves a single massive query (training job) rather than many small QPS queries, making performance critical and fault tolerance difficult.
- Parallelism strategies include Data Parallelism (splitting batch dimension), Tensor Parallelism (splitting task dimensions), and Pipeline Parallelism (splitting time dimensions).
- Pipeline parallelism in ML is complicated by the alternating forward/backward pass pattern, requiring complex scheduling to overlap communication and computation.
- Large-scale training (e.g., Llama-3) combines all parallelism types plus context parallelism, but compilers struggle because innovation often comes from expanding the search space rather than optimizing within it.
- Fault tolerance becomes a critical failure point at scale; failure rates increase inversely with the number of GPUs (e.g., one failure every 15 minutes on 131k GPUs).
- The speaker concludes that the primary challenge in ML systems is not just building better optimizations, but designing robust programming models that allow users to express complex optimizations reliably.