Build an LLMAdvancedLesson 467 min read

Step 10 · Making it fast

Same model, same GPUs, three times the throughput. Most training runs leave that on the table.

🧪 Lab · 20 min 📝 Lesson test

Lesson in motion

In 60 seconds

Step 10 · Making it fast

Same model, same GPUs, three times the throughput. Most training runs leave that on the table.

1/5

In simple words

Two cooks with the same kitchen: one wastes half the day walking to the fridge. Speed here is almost entirely about not walking to the fridge.

Modern GPUs can do arithmetic far faster than they can move data. Almost every optimisation below is really about moving less data.

The ranked list

Technique	Typical gain	Effort
bfloat16 mixed precision	2–3x	One line
torch.compile	1.3–2x	One line
FlashAttention	1.5–3x at long context	Free via scaled_dot_product_attention
Fused optimiser	5–10%	One argument
Gradient checkpointing	Enables larger batches	One line, costs ~30% compute
Sequence packing	10–40% on varied-length data	Moderate
Better data loading	Sometimes 2x	Pre-tokenise and memory-map
FP8 on recent hardware	1.3–1.8x	Significant care required

Why FlashAttention matters so much

A naive implementation builds the full sequence-by-sequence attention matrix in memory. At 8k context with 32 heads that is billions of numbers written to memory and read straight back.

FlashAttention never materialises that matrix. It processes attention in tiles that stay in fast on-chip memory, computing the softmax incrementally. Same mathematical result, far less memory traffic, and memory use grows linearly with sequence length instead of quadratically.

An exact algorithm, not an approximation. The output is bit-comparable to the naive version. It is purely a better memory access pattern — which is why it was adopted universally within months.

Sequence packing

If your examples vary in length and you pad to the longest, you may be spending 40% of your compute on padding. Packing concatenates short examples into full-length sequences with an attention mask that prevents them attending across the boundary.

Pack examples into full sequencespython

def pack(sequences, block_size, eos_id):
    buf, out = [], []
    for seq in sequences:
        buf.extend(seq + [eos_id])
        while len(buf) >= block_size:
            out.append(buf[:block_size])
            buf = buf[block_size:]
    return out
# NOTE: pass document boundaries to the attention mask, or examples
# will attend across each other and quietly learn nonsense.

Danger

That note is a real and common bug. Packing without a document mask lets the model attend from one example into an unrelated one. Training still "works" — the loss falls — and the model is subtly worse in ways no unit test catches.

Measure before you optimise

Tokens per second, honestlypython

import time
torch.cuda.synchronize(); t0 = time.time()
for _ in range(20):
    x, y = get_batch("train")
    with torch.autocast("cuda", dtype=torch.bfloat16):
        _, loss = model(x, y)
    loss.backward(); opt.step(); opt.zero_grad(set_to_none=True)
torch.cuda.synchronize()

tok_per_s = 20 * x.numel() / (time.time() - t0)
flops = 6 * n_params * tok_per_s
print(f"{tok_per_s:,.0f} tok/s   MFU {flops / peak_flops:.1%}")

Do this

Always compute MFU. A number like "12% utilisation" turns a vague sense that things are slow into a specific, fixable engineering problem — and it is usually the data loader.

Watch and read more

FlashAttention explainedYouTube search GPU performance for deep learningYouTube search

Lab

A 2x speedup on the same hardware, itemised.

~20 min

The problem

Baseline your training throughput. Then apply, one at a time: bf16 autocast, torch.compile, FlashAttention via SDPA, a fused optimiser, and sequence packing. Measure after each and attribute the gain.

You are done when

A table of tok/s after each change
Total speedup at least 2x
You can name which change gave the most

Hard questions

Try to answer before you reveal. If you can answer these, you understood the lesson.

Q1Packing gave you 35% on varied-length data. Name the bug it introduces and how you detect it.Reveal

Cross-contamination: without a document mask, tokens attend across the boundary into an unrelated example, so the model learns from sequences that never existed. It is silent — loss falls, nothing crashes. Detect it by constructing a batch where document A ends with a distinctive token and document B begins with a continuation of it, then checking whether the model's prediction at B's start is influenced by A. Or simply assert the attention mask is block-diagonal, which is cheaper and catches it every time.

Please sign in to continue.

Questions people ask

bf16 or fp16?

bfloat16 on anything from Ampere onward. It has the same exponent range as fp32, so it does not overflow and needs no loss scaling. fp16 has more mantissa precision but a narrow range, and it makes large-model training fragile.

Is FP8 worth it?

On hardware that supports it, yes, for large runs — but it needs per-tensor scaling and careful handling of sensitive layers. Do not start here. Get bf16 and compile working first.

Why is my GPU utilisation high but throughput low?

Utilisation percentage only says the GPU is busy, not that it is doing useful arithmetic. It can be busy moving memory. MFU is the honest metric.

Does torch.compile ever hurt?

It adds compile time on first run and recompiles on every new input shape. With highly dynamic shapes you can spend more time compiling than computing. Fixed shapes plus compile is the fast path.

Lesson test

5 questions. Get 3 right (60%) to pass and complete this lesson.