Profiling, MFU and compilation

Before this lesson: lesson 2 (roofline), lesson 3 (launch cost and fusion), lesson 6 (activation recomputation), lesson 7 (the cost of a collective) and lesson 8 (overlap of communication with compute). The compute budget $C \approx 6\Psi D$ of scaling laws (LLMs track, in preparation) is the same count as in section 3. Time: about 75 minutes, plus the notebook.

---

Why this matters #

"Our run reaches 30% MFU. Is that good, and what would you change?" This is a common question in interviews and in real work. For scale: PaLM 540B reached 46.2% MFU on 6144 TPU v4 chips, GPT-3 reached 21.3% on V100 GPUs (both numbers from the PaLM paper), and Llama 3 405B reached 38–43% on 8,192 to 16,384 H100 GPUs. Three tools answer the question:

• a FLOP count says what the hardware could do;
• a profile says where the time goes;
• a compiler removes part of the waste, but only in some regimes.

The same arithmetic also plans a run: how many GPU-hours 1 trillion tokens cost, and whether a cheaper or a faster setup is better.

1. Predict first #

2. The picture #

Where the time of a step goes. The CPU runs the Python program. Every PyTorch operation on a GPU tensor puts a kernel into a queue on the GPU (a CUDA stream) and returns at once; the GPU runs the queued kernels in order. NCCL communication runs on its own stream, so it can overlap with compute. So the step time splits into three parts:

• $t_\text{compute}$: compute kernels run. Each one is limited by its FLOPs or its bytes (the roofline of lesson 2).
• $t_\text{comm, exposed}$: communication runs and no compute runs next to it.
• $t_\text{idle}$: the GPU has nothing to do. The CPU is too slow to launch kernels, or it waits for the GPU (a synchronization), or it waits for data.

A timeline (from Nsight Systems or the PyTorch profiler) shows these parts directly: one row for each CPU thread and one row for each GPU stream.

Interactive figure: loading…

Three regimes. Horace He's way to think about any piece of work on a GPU: its time is set by one of three things, and each needs a different fix.

Regime	Time is about	Typical case
overhead-bound	CPU time per kernel × number of kernels	many tiny ops, small batch
memory-bound	bytes / bandwidth	elementwise ops, norms, small-batch matmuls
compute-bound	FLOPs / peak	large matmuls, attention at long sequences

A profile tells you which regime you are in. A compiler can help the memory-bound part (fusion) and the overhead-bound part (fewer kernels, graph capture), but not the compute-bound part.

Interactive figure: loading…

One number for the whole step. MFU (model FLOPs utilization) is the fraction of the peak FLOP/s that goes into the FLOPs the model needs. It is the first number to compute and the number to compare between runs.

3. Definitions and main results #

Notation. $\Psi$ parameters used in matmuls; $L$ layers; $h$ the attention width $a \cdot d_h$ (heads × head dimension; it equals the hidden size in most models); $s$ the sequence length; $\tau$ tokens per second over all GPUs; $n_\text{gpu}$ GPUs; $P$ the dense peak FLOP/s of one GPU in the format used (H100 SXM bf16: 989 TFLOP/s). Papers write $6N + 12LHQT$ for the formula below; we write $6\Psi + 12Lhs$. An input embedding is a table lookup and costs no FLOPs; papers usually include it in the parameter count, and for large models the difference is small.

Theorem (training FLOPs per token, PaLM convention) #

A dense decoder trained on sequences of $s$ tokens needs, per token,

FLOPs. With full activation recomputation, the hardware executes

Proof.

1. Forward, weights. In a matmul $Y = XW$, each weight takes part in one multiply and one add per token: $2\Psi$ FLOPs per token.
2. Backward, weights. For each forward matmul $Y = XW$, the backward computes $dX = dY\,W^\top$ and $dW = X^\top dY$: two matmuls of the same size, $4\Psi$.
3. Forward, attention. Per layer and per token: the scores of the token's query against $s$ keys are $s$ dot products of length $h$ (over all heads), $2sh$ FLOPs. The weighted sum of $s$ value vectors costs another $2sh$. Over $L$ layers: $4Lhs$.
4. Backward, attention. The two forward products need four backward products of the same size: the gradients of the attention weights and of the values (from the output gradient), and of the queries and of the keys (from the score gradient). Twice the forward: $8Lhs$.
5. Sum: $2\Psi + 4\Psi + 4Lhs + 8Lhs = 6\Psi + 12Lhs$.
6. Full recomputation runs the forward once more during the backward: $+(2\Psi + 4Lhs)$, so $8\Psi + 16Lhs$. Every term grows by $4/3$. $\blacksquare$

