Mixed precision, gradient accumulation, and activation checkpointing

Track 03 - Training · Notebook 01 · Runtime: ≈8-12 min on T4, ≈4 min on CPU

Prerequisites: you’ve trained something with PyTorch once, even if it was MNIST. You know the difference between float32 and float16 exists but perhaps not when to use which.

What you’ll know by the end: the three tricks every LLM training run uses to fit into less memory than it “should” need - and why each of them comes with a tradeoff you can feel in the numbers.

---

The problem: training hates memory

A language model’s training memory budget is dominated by four things:

1. Parameters - the weights themselves.

2. Gradients - one per parameter, same size as the parameters.

3. Optimiser state - AdamW keeps two moments per parameter, so 2× params.

4. Activations - the intermediate tensors you store for the backward pass. The big surprise: on a transformer, activations can be larger than everything else combined.

For a 1 B-parameter model in fp32, the parameters alone are 4 GB. Add gradients and optimiser state and you’re at ≈16 GB before you’ve processed a single token. Activations for a reasonable batch and sequence length easily add another 10+ GB. This is why training is so much memory-hungrier than inference, and why single-GPU training on consumer hardware hits the wall quickly.

The three techniques in this notebook attack different parts of that budget:

• Mixed precision (bf16) halves the bytes per parameter, gradient, and activation.

• Gradient accumulation lets you simulate a big batch by running several small batches and summing their gradients, so peak activation memory stays small.

• Activation checkpointing drops most activations during the forward pass and recomputes them during backward, trading compute for memory.

Each one is a small idea with an outsized impact. We’ll ablate them in a controlled way on a mini GPT-2 and watch the numbers move.

Step 1 - Why bf16 works

Floating-point numbers are (sign, exponent, mantissa). fp32 has 8 bits of exponent and 23 bits of mantissa - 32 bits total. fp16 has 5 bits of exponent and 10 bits of mantissa. bf16 has 8 bits of exponent (the same as fp32) and 7 bits of mantissa.

Two consequences:

• bf16 has fp32’s range but much less precision. Any number fp32 can represent by magnitude, bf16 can represent too - you just lose precision within it. fp16’s smaller exponent would silently underflow tiny gradients and overflow big activations.

• Neural networks don’t need many mantissa bits. Once a weight is accurate to 1 % it’s close enough - stochastic gradient descent is already doing a noisy walk. The 7-bit mantissa of bf16 gives roughly 0.4 % precision, which works fine for every transformer trained in the last five years.

So bf16 gives you 2× memory savings and roughly 2× throughput on tensor-core hardware, at essentially zero quality cost. That’s a free lunch, which is rare in deep learning.

from llm_systems_cookbook.nb import bootstrap

import math
import time
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as ckpt

s = bootstrap("03_training_01_mixed_precision_accum_checkpointing")

Step 2 - A mini GPT-2 we can torture

The model here is deliberately modest: six transformer blocks, 256 hidden dim, 256 context. That’s small enough to run on CPU in a few minutes but big enough to make peak memory measurable on a T4.

The point isn’t to get good perplexity (we’d need far more training data for that) - it’s to have a realistic activation footprint so the ablations below show a real signal.

@dataclass(frozen=True)
class Cfg:
    vocab: int = 2048
    d_model: int = 256
    n_head: int = 8
    n_layer: int = 6
    seq: int = 256
    ff_mult: int = 4

CFG = Cfg(n_layer=6 if IS_CUDA else 3, seq=256 if IS_CUDA else 128)


class Block(nn.Module):
    def __init__(self, d_model: int, n_head: int, ff_mult: int) -> None:
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_head, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_model * ff_mult),
            nn.GELU(),
            nn.Linear(d_model * ff_mult, d_model),
        )

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        h = self.ln1(x)
        h, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
        x = x + h
        return x + self.ff(self.ln2(x))


class MiniGPT(nn.Module):
    def __init__(self, cfg: Cfg, use_checkpoint: bool = False) -> None:
        super().__init__()
        self.tok_emb = nn.Embedding(cfg.vocab, cfg.d_model)
        self.pos_emb = nn.Embedding(cfg.seq, cfg.d_model)
        self.blocks = nn.ModuleList([Block(cfg.d_model, cfg.n_head, cfg.ff_mult) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab, bias=False)
        self.use_checkpoint = use_checkpoint
        self.register_buffer(
            "mask", torch.triu(torch.ones(cfg.seq, cfg.seq), diagonal=1).bool(),
            persistent=False,
        )

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)[None]
        mask = self.mask[:T, :T]
        for block in self.blocks:
            if self.use_checkpoint:
                x = ckpt.checkpoint(block, x, mask, use_reentrant=False)
            else:
                x = block(x, mask)
        return self.head(self.ln_f(x))


