PyTorch Training Best Practices

Real-World Scenario

An ML engineer trains a text classification model on a single GPU. The first version of the training loop works, but it is slow, runs out of memory at the batch size they want, and loses hours of progress whenever the job is preempted. Mixed precision, gradient accumulation, a warmup schedule, and proper checkpoints fix each of those — they are the standard configuration of a modern PyTorch training loop, and each can be checked on a CPU.

Mixed Precision: What autocast Changes

torch.autocast runs selected operations in a lower-precision dtype. The examples here use bfloat16 on CPU; on a CUDA GPU the usual choice is float16.

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(256, 512), nn.ReLU(), nn.Linear(512, 10))
x, y = torch.randn(64, 256), torch.randint(0, 10, (64,))

with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    logits = model(x)
    loss = nn.functional.cross_entropy(logits, y)

print("weights:", model[0].weight.dtype)
print("logits: ", logits.dtype)
print("loss:   ", loss.dtype)

loss.backward()
print("grads:  ", model[0].weight.grad.dtype)

$ python autocast_dtypes.py
weights: torch.float32
logits:  torch.bfloat16
loss:    torch.float32
grads:   torch.float32

Only the operations autocast considers safe change dtype: the linear layers produce bfloat16 logits, but the cross-entropy loss is computed in float32, because summing many small log-probabilities in low precision loses accuracy. The weights themselves stay float32, and so do their gradients — autocast casts copies on the fly, it never converts the model.

Why float16 Needs a Gradient Scaler

import torch

grad = 1e-8                                   # a realistic tiny gradient value
print(torch.tensor(grad, dtype=torch.float16))
print(torch.tensor(grad, dtype=torch.bfloat16))
print(torch.tensor(grad * 65536, dtype=torch.float16))   # scaled before casting

for dt in (torch.float16, torch.bfloat16, torch.float32):
    fi = torch.finfo(dt)
    print(f"{str(dt):15s} max {fi.max:.3g}  smallest normal {fi.tiny:.3g}  eps {fi.eps:.3g}")

$ python underflow.py
tensor(0., dtype=torch.float16)
tensor(1.0012e-08, dtype=torch.bfloat16)
tensor(0.0007, dtype=torch.float16)
torch.float16   max 6.55e+04  smallest normal 6.1e-05  eps 0.000977
torch.bfloat16  max 3.39e+38  smallest normal 1.18e-38  eps 0.00781
torch.float32   max 3.4e+38  smallest normal 1.18e-38  eps 1.19e-07

A gradient of 1e-8 becomes exactly zero in float16, whose smallest normal value is 6.1e-05 (with subnormals reaching down only to about 6e-8). Multiply it by 65,536 first and it survives. That is what GradScaler does: scale the loss up before backward(), then divide the gradients back down before the optimizer uses them. bfloat16 keeps float32’s exponent range — the same 1.18e-38 smallest normal — and trades away precision instead (eps of 0.0078), so it doesn’t underflow and doesn’t need a scaler.

A Mixed-Precision Training Loop

One loop that uses float16 plus a scaler on CUDA and bfloat16 without one on CPU:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
amp_dtype = torch.float16 if device.type == "cuda" else torch.bfloat16
torch.manual_seed(0)

X = torch.randn(4000, 256)
y = (X[:, :10].sum(dim=1) > 0).long()
loader = DataLoader(TensorDataset(X, y), batch_size=64, shuffle=True)
model = nn.Sequential(nn.Linear(256, 512), nn.ReLU(), nn.Linear(512, 2)).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler = torch.amp.GradScaler(device.type, enabled=(amp_dtype == torch.float16))

for epoch in range(1, 4):
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()
        with torch.autocast(device_type=device.type, dtype=amp_dtype):
            loss = nn.functional.cross_entropy(model(xb), yb)
        scaler.scale(loss).backward()          # multiplies by the scale when enabled
        scaler.unscale_(optimizer)             # back to true gradients before clipping
        norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)                 # skips the step if grads hold inf/NaN
        scaler.update()
    print(f"epoch {epoch}  loss {loss.item():.4f}  grad norm before clip {norm:.3f}")
print(device, amp_dtype, "scaler enabled:", scaler.is_enabled())