The attention term counts the full $s \times s$ score matrix. A causal kernel skips about half of it, but the PaLM convention counts all of it; say which convention you use. Norms, softmax and activations are not counted: they are a few FLOPs per element, small next to the matmuls. For a 7B model with $s = 4096$, the attention term is 13.7% of $F_\text{model}$.

MFU and HFU #

HFU (hardware FLOPs utilization) counts every FLOP the hardware runs, recomputation included. MFU is the number to compare (PaLM paper, §4.1). Its numerator is the measured tokens per second times a constant of the architecture, and its denominator is the published peak. So it does not depend on implementation choices. HFU rewards repeated work: full recomputation raises HFU by 4/3, while the run becomes slower.

Rules:

• Use the dense peak (989 TFLOP/s for H100 bf16), not the "with sparsity" number (1,979). For an fp8 run, say which peak you use.
• Count tokens and peak over the same set of GPUs.
• Measure steady-state steps only (not the first steps, which include compilation and warm-up).
• FlashAttention's backward recomputes the scores (lesson 4). These FLOPs belong to HFU, not to MFU.

From MFU to training time and cost #

Read the MFU formula the other way. To train on $D$ tokens:

Without the attention term, $F_\text{model} D = 6\Psi D$: the compute budget $C \approx 6\Psi D$ of the scaling laws. Note that the GPU-hours do not depend on $n_\text{gpu}$ at a fixed MFU: more GPUs finish sooner but cost the same. In practice MFU usually falls a little as $n_\text{gpu}$ grows (more communication, smaller batch per GPU), so more GPUs cost a little more.

Example. A 7B model (Llama-2 shape: $\Psi = 6.74 \times 10^9$, $L = 32$, $h = 4096$, $s = 4096$) on $D = 10^{12}$ tokens, on 256 H100s at 40% MFU.

1. $F_\text{model} = 6 \cdot 6.74 \times 10^9 + 12 \cdot 32 \cdot 4096 \cdot 4096 = 4.04 \times 10^{10} + 6.44 \times 10^9 = 4.69 \times 10^{10}$ FLOP/token.
2. Total model FLOPs: $4.69 \times 10^{10} \times 10^{12} = 4.69 \times 10^{22}$.
3. Delivered FLOP/s: $256 \times 989 \times 10^{12} \times 0.40 = 1.01 \times 10^{17}$.
4. Time: $4.69 \times 10^{22} / 1.01 \times 10^{17} = 4.63 \times 10^5$ s $= 5.36$ days (4.62 days without the attention term).
5. GPU-hours: $256 \times 4.63 \times 10^5 / 3600 = 32{,}900$.

Interactive figure: loading…

Proposition (launch-bound criterion) #