set_seed(0)
model_probe = MiniGPT(CFG).to(DEVICE)
n_params = sum(p.numel() for p in model_probe.parameters())
print(f"model: {n_params:,} params  (~{n_params * 4 / 1024**2:.1f} MiB in fp32)")

Step 3 - Synthetic data (keep the spotlight on memory, not on corpora)

We train on random token sequences. The point is that each variant sees the same sequence of batches so any difference in final loss is attributable to the precision / accumulation / checkpointing choice, not to data variance.

set_seed(0)
BATCH = 16 if IS_CUDA else 4
toy = torch.randint(0, CFG.vocab, (256, CFG.seq + 1), device=DEVICE)


def iter_batches(batch_size: int, n_steps: int) -> list[torch.Tensor]:
    out: list[torch.Tensor] = []
    for i in range(n_steps):
        start = (i * batch_size) % (toy.shape[0] - batch_size)
        out.append(toy[start : start + batch_size])
    return out


def step_loss(model: nn.Module, batch: torch.Tensor) -> torch.Tensor:
    x, y = batch[:, :-1], batch[:, 1:]
    logits = model(x)
    return F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))

Step 4 - Four variants, one loop

run_variant trains a fresh model with different (precision, accumulation, checkpointing) combinations and reports peak memory, step time, and final loss. We run four configurations:

Name	What changes
fp32	Baseline. Everything in fp32.
bf16	Autocast wraps every op, activations and intermediate products use bf16. Parameters stay fp32.
bf16 + accum=4	Same precision, but we accumulate gradients over 4 micro-batches of size BATCH/4 before stepping.
bf16 + accum=4 + checkpoint	Same plus activation checkpointing - drop most activations in forward, recompute them in backward.

Gradient accumulation is the cleanest way to train with an effective big batch on small hardware: each micro-batch’s activations fit in memory, and because we scale the loss by 1/accum_steps, the total gradient is identical (up to numerical noise) to one big-batch pass.

def reset_memory() -> None:
    if IS_CUDA:
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()


def peak_mb() -> float:
    return torch.cuda.max_memory_allocated() / 1024**2 if IS_CUDA else 0.0


def run_variant(name: str, use_bf16: bool, accum_steps: int, use_checkpoint: bool,
                 n_steps: int = 4) -> dict[str, float]:
    set_seed(0)
    model = MiniGPT(CFG, use_checkpoint=use_checkpoint).to(DEVICE)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
    micro = BATCH // accum_steps
    batches = iter_batches(batch_size=micro, n_steps=n_steps * accum_steps)
    reset_memory()
    losses: list[float] = []
    t0 = time.perf_counter()
    for step in range(n_steps):
        opt.zero_grad(set_to_none=True)
        step_sum = 0.0
        for mi in range(accum_steps):
            b = batches[step * accum_steps + mi]
            ctx = torch.autocast(device_type=DEVICE.type, dtype=torch.bfloat16) if use_bf16 \
                  else torch.autocast(device_type=DEVICE.type, enabled=False)
            with ctx:
                loss = step_loss(model, b) / accum_steps
            loss.backward()
            step_sum += loss.item() * accum_steps
        opt.step()
        losses.append(step_sum / accum_steps)
    dt = time.perf_counter() - t0
    return {"name": name, "peak_mb": peak_mb(), "step_s": dt / n_steps, "final_loss": losses[-1]}


results: dict[str, dict[str, float]] = {}
for cfg_name, bf16, accum, ckp in [
    ("fp32",                   False, 1, False),
    ("bf16",                   True,  1, False),
    ("bf16+accum4",            True,  4, False),
    ("bf16+accum4+checkpoint", True,  4, True),
]:
    results[cfg_name] = run_variant(cfg_name, bf16, accum, ckp)
    r = results[cfg_name]
    print(f"{cfg_name:>24}  peak={r['peak_mb']:7.2f} MiB  step={r['step_s'] * 1000:6.1f} ms  loss={r['final_loss']:.3f}")

Reading the table

On T4 (where the peak-memory numbers are real) you should see something like:

• bf16 cuts peak memory by ≈35 % vs fp32 - exactly what the “half the bytes per activation” argument predicts.

• bf16+accum4 doesn’t save much more memory (activations for a quarter-sized micro-batch are ≈25 % of the original, but other allocations stay the same).

• bf16+accum4+checkpoint drops another ≈30 %+ because most activations are thrown away during forward and recomputed on demand in backward. The tradeoff is ≤75 % more time per step (we assert a ceiling below).

On CPU, PyTorch doesn’t track max_memory_allocated the same way, so the memory checks are skipped. The gradient-accumulation correctness check still runs.