$ python amp_loop.py
epoch 1  loss 0.2592  grad norm before clip 0.889
epoch 2  loss 0.1299  grad norm before clip 0.850
epoch 3  loss 0.0582  grad norm before clip 0.569
cpu torch.bfloat16 scaler enabled: False

The order of calls matters: scaler.scale(loss).backward() produces scaled gradients, unscale_ divides them back so that clip_grad_norm_ compares the real norm against max_norm, and scaler.step skips the update entirely if the gradients contain inf or NaN — the sign that the scale was too high, which update() then lowers. With enabled=False, as here on CPU, every scaler call passes straight through, so the same code runs on both devices. clip_grad_norm_ returns the norm before clipping, which is worth logging: a sudden spike is often the first sign of a diverging run.

Gradient Accumulation

Gradients add up across backward() calls until they are zeroed. Accumulating over several small batches reproduces one large batch — if the losses are scaled correctly.

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Linear(8, 3)
X, y = torch.randn(64, 8), torch.randint(0, 3, (64,))
loss_fn = nn.CrossEntropyLoss()                 # averages over the batch

loss_fn(model(X), y).backward()                 # one batch of 64
full = model.weight.grad.clone()

for divide in (True, False):
    model.zero_grad()
    for xb, yb in zip(X.split(16), y.split(16)):  # 4 micro-batches of 16
        loss = loss_fn(model(xb), yb)
        (loss / 4 if divide else loss).backward()  # grads add up across calls
    ratio = (model.weight.grad / full).mean().item()
    print(f"divide by 4={divide!s:5}  matches full batch: "
          f"{torch.allclose(model.weight.grad, full, atol=1e-6)}  ratio {ratio:.2f}")

$ python accumulation.py
divide by 4=True   matches full batch: True  ratio 1.00
divide by 4=False  matches full batch: False  ratio 4.00

CrossEntropyLoss averages over the batch, so each micro-batch’s loss is an average over 16 examples. Summing four of those gives four times the full-batch average — the ratio 4.00 line. For SGD that is the same as silently multiplying the learning rate by four; Adam and AdamW partly cancel a constant factor, which hides the bug rather than fixing it. Dividing each loss by the number of accumulation steps makes the accumulated gradient equal to the 64-example batch’s. In a real loop, call optimizer.step() and optimizer.zero_grad() once every four micro-batches.

Learning Rate Warmup

Large learning rates at the start of training — when the weights are random and Adam’s statistics are empty — can destabilize a run. Warmup ramps the learning rate up first.

import math
import torch
import torch.nn as nn

opt = torch.optim.AdamW(nn.Linear(4, 1).parameters(), lr=1e-3)
WARMUP, TOTAL = 100, 1000

def warmup_cosine(step):
    if step < WARMUP:
        return (step + 1) / WARMUP                             # linear ramp up
    progress = (step - WARMUP) / (TOTAL - WARMUP)
    return 0.5 * (1 + math.cos(math.pi * progress))            # cosine down to 0

sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda=warmup_cosine)
for step in range(TOTAL):
    if step in (0, 49, 99, 100, 325, 550, 775, 999):
        print(f"step {step:4d}  lr {sched.get_last_lr()[0]:.2e}")
    opt.step()
    sched.step()                                               # once per batch

$ python warmup.py
step    0  lr 1.00e-05
step   49  lr 5.00e-04
step   99  lr 1.00e-03
step  100  lr 1.00e-03
step  325  lr 8.54e-04
step  550  lr 5.00e-04
step  775  lr 1.46e-04
step  999  lr 3.05e-09

LambdaLR multiplies the base learning rate by whatever the function returns. For the first 100 steps that ramps linearly to 1e-3; after that, a cosine curve brings it down, passing half the peak at step 550, the midpoint of the decay, and reaching nearly zero at the end. This scheduler is stepped once per batch, not per epoch — getting that wrong stretches the warmup over 100 epochs.

OneCycleLR packages the same idea:

import torch
import torch.nn as nn

opt = torch.optim.AdamW(nn.Linear(4, 1).parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=1e-2, total_steps=100, pct_start=0.3)

for step in range(100):
    if step % 15 == 0 or step == 99:
        print(f"step {step:3d}  lr {sched.get_last_lr()[0]:.2e}")
    opt.step()
    sched.step()

