Making Deep Learning Go Brrrr: A First-Principles Guide to GPU Performance
Optimizing the performance of a deep learning model often feels like alchemy. Developers frequently resort to a "grab-bag of tricks"—switching PyTorch versions, setting gradients to None, or using in-place operations—without a clear understanding of why these changes work. However, by reasoning from first principles, we can move away from guesswork and systematically identify the actual bottlenecks hindering performance.
To optimize a system, you first need to know which "regime" you are in. Just as training loss versus validation loss tells you whether you are overfitting or underfitting, analyzing your system's resource utilization tells you where your time is actually being spent. In deep learning, efficiency is generally split into three components: Compute, Memory, and Overhead.
The Three Pillars of Performance
1. Compute (The Factory)
Compute refers to the time the GPU spends performing actual floating-point operations (FLOPS). The goal of most optimization is to maximize the time spent in the compute-bound regime. You pay for the massive TFLOPs of a modern GPU (like the A100's 312 TeraFLOPS), and you want to actually utilize them.
It is important to note that these peak numbers usually refer to specialized hardware like Tensor Cores, which are designed specifically for matrix multiplication. If your operations aren't matrix multiplications, you will operate at a fraction of that peak performance. However, in most deep learning models (like BERT), non-matmul operations (layer norm, activations) make up a tiny fraction of total FLOPS, meaning the inefficiency of non-matmul ops isn't usually the primary bottleneck for compute—but it is for memory.
2. Memory Bandwidth (The Warehouse)
Memory bandwidth is the cost of moving data from one place to another—specifically from the GPU's DRAM (the "warehouse") to the compute units/SRAM (the "factory").
Many operations are memory-bound, meaning the GPU spends more time shipping data than actually computing. A simple unary operation like torch.cos is a prime example: the GPU reads the data from DRAM, performs a tiny calculation, and writes it back. The computation is so fast that the GPU spends nearly all its time waiting for the memory transfer.
The Power of Operator Fusion
To combat memory bottlenecks, we use operator fusion. Instead of writing the result of every single operation back to global memory only to read it again for the next step, fusion combines multiple operations into a single GPU kernel.
For example, x.cos().cos() normally requires four global memory accesses (two reads, two writes). With fusion, it requires only two (one read, one write). This is why complex activation functions like GELU often cost the same as simple ones like ReLU; the bottleneck is the memory access, not the number of mathematical operations.
3. Overhead (The Manager)
Overhead is everything that isn't compute or memory transfer. This includes the Python interpreter, the PyTorch framework dispatch logic, and the time it takes to launch CUDA kernels.
Modern GPUs are so fast that Python becomes a massive bottleneck. In the time it takes Python to perform a single FLOP, an A100 could have processed millions. PyTorch mitigates this by executing kernels asynchronously; the CPU "runs ahead" of the GPU, queuing up work so the GPU never sits idle.
However, if your tensors are too small, the GPU finishes the work faster than the CPU can queue the next task. In this regime, your GPU becomes an "expensive paperweight."
Identifying Your Bottleneck
Knowing which regime you are in determines your solution. If you double your batch size and the runtime barely increases, you are likely overhead-bound. If you increase the complexity of your operations but the runtime stays flat, you are memory-bandwidth bound.
| Performance Regime | Plausible Solutions |
|---|---|
| Overhead-Bound | Tracing (jit.trace, FX), CUDA Graphs, or moving to a JIT compiler like TorchDynamo |
| Bandwidth-Bound | Operator Fusion (Triton, NVFuser, XLA) |
| Compute-Bound | Utilizing Tensor Cores, upgrading hardware |
Synthesis and Critical Perspectives
While the first-principles approach provides a clear mental model, the practical application can be messy. As noted in community discussions, performance is rarely portable. A model exported to ONNX may behave differently depending on whether it is run via ONNX Runtime or TensorRT, and results can vary based on target hardware and memory tuning.
Furthermore, the "bitter lesson" of scaling suggests that while human ingenuity in operator fusion is valuable, the long-term trend favors massive increases in raw TFLOPs and bandwidth. NVIDIA's ability to maintain exponential growth in both compute and interconnects ensures that the hardware continues to push the boundaries of what is possible, even as the software layers struggle to keep up.
Ultimately, the goal is to increase compute intensity—the ratio of compute to memory access. By reducing overhead and fusing operators, we clear the path for the GPU to do what it does best: perform massive matrix multiplications at peak speed.