The CPU needs $t_\text{cpu}$ to issue one kernel (Python, the framework's dispatcher, the launch call), and the kernel runs $t_k$ on the GPU. In a long stream of such kernels, the GPU is busy at most a fraction

of the time.

Proof.

1. Kernels enter the queue at most once every $t_\text{cpu}$.
2. The GPU finishes at most one every $t_k$.
3. If $t_\text{cpu} > t_k$, the queue empties: after each kernel, the GPU waits $t_\text{cpu} - t_k$ for the next one. In each period of $t_\text{cpu}$ it is busy $t_k$: busy fraction $t_k / t_\text{cpu}$.
4. If $t_\text{cpu} \le t_k$, the CPU runs ahead, the queue never empties, and the GPU is always busy: fraction 1. $\blacksquare$

Eager PyTorch costs a few µs of CPU time per operation (about 3.6 µs, lesson 3). With kernels of 2 µs the GPU is busy at most $2/3.6 = 56\%$ of the time. The fixes follow from the formula: longer kernels (a larger batch), fewer kernels (fusion), or no CPU work per kernel (CUDA graphs).

Host-device synchronization. Some calls make the CPU wait until the GPU has finished all queued work: .item(), .cpu(), print(tensor), if tensor: and ops whose output size depends on the data, such as nonzero. Such a sync breaks step 4 of the proof: the queue is empty when the CPU continues, so the GPU idles until the CPU has launched new work.

Goodput: failures, checkpoints and restarts

MFU measures the steps that run. A long run also loses time to failures. Goodput (or effective training time) is the useful training time divided by the elapsed time. Llama 3 405B had 466 job interruptions in a 54-day period of pre-training on up to 16,384 H100s (419 unexpected, mostly hardware), and still reached more than 90% effective training time.

A simple model shows the trade-off. A checkpoint costs $\delta$ seconds and is written every $\tau_c$ seconds; failures come on average every $M$ seconds; after a failure, the work since the last checkpoint (on average $\tau_c / 2$) is lost. The lost fraction is about

$$\frac{\delta}{\tau_c} + \frac{\tau_c}{2M},$$

which is smallest at $\tau_c^* = \sqrt{2\delta M}$ (Young's approximation, 1974). Example: 419 failures in 54 days is one every $M \approx 3.1$ h. With $\delta = 60$ s, $\tau_c^* = \sqrt{2 \cdot 60 \cdot 11{,}135} \approx 1{,}156$ s, about 19 minutes, and the loss is about $5.2\% + 5.2\% = 10.4\%$, plus the restart time. This is why large runs use fast asynchronous checkpoints (small $\delta$) and fast failure detection.

A plan for a performance analysis #

1. Estimate. Compute the hardware FLOPs per GPU per step from the theorem, and their time at peak. Real kernels stay below the peak, so also measure your GPU's rate on one large bf16 matmul and use it as a realistic ceiling. Estimate the communication bytes per GPU from the parallel scheme (lesson 8) and their time with the $\alpha$–$\beta$ model of lesson 7: can they overlap with compute? Estimate the bytes per second the input pipeline must deliver.
2. Measure the step. Use steady-state steps, after warm-up and compilation. Time them with synchronization (torch.cuda.synchronize() or CUDA events), over several steps. Compute MFU.
3. Profile a few steps. Split each step into compute kernels (matmul, attention, elementwise and norms), exposed communication and idle time. Name the cause of each idle gap.
4. Compare each part with its estimate, and fix the largest gap first. Change one thing, then measure again.
5. Stop when the step is close to the estimate, or when the next fix costs more than it saves.

Reading a timeline #

Nsight Systems (nsys) records a timeline of the whole system: CPU threads with their CUDA API calls (cudaLaunchKernel, cudaMemcpyAsync, cudaStreamSynchronize), named NVTX ranges, and one row of kernels per GPU stream. NCCL kernels have names that start with nccl. Profile a few steady steps only, for example:

nsys profile -t cuda,nvtx,osrt --capture-range=cudaProfilerApi -o step python train.py

with torch.cuda.cudart().cudaProfilerStart() and cudaProfilerStop() around the steps, and torch.cuda.nvtx.range_push("forward") / range_pop() to name the phases. The PyTorch profiler (torch.profiler.profile with ProfilerActivity.CPU and ProfilerActivity.CUDA, schedule(wait, warmup, active), and record_function for named ranges) gives a table of operators and a trace that links each CPU launch to its GPU kernel.

What you see	Cause	How to confirm	Fix
short gaps between short kernels; the CPU thread is busy all the time	launch-bound: CPU time per kernel > kernel time	gaps line up with cudaLaunchKernel calls; kernels of a few µs	larger batch per GPU, fusion (torch.compile), CUDA graphs, a fused or foreach optimizer
NCCL kernels with an empty compute stream next to them, often at the end of the backward	exposed communication	step time minus compute time; NCCL stream busy, compute stream empty	overlap: DDP bucket size, FSDP backward prefetch and unit size; HSDP (most traffic stays inside a node); gradient accumulation (fewer syncs per token)
GPU idle at the start of each step; the main thread waits in the data loader	input stall	idle time matches the loader wait; the host-to-device copy comes just before the first kernel	more workers, pin_memory=True with non_blocking=True copies, prefetching, pre-processed data in chunk-aligned files
cudaStreamSynchronize or a device-to-host copy on the CPU row, then a GPU gap	host-device sync (.item(), .cpu(), print(tensor), if tensor:)	the call stack of the sync (profiler with_stack=True)	log every $k$ steps, keep metrics on the GPU, copy asynchronously, no Python control flow on GPU values
long elementwise, norm or softmax kernels	memory-bound kernels	their intensity (lesson 2); Nsight Compute's "Speed Of Light" section	fusion, fused kernels (FlashAttention, fused optimizer)
one rank slower than the others; the others wait inside collectives	straggler	per-rank step times; long NCCL kernels on the fast ranks	find the slow node (thermal limits, memory errors, a bad link) and replace it

Compilation #

Eager mode runs each operation at once: Python → dispatcher → one kernel. It is flexible and easy to debug, but nothing is optimized across operations. Each operation reads its inputs from HBM, writes its outputs back, and costs one launch.

torch.compile (Ansel et al., 2024) has three stages.

1. TorchDynamo captures graphs. It hooks into CPython's frame evaluation (PEP 523), runs the function's bytecode symbolically, and records the PyTorch operations into an FX graph. It installs guards: cheap checks of everything the graph was specialized on (shapes, dtypes, devices, requires_grad, Python constants, module attributes such as training). On the next call, the guards run; if they all pass, the compiled code runs.
   • A graph break happens when Dynamo meets code it cannot trace: a Python branch on a tensor value, .item(), print, an unsupported library call. Dynamo ends the graph, runs that code in normal Python, and starts a new graph after it. The function still works, but there are fewer chances to fuse and more overhead.
   • A recompilation happens when a guard fails. With the default settings, a new size first triggers a recompile with that dimension marked dynamic (symbolic), and later sizes reuse it. A new dtype or device, a changed Python constant, or model.eval() on a module that reads training (dropout, batch norm) triggers another compile. After 8 compiles of one function (torch._dynamo.config.recompile_limit), Dynamo stops compiling it and runs it eagerly.
2. AOTAutograd builds the backward. It traces forward and backward together, ahead of time, on fake tensors (shapes without data). A partitioner splits the joint graph into a forward and a backward graph and decides which activations to save and which cheap ones to recompute.
3. TorchInductor (the default backend) generates code. It lowers the graphs to a loop-level representation, fuses pointwise operations and reductions, and writes Triton kernels on GPUs (C++ with OpenMP on CPUs). Matmuls and convolutions call vendor libraries; mode="max-autotune" also tries tuned Triton templates. mode="reduce-overhead" replays the compiled code with CUDA graphs.

Tools: torch._dynamo.explain(fn)(*args) counts graphs and graph breaks; TORCH_LOGS="graph_breaks,recompiles" prints the reasons; fullgraph=True turns a graph break into an error. Over more than 180 models on an A100, the PyTorch 2 paper reports geometric-mean speed-ups of 2.27× for inference and 1.41× for training.

Interactive figure: loading…

CUDA graphs record a sequence of kernel launches once (stream capture) and replay it later with one launch, so the CPU cost per kernel disappears. The restrictions: fixed shapes and fixed memory addresses (inputs are copied into static buffers before each replay), no CPU synchronization and no CPU-side decisions inside the graph, and memory held by the graph's pool. Serving engines capture one graph for each of a list of decode batch sizes and pad a batch up to the next captured size (lesson 12).

JAX jit traces the function once per input signature with abstract values (shape and dtype, no data) into a jaxpr, lowers it to StableHLO, and the XLA compiler compiles the whole program (fusion, layouts, scheduling). The result is cached. A new shape, dtype or value of a static argument traces and compiles again. Python side effects (print, appending to a list) run only during tracing. A Python branch on a traced value raises an error: use jax.lax.cond, lax.while_loop, lax.scan or jnp.where, or mark the argument as static (static_argnums). All shapes are static.

Which tool helps which workload. Each regime has its own time formula, and a tool helps only if it shrinks the term that sets the time:

Workload	Time is set by	Helps	Does not help
overhead-bound: small batch, many tiny kernels (small-batch decode, a small MLP or PINN at small batch, an optimizer over many parameter tensors)	CPU time per kernel × kernels	CUDA graphs (also via mode="reduce-overhead"), JAX jit, fusion (fewer kernels), a larger batch	faster kernels
memory-bound: norms, activations, residual adds, losses, optimizer updates, physics losses on large grids	bytes / bandwidth	fusion: TorchInductor, XLA, hand-fused kernels	CUDA graphs (launches are not the limit)
compute-bound: large matmuls, FlashAttention at long sequences	FLOPs / peak	better kernels, lower precision (fp8), fewer FLOPs (a better algorithm)	compilation, mostly (autotuning can pick a slightly better kernel, but it cannot remove FLOPs)

Answers to section 1. (1) About 43% (section 4): a good number for dense training. (2) (a) speeds up much more. Its time is the memory traffic of the elementwise chain, and fusion removes most of it: a chain of $k$ ops moves $2kN$ elements unfused and $2N$ fused (lesson 3). (b) gains much less: compilation fuses the norms, activations and residual adds around the matmuls, but the matmuls and attention, which take most of the time, already run close to their roofline.

4. Worked example #

Setup. Llama-2-7B shape: $\Psi = 6.74 \times 10^9$, $L = 32$, $h = 4096$, $s = 4096$. FSDP (ZeRO-3) on 32 H100s (4 nodes × 8 GPUs), full activation checkpointing, global batch 256 sequences. Log line: 3.60 s per step.

Step 1: MFU.

1. Tokens per step: $256 \times 4096 = 1{,}048{,}576$, so $\tau = 1{,}048{,}576 / 3.60 = 291{,}271$ tokens/s (9,102 per GPU).
2. $F_\text{model} = 6\Psi + 12Lhs = 4.044 \times 10^{10} + 6.44 \times 10^9 = 4.688 \times 10^{10}$ FLOP/token (attention is 13.7% of it).
3. $\text{MFU} = \dfrac{291{,}271 \times 4.688 \times 10^{10}}{32 \times 989 \times 10^{12}} = 43.1\%$ (37.2% without the attention term).
4. $\text{HFU} = \tfrac{4}{3} \times 43.1\% = 57.5\%$.

Step 2: the estimate.

1. Per GPU and step: $1{,}048{,}576 / 32 = 32{,}768$ tokens.
2. Hardware FLOPs: $32{,}768 \times (8\Psi + 16Lhs) = 2.05 \times 10^{15}$, which take 2.07 s at peak. The model FLOPs alone would take 1.55 s.
3. Communication: FSDP all-gathers the bf16 weights in the forward, all-gathers them again in the backward, and reduce-scatters the gradients (in bf16 here). Each GPU sends about $3 \times 2\Psi \times 31/32 = 39.2$ GB per step. Even if all of it crossed the 50 GB/s network link of each GPU (an upper bound: part of every ring runs over NVLink), that is 0.78 s. It can hide behind 2 s of compute when the overlap works.

Step 3: the measured breakdown. Suppose a profile of five steady-state steps on one rank gives these numbers (an invented but typical example):

Part	Measured	Estimate
compute kernels	2.90 s	2.07 s at peak
exposed NCCL	0.40 s	0 (should overlap)
idle: data-loader wait at the start of the step	0.18 s	0
idle: .item() sync in the logging code	0.12 s	0
step	3.60 s

The compute kernels run at $2.05 \times 10^{15} / (2.90 \times 989 \times 10^{12}) = 71\%$ of peak.

Step 4: the largest gap first.

1. Exposed communication, 0.40 s. The last reduce-scatters have no compute left to hide behind, or prefetching is off. Check FSDP backward prefetch and the unit size; try HSDP (shard inside a node, replicate across nodes), so most traffic stays on NVLink.
2. Idle, 0.30 s. More loader workers with pinned memory; log the loss every 50 steps instead of calling .item() at every step.
3. Compute, 2.90 s against 2.07 s at peak. Full checkpointing costs one extra forward (4/3 of the model FLOPs). Selective recomputation (lesson 6) keeps most of the memory saving for a few percent of extra FLOPs, and torch.compile fuses the norms and activations.

If fixes 1 and 2 remove all exposed and idle time, the step takes 2.90 s and MFU rises to 53.6%. Fixing only the communication gives 3.20 s and 48.5%.

5. In code #

The first block checks the numbers of sections 3 and 4 with plain Python. The notebook for this lesson has these functions as exercises.

def train_flops_per_token(psi, L, h, s):
    """Model FLOPs per token, PaLM convention: 6 Psi + 12 L h s (h = attention width)."""
    return 6 * psi + 12 * L * h * s

def mfu(tokens_per_s, flops_per_token, n_gpu, peak):
    return tokens_per_s * flops_per_token / (n_gpu * peak)

P = 989e12                                      # H100 SXM, dense bf16

# Worked example: 7B (Llama-2 shape), 32 GPUs, 256 x 4096 tokens per 3.60 s step
F = train_flops_per_token(6.74e9, 32, 4096, 4096)
tau = 256 * 4096 / 3.60
print(f"{F:.4g} FLOP/token, {tau:,.0f} tokens/s")             # 4.688e+10 FLOP/token, 291,271 tokens/s
print(f"MFU {mfu(tau, F, 32, P):.1%}, HFU {4 / 3 * mfu(tau, F, 32, P):.1%}")   # MFU 43.1%, HFU 57.5%

# PaLM 540B on 6144 TPU v4 chips (275 TFLOP/s): attention width 48 x 256 = 12,288
F_palm = train_flops_per_token(540.35e9, 118, 48 * 256, 2048)
print(f"PaLM MFU {mfu(238.3e3, F_palm, 6144, 275e12):.1%}")   # PaLM MFU 46.2%

# Planning: 7B on 1T tokens, 256 GPUs at 40% MFU
D, n, u = 1e12, 256, 0.40
t = F * D / (n * P * u)
print(f"{t / 86400:.2f} days, {n * t / 3600:,.0f} GPU-hours")  # 5.36 days, 32,919 GPU-hours

# Launch-bound criterion: busy fraction of the GPU
busy = lambda t_k, t_cpu: t_k / max(t_k, t_cpu)
print(busy(2e-6, 3.6e-6), busy(20e-6, 3.6e-6))                 # 0.5555555555555556 1.0

PyTorch's FLOP counter checks the theorem on a small model. The two blocks below need PyTorch (they run on a CPU); they are not in the notebook.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.flop_counter import FlopCounterMode

class Layer(nn.Module):                       # matmul part of a decoder layer, no biases or norms
    def __init__(self, d, n_heads):
        super().__init__()
        self.h = n_heads
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.o = nn.Linear(d, d, bias=False)
        self.up = nn.Linear(d, 4 * d, bias=False)
        self.down = nn.Linear(4 * d, d, bias=False)

    def forward(self, x):                     # x: (batch, s, d)
        b, t, d = x.shape
        q, k, v = self.qkv(x).view(b, t, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        scores = q @ k.transpose(-1, -2) / (d // self.h) ** 0.5   # full s x s scores, as in PaLM's count
        a = scores.softmax(dim=-1) @ v
        x = x + self.o(a.transpose(1, 2).reshape(b, t, d))
        return x + self.down(F.gelu(self.up(x)))

L, d, n_heads, s, batch = 4, 256, 4, 128, 2
model = nn.Sequential(*[Layer(d, n_heads) for _ in range(L)])
psi = sum(p.numel() for p in model.parameters())             # 12 L d^2 = 3,145,728
x = torch.randn(batch, s, d, requires_grad=True)             # so the first layer's input gradient counts too
with FlopCounterMode(display=False) as counter:
    model(x).sum().backward()
per_token = counter.get_total_flops() / (batch * s)
print(psi, int(per_token), 6 * psi + 12 * L * d * s)         # 3145728 20447232 20447232
print(f"attention share {12 * L * d * s / per_token:.1%}")   # 7.7%

The count matches $6\Psi + 12Lhs$ exactly ($h = d$ here). Without requires_grad=True on the input, autograd skips $dX$ of the first matmul, and the count is $2 \cdot 256 \cdot 768 = 393{,}216$ FLOPs per token lower.

The next block counts compilations and graph breaks. A backend that only records the captured graph makes the count visible.

import torch
import torch._dynamo as dynamo

graphs = []
def counting_backend(gm, example_inputs):     # records each captured FX graph, then runs it
    graphs.append(gm)
    return gm.forward

def f(x):
    return torch.relu(x) * 2 + 1

cf = torch.compile(f, backend=counting_backend)
for n in (8, 16, 32):
    cf(torch.randn(n))
print(len(graphs))       # 2: size 8 is compiled as static; 16 fails the guard and is
                         #    recompiled with a dynamic size, which 32 then reuses
cf(torch.randn(32, dtype=torch.float64))
print(len(graphs))       # 3: a new dtype fails the guards again

def g(x):
    y = x.sin()
    if y.sum() > 0:      # Python branch on a tensor value
        return y.cos()
    return y.exp()

def h(x):
    y = x.sin()
    s = y.sum().item()   # copies a value to the host as a Python float
    return y * s

for fn in (f, g, h):
    e = dynamo.explain(fn)(torch.randn(10))
    print(fn.__name__, e.graph_count, e.graph_break_count)
# f 1 0
# g 2 1
# h 2 1

These outputs were checked with PyTorch 2.15 on a CPU. Older versions may count differently.

6. Where it shows up #

• LLM pretraining reports give MFU as the main efficiency number. A drop in MFU, or one rank slower than the others, is the first alarm in a long run.
• Serving engines capture CUDA graphs for the small-batch decode step (overhead-bound and memory-bound) and fuse norms, rotary embeddings and activations (lesson 12).
• Clusters. On many nodes, profile one or two ranks with nsys, and compare the step times of all ranks to find stragglers.
• ML for science. Physics-informed losses and PDE residuals are long elementwise chains: memory-bound, so fusion helps a lot. Graph networks on meshes, or samples of different sizes, change shapes from step to step: pad to a few bucket sizes or mark the dimensions dynamic, or torch.compile keeps recompiling. Large gridded datasets (Zarr, NetCDF) often cause input stalls; chunks that match the training samples fix most of them. Several weather and climate models (GraphCast, NeuralGCM) are written in JAX with jit and static shapes.

7. Common mistakes #

• Using the sparse peak (1,979 TFLOP/s for H100 bf16) instead of the dense one (989), or a bf16 peak for an fp32 run.
• Mixing scopes: tokens/s of one GPU against the peak of all GPUs, or the other way round.
• Reporting HFU as MFU, or counting recomputation in MFU.
• Using the hidden size for the attention term when the attention width differs (PaLM: 12,288 against 18,432).
• Timing GPU code without synchronization. Launches return at once, so a CPU timer without torch.cuda.synchronize() measures the launches, not the work.
• Profiling the first steps. They include compilation, library autotuning and allocator warm-up. Skip them (schedule(wait=..., warmup=..., active=...)).
• A .item() or a printed loss at every step: one host-device sync per step.
• Expecting torch.compile to speed up a matmul-bound model, or counting the compile time of the first steps as step time.
• Shapes that change every step (sequence lengths, the last batch): recompiles, and after 8 of them an eager fallback. With CUDA graphs, a changing shape or a CPU sync inside the captured region breaks the capture.

8. Check yourself #

1. A 1.2B decoder ($L = 24$, attention width 2048, $s = 2048$) trains on 8 H100s at 400,000 tokens/s. Compute the MFU with and without the attention term.
   Answer

   $6\Psi = 7.2 \times 10^9$; $12Lhs = 12 \cdot 24 \cdot 2048 \cdot 2048 = 1.21 \times 10^9$; sum $8.41 \times 10^9$ FLOP/token. MFU $= 4 \times 10^5 \times 8.41 \times 10^9 / (8 \times 989 \times 10^{12}) = 42.5\%$; without attention 36.4%.
2. The same run uses full activation checkpointing. What is its HFU, and why do you report the MFU?
   Answer

   $F_\text{hw} = \tfrac{4}{3} F_\text{model}$, so HFU $= 56.7\%$. MFU depends only on tokens/s, the architecture and the published peak. HFU grows with recomputation, an implementation choice that makes the run slower, not faster.
3. How long does the 1.2B run of question 1 need for $D = 2.4 \times 10^{10}$ tokens (20 tokens per parameter), and how many GPU-hours is that?
   Answer

   $t = D/\tau = 2.4 \times 10^{10} / 4 \times 10^5 = 6 \times 10^4$ s $= 16.7$ h. GPU-hours: $8 \times 16.7 = 133$.
4. Batch-1 decode of a 1B model: the CPU row is full of cudaLaunchKernel calls, and the GPU runs kernels of 3–5 µs with gaps between them. Which regime is this, which tool fits, and why does torch.compile in its default mode help only partly?
   Answer

   Overhead-bound (launch-bound): the CPU time per kernel exceeds the kernel time. CUDA graphs (by `mode="reduce-overhead"` or a manual capture) remove the CPU cost per kernel. Default `torch.compile` fuses ops and so cuts the number of kernels, but each remaining kernel is still launched from the CPU, and the compiled call has its own overhead.
5. A compiled module with dropout receives inputs of shape (4, 8), then (128, 8), (256, 8) and (512, 8). Then you call model.eval() and run it again. How many compilations happen?
   Answer

   Three. (4, 8) compiles a static graph; (128, 8) fails the shape guard and recompiles with a dynamic first dimension, which (256, 8) and (512, 8) reuse; dropout reads the `training` attribute, `eval()` changes it, a guard fails, and it compiles again. (A module that never reads `training`, such as a plain `nn.Linear`, does not recompile.)
6. Inside a function under jax.jit you write if x.sum() > 0: and print(x). What happens with each, and what do you use instead?
   Answer

   The `if` raises an error: during tracing `x` is an abstract value (shape and dtype only), so it has no truth value. Use `jax.lax.cond` or `jnp.where`, or make the argument static. The `print` runs once, during tracing, and prints a tracer, not the values; use `jax.debug.print` to print at run time.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. How many training FLOPs per token does a dense decoder need (PaLM convention)?

   Show answer

2. Define MFU.

   Show answer

3. What is the difference between HFU and MFU, and which one do you report?

   Show answer

4. How long does training on $D$ tokens take at a given MFU?

   Show answer

5. State the launch-bound criterion.

   Show answer

6. Compute the GPU busy fraction for 2 µs kernels and 3.6 µs of CPU time per launch.

   Show answer

7. Name the three regimes and what fixes each.

   Show answer

8. What is a graph break in torch.compile?

   Show answer

9. What does TorchInductor generate?

   Show answer

10. What does a profiler timeline show for exposed communication?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Our run reaches 30% MFU. How would you improve it?" — Estimate the step from FLOPs and communication bytes; measure; split the profile into compute, exposed communication and idle time; fix the largest gap first; measure again.
• "The GPU is idle 20% of the time. Why?" — Launch overhead, host-device syncs, data-loader stalls. Exposed communication and straggler waits show up as NCCL kernels with no compute next to them.
• "How long will it take to train a 7B model on 1T tokens on 256 H100s?" — $6\Psi D$ plus the attention term, divided by $n_\text{gpu} P$ times a realistic MFU: about 5.4 days at 40% MFU (4.8 days at 45%, 6.1 days at 35%).
• "When does torch.compile not help?" — On matmul-bound models, and when graph breaks or recompiles split the model into small pieces.
• "CUDA graphs or JAX jit?" — CUDA graphs only remove launch overhead (static shapes and addresses); jit traces abstract values into one XLA program and fuses, with static shapes and structured control flow.

10. Practice #

The notebook for this lesson runs in Colab, with NumPy only. It has five exercises, each with a check cell and a folded solution:

1. train_flops_per_token(...) and mfu(...), checked on the 7B example and on PaLM 540B.
2. The training time and GPU-hours of a planned run, and the number of GPUs needed for a deadline.
3. A queue simulator of CPU launches and GPU kernels: the busy fraction against $t_k$ and $t_\text{cpu}$, checked against the proposition.
4. Read a small synthetic profiler trace and compute the compute, exposed-communication and idle shares of a step.
5. Fit the fixed cost per call and the bandwidth of your own CPU from NumPy timings: the overhead regime and the bandwidth regime (under 20 seconds).

11. Watch and read #

Lecture 1 How to profile CUDA kernels in PyTorchMark Saroufim · GPU MODE lecture 1 · 56 min

The PyTorch profiler, a trace in the browser, and how to see which kernels a PyTorch program runs.

torch.compile: The Missing ManualEdward Yang · PyTorch · 15 min

How to find graph breaks and recompilations, and what to do about them.

Let's reproduce GPT-2 (124M)Andrej Karpathy · 241 min

From 1:48:15, torch.compile, Python overhead and kernel fusion cut a training step from about 300 ms to 130 ms.

Performance Analysis with NVIDIA Nsight Systems Timeline | CUDA Developer ToolsNVIDIA Developer · 13 min

How to read the Nsight Systems timeline: CPU threads, CUDA API calls and GPU rows.

• Chowdhery et al. (2022), PaLM: Scaling Language Modeling with Pathways: §4.1 (why MFU and not HFU) and Appendix B (the MFU formula and its numbers).
• Ansel et al. (2024), PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation (ASPLOS): TorchDynamo, guards, graph breaks, AOTAutograd and TorchInductor.
• Llama Team (2024), The Llama 3 Herd of Models: §3.3 (MFU of 38–43% on H100s, and the interruptions of a long run).
• PyTorch, Profiler recipe: execution time, memory and traces.
• NVIDIA, Nsight Systems user guide: the section on profiling from the command line.
• JAX, Just-in-time compilation: tracing, static arguments and caching.
• Horace He, Making Deep Learning Go Brrrr From First Principles: the three regimes, with pictures.