$ python onecycle.py
step   0  lr 4.00e-04
step  15  lr 5.46e-03
step  30  lr 9.99e-03
step  45  lr 8.77e-03
step  60  lr 5.89e-03
step  75  lr 2.63e-03
step  90  lr 4.02e-04
step  99  lr 4.00e-08

It starts at max_lr / 25, rises to max_lr over the first 30% of steps, then anneals to a value far below the starting point. It also ignores the optimizer’s lr=1e-3 — the schedule is fully defined by max_lr.

Reduce on Plateau

import torch
import torch.nn as nn

opt = torch.optim.AdamW(nn.Linear(4, 1).parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, mode="min", patience=2, factor=0.5)

val_losses = [0.90, 0.70, 0.60, 0.61, 0.60, 0.62, 0.55, 0.56, 0.57, 0.58, 0.56]
for epoch, val_loss in enumerate(val_losses):
    sched.step(val_loss)                     # pass the validation metric
    print(f"epoch {epoch:2d}  val_loss {val_loss:.2f}  lr {opt.param_groups[0]['lr']:.2e}")

$ python plateau.py
epoch  0  val_loss 0.90  lr 1.00e-03
epoch  1  val_loss 0.70  lr 1.00e-03
epoch  2  val_loss 0.60  lr 1.00e-03
epoch  3  val_loss 0.61  lr 1.00e-03
epoch  4  val_loss 0.60  lr 1.00e-03
epoch  5  val_loss 0.62  lr 5.00e-04
epoch  6  val_loss 0.55  lr 5.00e-04
epoch  7  val_loss 0.56  lr 5.00e-04
epoch  8  val_loss 0.57  lr 5.00e-04
epoch  9  val_loss 0.58  lr 2.50e-04
epoch 10  val_loss 0.56  lr 2.50e-04

ReduceLROnPlateau watches the metric you pass to step() — here, a fixed list standing in for validation losses. The best value is 0.60 at epoch 2; epoch 4 ties it but does not improve on it, so epochs 3, 4, and 5 count as three bad epochs, one more than patience=2, and the rate halves at epoch 5. The same thing happens after the new best of 0.55 at epoch 6. Pass validation loss, not training loss, which keeps falling while the model overfits.

Checkpoint and Resume — Exactly

A checkpoint is only complete if resuming from it gives the same result as never stopping.

import torch
import torch.nn as nn

def setup():
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
    opt = torch.optim.AdamW(model.parameters(), lr=1e-2)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10)
    return model, opt, sched

def train(model, opt, sched, epochs):
    for epoch in epochs:
        g = torch.Generator().manual_seed(epoch)          # data for this epoch
        X = torch.randn(64, 10, generator=g)
        opt.zero_grad()
        loss = nn.functional.mse_loss(model(X), X.sum(1, keepdim=True))
        loss.backward()
        opt.step()
        sched.step()
    return loss.item()

model, opt, sched = setup()
print(f"uninterrupted: {train(model, opt, sched, range(10)):.6f}")

model, opt, sched = setup()
train(model, opt, sched, range(5))
torch.save({"epoch": 4, "model": model.state_dict(), "opt": opt.state_dict(),
            "sched": sched.state_dict()}, "ckpt.pt")

model, opt, sched = setup()                                # a fresh process would do this
ckpt = torch.load("ckpt.pt", weights_only=True)
model.load_state_dict(ckpt["model"]); opt.load_state_dict(ckpt["opt"])
sched.load_state_dict(ckpt["sched"])
print(f"resumed:       {train(model, opt, sched, range(ckpt['epoch'] + 1, 10)):.6f}")

$ python resume.py
uninterrupted: 7.319561
resumed:       7.319561

The run was trained for five epochs, saved, rebuilt from scratch, restored, and trained for the remaining five — and it ends at exactly the same loss, to six decimal places, as the uninterrupted run. That requires three state_dicts: the model’s weights, the optimizer’s Adam statistics, and the scheduler’s position in its cosine curve. Leave out either of the last two and the resumed run silently diverges. Seeding each epoch’s data by its epoch number makes the data reproducible too. This model has no dropout; one that does also needs torch.get_rng_state() saved and restored with torch.set_rng_state(), or the resumed run draws different dropout masks.

