Mixed Precision and Gradient Checkpointing

Two memory + speed tools you'll always reach for

Already covered AMP basics in Track 2. This lesson gathers the modern combinations and adds gradient checkpointing — the trick that trades compute for memory.

AMP recap — the modern bf16 path

On Ampere+ GPUs (A100, RTX 30/40, H100) and Apple Silicon, prefer autocast(device_type='cuda', dtype=torch.bfloat16) with no GradScaler. Same speed as fp16, none of the underflow drama. fp16 + GradScaler is still around for older hardware (V100, T4) but for new training runs on modern GPUs, bf16 wins.

Gradient checkpointing

Normally, the forward pass stores every intermediate activation so backward can use it. For deep networks (a 24-layer Transformer with sequence length 4K) those activations dominate memory. Gradient checkpointing doesn't store some activations — it recomputes them on demand during backward. Memory drops by 50-70%, training time goes up by ~30%. Almost always worth it when you'd otherwise OOM.

Composing them

AMP + gradient checkpointing + torch.compile + FSDP all compose. The order to add them when scaling up:

Start eager fp32 — get correctness.
Add bf16 autocast — free 1.5–2x speed.
Add torch.compile — another ~1.5–2x.
If still OOM at desired batch size: gradient checkpointing.
If model itself doesn't fit: FSDP.

Code

bf16 mixed precision — the modern recipe·python

import torch
from torch.amp import autocast

# bf16 — preferred on Ampere+ and Apple Silicon. NO GradScaler.
for x, y in loader:
    x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.bfloat16):
        out = model(x)
        loss = criterion(out, y)

    loss.backward()
    optimizer.step()

fp16 mixed precision — older GPUs, GradScaler required·python

import torch
from torch.amp import autocast, GradScaler

scaler = GradScaler('cuda')

for x, y in loader:
    x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.float16):
        out = model(x)
        loss = criterion(out, y)

    scaler.scale(loss).backward()    # scale loss to avoid fp16 underflow
    scaler.step(optimizer)            # unscale + step (skip on inf/nan)
    scaler.update()                   # adjust scale dynamically

Gradient checkpointing — recompute instead of store·python

import torch.nn as nn
from torch.utils.checkpoint import checkpoint

class CheckpointedTransformer(nn.Module):
    def __init__(self, num_layers):
        super().__init__()
        self.layers = nn.ModuleList([TransformerBlock() for _ in range(num_layers)])

    def forward(self, x):
        for layer in self.layers:
            # Don't store this layer's activations; recompute in backward.
            x = checkpoint(layer, x, use_reentrant=False)
        return x

# For HuggingFace transformers, there's a one-liner:
# model.gradient_checkpointing_enable()

Combining bf16 + compile + checkpointing — the production stack·python

import torch
import torch.nn as nn
from torch.amp import autocast

model = MyTransformer().cuda()
model.gradient_checkpointing_enable()       # if available
model = torch.compile(model, mode='default')

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

for epoch in range(num_epochs):
    for x, y in loader:
        x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
        optimizer.zero_grad(set_to_none=True)

        with autocast(device_type='cuda', dtype=torch.bfloat16):
            out = model(x)
            loss = criterion(out, y)

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

Mixed Precision and Gradient Checkpointing

Two memory + speed tools you'll always reach for

AMP recap — the modern bf16 path

Gradient checkpointing

Composing them

Code

External links

Exercise

Progress

Comments 0