Case Study: Making a Transformer Fast

The Method, on a Real Model

Nothing new here — that is the point. Six sections built a method and a toolbox; now run the whole loop on ch. 11’s GPT and take it down a waterfall, one rung per technique, each measured and attributed.

Subject: d2l.GPT, width 512 (~19M params). Why not 256? At 256, bf16 goes backwards — matmuls too small for the tensor cores. The width choice is itself a measurement.

Profile Your Experiment First

Two traps corrupted an early draft by tens of percent:

• a ragged final batch put a torch.compile retrace inside the timing window → keep shapes constant
• the profiler’s instrumentation outlives its cell → TEARDOWN_CUPTI=1, or every later timing pays a launch tax

The metric is end-to-end tokens/s — DataLoader and H2D included.

Rung 0: Baseline, Profiled

Classify before you fix.

opt = torch.optim.AdamW(model.parameters(), lr=1e-3)

def step_eager(X, Y):
    opt.zero_grad(set_to_none=True)
    loss = F.cross_entropy(model(X).reshape(-1, vocab_size), Y.reshape(-1))
    loss.backward(); opt.step()

for _ in range(3):
    step_eager(*next(stream))
with torch.profiler.profile(activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA], acc_events=True) as prof:
    for _ in range(5):
        step_eager(*next(stream))
print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=6))
tput0 = throughput(step_eager)
print(f'R0 eager: {tput0:.0f} tokens/s')

------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ...
Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg # of Calls
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ...
aten::mm 9.84% 13.373ms 14.07% 19.133ms 41.146us 67.959ms 53.15% 67.959ms 146.148us 465
autograd::engine::evaluate_function: MmBackward0 1.60% 2.176ms 15.00% 20.388ms 131.537us 0.000us 0.00% 45.438ms 293.149us 155
MmBackward0 1.44% 1.963ms 13.40% 18.212ms 117.498us 0.000us 0.00% 45.438ms 293.149us 155
...
aten::matmul 0.53% 715.844us 5.41% 7.357ms 47.462us 0.000us 0.00% 22.521ms 145.295us 155
aten::mul 6.15% 8.366ms 10.25% 13.933ms 21.111us 16.700ms 13.06% 16.700ms 25.303us 660
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ...
Self CPU time total: 135.957ms
Self CUDA time total: 127.873ms
R0 eager: 262957 tokens/s

Matmuls ≈ half the device time; the rest is a fusible elementwise tail, with dispatch busy most of the step. That predicts compile pays, then precision.

The Rungs

• R1 compile — fuses the tail: ~1.3× (and asserts compiled ≡ eager first)
• R2 bf16 — tensor cores, matmuls wide enough: ~1.4×
• R3 batch-up — climb the roofline: ~1.3× (fp32-512 control: bf16 bought headroom, not admission)
• R4 checkpoint — negative for speed (−~10%), but cuts peak memory ~3× (unneeded here)

big = d2l.TimeMachine(batch_size=512, num_steps=128, tokenization='char')
big_stream = batches(big)

def throughput_big(step_fn, warmup=60, timed=50):
    for _ in range(warmup):
        step_fn(*next(big_stream))
    torch.cuda.synchronize(); t0 = time.perf_counter(); n = 0
    for _ in range(timed):
        X, Y = next(big_stream); step_fn(X, Y); n += X.numel()
    torch.cuda.synchronize()
    return n / (time.perf_counter() - t0)

torch.cuda.reset_peak_memory_stats()      # control: batch 512 in fp32
for _ in range(3):
    step_eager(*next(big_stream))
mem_fp32 = torch.cuda.max_memory_allocated() / 2**30
torch.cuda.reset_peak_memory_stats()
tput3 = throughput_big(step_bf16)
print(f'R3 +batch-up (512): {tput3:.0f} tokens/s ({tput3 / tput2:.2f}x), '
      f'peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB '
      f'(fp32 control: {mem_fp32:.1f} GiB)')

R3 +batch-up (512): 622604 tokens/s (1.33x), peak 8.6 GiB (fp32 control: 16.6 GiB)

The Waterfall

rungs = ['R0\neager', 'R1\ncompile', 'R2\n+bf16', 'R3\n+batch', 'R4\n+ckpt']
tputs = [tput0, tput1, tput2, tput3, tput4]
print(f'cumulative, R0 -> R3: {tputs[3] / tputs[0]:.2f}x')
d2l.plt.figure(figsize=(6, 3.5))
bars = d2l.plt.bar(rungs, [t / 1e3 for t in tputs],
                   color=['#7f7f7f', '#1f77b4', '#1f77b4', '#2ca02c', '#d62728'])
d2l.plt.ylabel('throughput (k tokens/s)')
d2l.plt.title('Making a Transformer fast: cumulative rungs')
for b, t in zip(bars, tputs):
    d2l.plt.text(b.get_x() + b.get_width() / 2, t / 1e3,
                 f'{t / tputs[0]:.1f}x', ha='center', va='bottom')
d2l.plt.show()

cumulative, R0 -> R3: 2.37x

Cumulative — each bar inherits every choice to its left. No dominant rung; three modest wins multiply to ~2.4×. Checkpointing is red: a technique that helped a different model hurts this one. A 300-step run confirms the fast configuration still learns.

Predict, Then Measure: Data Parallel

Transformer params ∝ compute ⇒ ~76 MB of gradients per step with little compute to hide them. The cost model (§13.5) gives a no-overlap floor; DDP’s bucketing buys back some overlap.

n_bytes = 4 * sum(p.numel() for p in model.parameters())
beta = 4.5e9   # NCCL allreduce, effective bytes/device/s on this box (13.5)
t_comm, t_cmp = 2 * n_bytes / beta, 64 * 128 / tput0
print(f'per step: t_comm ~ {1e3 * t_comm:.0f} ms vs '
      f't_compute ~ {1e3 * t_cmp:.0f} ms')
torch.cuda.empty_cache()               # hand the ranks the parent's cache
for k in (k for k in (2, 4) if k <= d2l.num_gpus()):
    floor = k * 64 * 128 / (t_cmp + t_comm)
    meas = ddp_gpt_tokens(k)
    print(f'R5 k={k}: no-overlap floor {floor / 1e3:.0f}k tokens/s, '
          f'measured {meas / 1e3:.0f}k ({meas / tput0:.2f}x of one GPU)')

per step: t_comm ~ 34 ms vs t_compute ~ 31 ms
R5 k=2: no-overlap floor 253k tokens/s, measured 297k (1.13x of one GPU)
R5 k=4: no-overlap floor 506k tokens/s, measured 455k (1.73x of one GPU)

Measured lands at the floor — just above it at k=2 (overlap), just below at k=4 (the staged fabric sags) — priced before the launch. NVLink changes the constant, not the method.

The Lore, and the Ladder Beyond

modded-nanoGPT’s speedrun = this chapter’s contents stacked to the ceiling: compiled kernels, FlashAttention, a better optimizer (Muon), fp8. Each record a new rung.

measure → classify → fix → re-measure. Two budgets, three regimes, one loop — now shown on a real model, seven sections deep.