torch.compile

torch.compile traces a model and generates fused kernels for it. The first call pays for the compilation.

import time
import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(512, 1024), nn.GELU(), nn.Linear(1024, 1024), nn.GELU(),
                      nn.Linear(1024, 10)).eval()
x = torch.randn(64, 512)
compiled = torch.compile(model)

def ms_per_call(fn, n=200):
    t0 = time.perf_counter()
    with torch.no_grad():
        for _ in range(n):
            fn(x)
    return (time.perf_counter() - t0) / n * 1000

t0 = time.perf_counter()
with torch.no_grad():
    compiled(x)                                  # first call compiles
print(f"first call (compile): {time.perf_counter() - t0:.1f} s")
ms_per_call(model, 20)                           # warm up eager too
print(f"eager:    {ms_per_call(model):.3f} ms")
print(f"compiled: {ms_per_call(compiled):.3f} ms")
with torch.no_grad():
    print("same result:", torch.allclose(model(x), compiled(x), atol=1e-5))

$ python compile_speed.py
first call (compile): 2.5 s
eager:    1.658 ms
compiled: 1.476 ms
same result: True

The results match. The first call took 2.5 seconds here because Inductor caches compiled kernels on disk and this machine had compiled the model before; with an empty cache the first compile takes considerably longer. The payoff is modest: about 10% faster per call on this CPU, and the size of that gap varies from run to run. This model is three large matrix multiplications, which eager PyTorch already hands to optimized BLAS routines, so there is little overhead left for fusion to remove. torch.compile pays off where there is more: many small element-wise operations, Python overhead between them, and especially GPUs. Measure it on your own model and hardware, and remember that a change in input shapes can trigger recompilation.

Practice

import torch

for dt in (torch.float32, torch.float16, torch.bfloat16):
    one = torch.tensor(1.0, dtype=dt)
    print(f"{str(dt):15s} 1 + 0.001 = {(one + 0.001).item():.6f}   "
          f"1000 + 1 = {(torch.tensor(1000.0, dtype=dt) + 1).item():.1f}")

$ python p1_precision.py
torch.float32   1 + 0.001 = 1.001000   1000 + 1 = 1001.0
torch.float16   1 + 0.001 = 1.000977   1000 + 1 = 1001.0
torch.bfloat16  1 + 0.001 = 1.000000   1000 + 1 = 1000.0

val_losses = [0.90, 0.70, 0.60, 0.61, 0.58, 0.59, 0.60, 0.61, 0.57]
PATIENCE = 3

best, bad_epochs = float("inf"), 0
for epoch, loss in enumerate(val_losses):
    if loss < best:
        best, bad_epochs = loss, 0
        print(f"epoch {epoch}: {loss:.2f} new best -> save checkpoint")
    else:
        bad_epochs += 1
        print(f"epoch {epoch}: {loss:.2f} no improvement ({bad_epochs}/{PATIENCE})")
        if bad_epochs == PATIENCE:
            print(f"stopping; best was {best:.2f}")
            break

$ python p2_early_stop.py
epoch 0: 0.90 new best -> save checkpoint
epoch 1: 0.70 new best -> save checkpoint
epoch 2: 0.60 new best -> save checkpoint
epoch 3: 0.61 no improvement (1/3)
epoch 4: 0.58 new best -> save checkpoint
epoch 5: 0.59 no improvement (1/3)
epoch 6: 0.60 no improvement (2/3)
epoch 7: 0.61 no improvement (3/3)
stopping; best was 0.58

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Linear(8, 3)
X, y = torch.randn(70, 8), torch.randint(0, 3, (70,))       # 70 = 16+16+16+16+6
loss_fn = nn.CrossEntropyLoss()

loss_fn(model(X), y).backward()
full = model.weight.grad.clone()

for name, weight in [("divide by 5", lambda n: 1 / 5), ("weight by size", lambda n: n / 70)]:
    model.zero_grad()
    for xb, yb in zip(X.split(16), y.split(16)):
        (loss_fn(model(xb), yb) * weight(len(yb))).backward()
    print(f"{name:15s} matches full batch: {torch.allclose(model.weight.grad, full, atol=1e-6)}")

$ python p3_uneven.py
divide by 5     matches full batch: False
weight by size  matches full batch: True