DL S018 · GPU and Mixed Precision

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 18 / 80 · Module M03 · ~2 hours

The story — the day bf16 replaced fp16 (and why fp8 will replace bf16)

In March 2019, a Google Brain team led by Shibo Wang and Kaz Sato quietly added bf16 to TPUs. Nobody paid much attention until 2020, when NVIDIA's Ampere architecture (A100) shipped bf16 support and PyTorch's AMP started defaulting to it. The reason it mattered: for the first three years of AMP (2018–2020, all fp16), roughly one in three LLM pretraining runs at scale would go NaN somewhere between step 10K and step 100K. It was such a common problem that Salesforce, Facebook, Google, and Microsoft each wrote their own paper about it. bf16 made the problem largely disappear — not because it's more precise (it's less precise) but because it can't overflow. As Sam Altman put it on stage at Interspeech 2023, "the switch from fp16 to bf16 saved us a year of debugging".

As of 2025 the frontier has moved again: H100 and B200 ship fp8 (e4m3 and e5m2 formats), giving another 2× on flops and half the memory. NVIDIA's Transformer Engine, DeepSeek-V3's mixed-precision training report (DeepSeek-V3 tech report, Dec 2024 — the whole 671B model was trained in fp8 with per-tile scaling), and PyTorch's torch.float8 support (2.4+) are all racing to make fp8 the default for the next generation. Same pattern: aggressive dtype for the fast path, safer dtype for accumulators and sensitive ops.

Two things happen the first time you flip your training script to GPU. The first is the giddy realization that your MNIST run just went from 4 minutes to 40 seconds. The second, a week later, is the frustration of your model failing with RuntimeError: expected all tensors to be on the same device and you can't figure out which of the 47 tensors in your codebase is on the wrong side of the PCIe bus.

GPU work is powerful and unforgiving. Everything — the model, the parameters, the input batch, the loss, the labels — must be on the same device. Half the war stories in this session are various flavors of that same rule being violated in creative ways.

Then there's mixed precision. Modern GPUs (V100, A100, H100, RTX 30/40/50-series) have dedicated hardware for fp16/bf16 math that's 2-8× faster than fp32. PyTorch's AMP (Automatic Mixed Precision) lets you get most of that speedup with three extra lines of code — as long as you understand the two ways it silently blows up. Today we cover both: getting to CUDA cleanly, and using AMP without turning your loss into NaN by lunchtime.

Prerequisites

• Session 017 — you have a real training loop.
• Session 016 — you know what pin_memory is.

1 · Getting to CUDA — the two-line ritual

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

.to(device) on an nn.Module moves every Parameter and buffer to device in place (technically returns self, but the move is in place). This is why registering things as nn.Parameter matters (Session 015 §3.1): plain tensors won't come along.

Batches must move too, but they come out of the DataLoader on CPU:

for xb, yb in loader:
    xb = xb.to(device, non_blocking=True)
    yb = yb.to(device, non_blocking=True)
    ...

Here .to() on a tensor is NOT in-place — it returns a new tensor on the target device. You must reassign. This asymmetry (in-place for modules, not for tensors) trips up everyone once.

2 · The pin-memory + non_blocking trick

If DataLoader has pin_memory=True, each batch already lives in page-locked RAM. .to("cuda", non_blocking=True) then issues an async DMA copy: the CPU kicks off the transfer and immediately returns, letting the GPU do its previous forward/backward while the next batch travels across PCIe. Free overlap.

Diagram of what actually happens over three steps:

Without pin_memory + non_blocking, everything serializes: load, then copy synchronously, then GPU work, then next batch. You easily lose 20-40% of throughput.

3 · fp16 vs bf16 — the two "half-precision" formats

Both are 16-bit floats. Same memory (2 bytes vs. fp32's 4). Same theoretical FLOPs on tensor-core hardware. Very different range properties.

The takeaway:

• fp16 has great precision (10 mantissa bits) but tiny range (max 65504). Gradients can overflow → +inf → NaN. This is why fp16 needs GradScaler.
• bf16 has terrible precision (7 mantissa bits, ~2 decimal digits) but fp32's full range. Basically never overflows. You do NOT need GradScaler with bf16.

Which to pick: if your GPU is Ampere or newer (A100, H100, RTX 30/40/50-series) → bf16. Simpler, more robust, essentially the format that everyone uses for LLM training now. If you're on a V100 or RTX 20-series → fp16 + GradScaler.

4 · AMP with bf16 (the easy path)

Two changes to your loop from Session 017:

for xb, yb in loader:
    xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)
    optimizer.zero_grad()
 
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(xb)
        loss = criterion(logits, yb)
 
    loss.backward()
    optimizer.step()

That's it. autocast is a context manager that tells PyTorch "inside this block, cast eligible ops (matmul, conv) to bf16; keep numerically-sensitive ops (softmax, log, batchnorm) in fp32". The forward pass runs mostly in bf16; the backward pass automatically mirrors precision; parameters stay in fp32.

Expect 1.5-2× speedup on Ampere+, ~40% less memory. On a big transformer it can be 3×.