if IS_CUDA:
    bf16_saving = 1.0 - results["bf16"]["peak_mb"] / results["fp32"]["peak_mb"]
    ckpt_saving = 1.0 - results["bf16+accum4+checkpoint"]["peak_mb"] / results["bf16+accum4"]["peak_mb"]
    step_overhead = results["bf16+accum4+checkpoint"]["step_s"] / results["bf16+accum4"]["step_s"] - 1
    print(f"bf16 saving: {bf16_saving:.1%}   ckpt saving: {ckpt_saving:.1%}   ckpt step overhead: {step_overhead:.1%}")
    s.check("bf16_saves_at_least_30pct_memory",       lambda: bf16_saving >= 0.30, msg=f"{bf16_saving:.1%}")
    s.check("checkpointing_saves_at_least_30pct",     lambda: ckpt_saving >= 0.30, msg=f"{ckpt_saving:.1%}")
    s.check("checkpoint_step_overhead_bounded",       lambda: step_overhead <= 0.75, msg=f"{step_overhead:.1%}")
else:
    s.skip("bf16_saves_at_least_30pct_memory",   "peak-memory metric requires CUDA")
    s.skip("checkpointing_saves_at_least_30pct", "peak-memory metric requires CUDA")
    s.skip("checkpoint_step_overhead_bounded",   "peak-memory metric requires CUDA")

loss_no_accum = results["bf16"]["final_loss"]
loss_accum    = results["bf16+accum4"]["final_loss"]
rel = abs(loss_accum - loss_no_accum) / max(abs(loss_no_accum), 1e-6)
s.check(
    "accum_and_big_batch_agree_within_5pct",
    lambda: rel < 0.05,
    msg=f"accum={loss_accum:.4f}  no-accum={loss_no_accum:.4f}  rel_err={rel:.1%}",
)

Peak memory and step time side by side

Each technique attacks a different corner of the budget, so the honest summary is two bars per config. Memory (left) shrinks from bf16 and again from checkpointing; step time (right) reveals the compute tax that checkpointing pays for the memory it saves. On CPU the peak-memory column is unavailable, so only the step-time panel is drawn.

import matplotlib.pyplot as plt

names = list(results.keys())
peaks_mb = [results[n]["peak_mb"] for n in names]
step_ms  = [results[n]["step_s"] * 1000 for n in names]
xs = range(len(names))

if IS_CUDA:
    fig, (axm, axt) = plt.subplots(1, 2, figsize=(9, 3.4))
    axm.bar(xs, peaks_mb, color="tab:blue")
    axm.set_xticks(list(xs)); axm.set_xticklabels(names, rotation=20, ha="right")
    axm.set_ylabel("peak MiB"); axm.set_title("peak GPU memory")
    axt.bar(xs, step_ms, color="tab:orange")
    axt.set_xticks(list(xs)); axt.set_xticklabels(names, rotation=20, ha="right")
    axt.set_ylabel("ms / step"); axt.set_title("step time")
    fig.suptitle("bf16 cuts memory; checkpointing cuts more at a time cost")
    fig.tight_layout(); plt.show()
else:
    fig, axt = plt.subplots(figsize=(5.5, 3.4))
    axt.bar(xs, step_ms, color="tab:orange")
    axt.set_xticks(list(xs)); axt.set_xticklabels(names, rotation=20, ha="right")
    axt.set_ylabel("ms / step")
    axt.set_title("step time per config (CPU; peak memory unavailable)")
    fig.tight_layout(); plt.show()

Exercises

1. Plot the tradeoff curve. For accum_steps ∈ {1, 2, 4, 8, 16}, record peak memory and step time. The curve is hyperbolic - doubling accumulation roughly halves the activation footprint but adds a launch-overhead penalty to step time.

2. Try use_reentrant=True checkpointing. On some PyTorch versions this is faster but incompatible with certain hooks. Measure the difference and read the checkpoint docs to understand why.

3. fp16 instead of bf16. Replace torch.bfloat16 with torch.float16. On a T4 you’ll need a torch.amp.GradScaler to prevent gradient underflow. Does the loss curve look the same?

4. Memory accounting. After the ablation, compute parameter memory, gradient memory, optimiser-state memory, and infer the activation memory from the peak. The first three are deterministic from n_params; the remainder is what checkpointing is fighting.

Further reading

• Micikevicius et al. 2017, Mixed Precision Training - the paper that introduced automatic loss scaling and made fp16 training practical.

• Chen et al. 2016, Training Deep Nets with Sublinear Memory Cost - the gradient-checkpointing idea.

• Rajbhandari et al. 2019, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models - the next step after checkpointing when you really need to go big.

What’s next

Notebook 02 in this track (02_ddp_vs_fsdp2) moves from single-GPU memory tricks to multi-GPU sharding. LoRA and QLoRA appear in notebooks 05 and 06.

s.summary()
s.save()