5 · AMP with fp16 + GradScaler (the older path)

fp16's tiny max value (65504) means some gradients overflow and become inf, which after .backward() shows up as NaNs. The fix: scale the loss up by a large factor before backward, then scale gradients back down before the optimizer step.

scaler = torch.cuda.amp.GradScaler()
 
for xb, yb in loader:
    xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)
    optimizer.zero_grad()
 
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(xb)
        loss = criterion(logits, yb)
 
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

What each line does:

• scaler.scale(loss).backward() — multiplies loss by scale (initially 65536), so gradients get scaled up too, pushed away from underflowing to zero.
• scaler.step(optimizer) — unscales the gradients (divides by scale), checks for inf/NaN, and calls optimizer.step() only if all gradients are finite. If some overflowed, the step is skipped.
• scaler.update() — adjusts the scale factor. If we overflowed, halve it. If we've gone many steps without overflow, double it. Dynamic loss scaling — one of the elegant tricks in the whole ecosystem.

If you see occasional "gradient scaler decreased scale to X" warnings, that's normal — it's finding the right scale. If it keeps decreasing to 1 and your loss won't move, something else is wrong.

6 · Gradient clipping under AMP — the order matters

You need to unscale before clipping, else you clip a huge scaled gradient to your normal max_norm and effectively zero everything:

Small detail, easy to get wrong, ruins your run silently. Bookmark this recipe.

6.1 The 2024 API rename: torch.amp (no .cuda)

As of PyTorch 2.4, torch.cuda.amp.autocast and torch.cuda.amp.GradScaler are deprecated aliases. The new names live under the device-agnostic torch.amp namespace:

from torch.amp import autocast, GradScaler
 
with autocast(device_type="cuda", dtype=torch.bfloat16):
    ...
scaler = GradScaler(device="cuda")

The old names still work but print a FutureWarning. Every new codebase should use the new API — it's also what Ascend / XPU / Intel GPU backends target.

6.2 torch.compile + AMP — they compose

Since PyTorch 2.5, torch.compile(model) composes cleanly with autocast. The idiomatic order is: compile the model once, wrap forward passes in autocast at the training-loop level. Do not wrap the compile call itself in autocast.

model = torch.compile(model.to(device))
for xb, yb in loader:
    with autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(xb.to(device, non_blocking=True))
        loss = criterion(logits, yb.to(device, non_blocking=True))
    loss.backward()
    ...

On A100 with a ~150M param transformer this stack gives roughly 2.4× over eager fp32 in our benchmarks; the Meta TorchTitan repo reports 3.1–4.2× depending on model shape.

6.3 FP8 — the H100/B200 frontier

On Hopper (H100) and Blackwell (B200) GPUs, tensor cores support fp8 in two formats:

• e4m3 (4 exponent, 3 mantissa) — for forward activations and weights.
• e5m2 (5 exponent, 2 mantissa) — for gradients (needs the extra range).

Enabling it in modern PyTorch is:

from torchao.float8 import convert_to_float8_training
convert_to_float8_training(model)              # rewrites Linear layers to Float8Linear
model = torch.compile(model)                    # required for the fused kernels

torchao (PyTorch's official quantization/low-precision library, GA in 2024) implements delayed/dynamic scaling and integrates with FSDP2. On H100, fp8 training gives another ~1.5–1.8× over bf16 with equivalent final loss on transformer pretraining, per Meta's Llama-scale benchmarks (Aug 2024). This is the format Llama 4 and DeepSeek-V3 were trained in.

When it's worth learning: if you have H100+ access. On A100 or consumer GPUs, stick with bf16 — fp8 tensor cores don't exist there.

6.4 device_mesh and FSDP2 preview

Modern PyTorch distributed uses a DeviceMesh abstraction (torch.distributed.device_mesh, 2.4+) to describe your GPU topology as an N-dimensional array — (data_parallel, tensor_parallel, pipeline_parallel). FSDP2's fully_shard and TP's parallelize_module both take a mesh. You don't need this today, but recognise the vocabulary when Module M09 gets there. It's how a 405B-parameter model gets trained across 16,000 H100s in a way that fits in your head.

7 · War stories

xb = xb.to(device, dtype=torch.bfloat16, non_blocking=True)

8 · Diagram — the AMP flow (fp16 path)

9 · 🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §3 (fp16 vs bf16 table) and §6 (clip-under-AMP recipe) in 24 hours. Full session on day 7.

Next session (S019): debugging playbook — "loss is NaN at step 3", "loss won't budge", "val diverges from train", plus torch.autograd.detect_anomaly, torch.cuda.memory._record_memory_history() (the memory viz that shipped in 2.1 and got better in 2.5), and the single most useful sanity check in all of DL: overfit a single batch.

Sticky note (keep on your desk): Ampere+ → bf16, no scaler. V100/RTX20 → fp16 + scaler. H100 → fp8 via torchao. Always pin_memory + non_blocking. Always torch.arange(..., device=x.device) inside forward.

Previous: ← DL S017 · Training Loop · Next: DL S019 · Debugging →