deep-learning — train a neural net in PyTorch without the silent bugs

You own training and understanding neural networks in PyTorch: the loop, autograd, mixed precision, optimizers/schedulers, multi-GPU, and the reproducibility/checkpoint hygiene that separates a real result from a lucky one. Nets from scratch, vision, custom architectures — all here. This is PyTorch-first by design: JAX and TensorFlow are real and fine, but the patterns, APIs, and gotchas below are Torch's.

Version reality (verify at author time). Current stable is PyTorch 2.x — ~2.13 as of mid-2026 (pytorch.org/get-started, releases move fast; don't hard-pin a minor). Everything below is stable 2.x API. The one namespace shift to know: AMP now lives under torch.amp (torch.amp.GradScaler("cuda")), not the old torch.cuda.amp.*.

Am I in the right skill?

You are doing…	Skill
Training/debugging a net in PyTorch (from scratch, vision, custom loop, AMP, multi-GPU)	deep-learning (here)
Adapting a pretrained LLM — LoRA/QLoRA, SFT, trl/peft	finetuning
Classic/tabular — sklearn, XGBoost/LightGBM, feature engineering	machine-learning
Tokenization, NLP task modeling, task metrics (F1/BLEU/ROUGE)	nlp
Envs, packaging, tests and hygiene around the model code	python

1. PyTorch essentials

Three objects carry everything.

• Tensor — an n-d array on a device (cpu/cuda/mps) with a dtype. requires_grad=True makes autograd track ops on it. .to(device) / .detach() / .item() are the moves you use constantly; .item() pulls a Python scalar and drops the graph (see the loop bugs below).
• autograd — a tape. Every op on a requires_grad tensor records a node; loss.backward() walks it and accumulates into each leaf's .grad. "Accumulates" is the word that bites people (§2). Wrap read-only regions in torch.no_grad() to skip taping.
• nn.Module — the model container. __init__ registers submodules/params; forward defines compute. model.parameters() feeds the optimizer; model.train() / model.eval() flip train-vs-eval behavior for Dropout and BatchNorm.

import torch, torch.nn as nn
device = "cuda" if torch.cuda.is_available() else "cpu"

class MLP(nn.Module):
    def __init__(self, d_in, d_h, d_out, p=0.1):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(d_in, d_h), nn.ReLU(),
                                 nn.Dropout(p), nn.Linear(d_h, d_out))
    def forward(self, x):
        return self.net(x)

model = MLP(784, 256, 10).to(device)

torch.compile — wrap the model once and get a JIT-optimized (Dynamo→Inductor) graph, often a real speedup with no code change: model = torch.compile(model). Modes: "default" (balanced), "reduce-overhead" (CUDA graphs, best for small batches), "max-autotune" (Triton autotuning, slowest to warm up). First step(s) are slow — that's compilation, not your model. Keep it after .to(device) and before the loop. (pytorch.org/docs/stable/generated/torch.compile.html)

2. The training loop — and its classic bugs

The whole game is five ordered steps. Get the order or the resets wrong and it fails silently — no exception, just a model that won't learn.

model.train()                                    # Dropout/BN in TRAIN mode
for epoch in range(epochs):
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad(set_to_none=True)    # clear LAST step's grads
        out  = model(x)                          # forward
        loss = loss_fn(out, y)                    # loss  (raw logits in, e.g. CrossEntropyLoss)
        loss.backward()                          # autograd fills .grad (ACCUMULATES)
        optimizer.step()                         # update weights from .grad
        scheduler.step()                         # if a per-step scheduler

The load-bearing traps, stated next to the lines they govern:

• Missing optimizer.zero_grad() → gradients silently add up across batches. PyTorch accumulates .grad by design (so you can split a batch); if you never zero it, every step updates on the sum of all prior batches' grads and training destabilizes. Zero every step. (pytorch.org/tutorials optimization tutorial; torch.optim docs.) set_to_none=True is the default in modern Torch and is slightly faster.
• CrossEntropyLoss eats raw logits. It applies log_softmax internally — do not put a Softmax/Sigmoid before it or you double-activate and learning stalls. Same for BCEWithLogitsLoss.
• Accumulating loss without .item() leaks memory. total += loss keeps the whole autograd graph alive across the epoch. Use total += loss.item() (or .detach()).
• backward() twice on one graph raises unless you meant it (retain_graph=True) — usually a sign you forgot to zero or re-ran forward.

3. Mixed precision (AMP)

Run forward/loss in low precision for speed and memory, keep a master copy for the update. Two pieces: torch.autocast (picks per-op precision) + torch.amp.GradScaler (rescues fp16 gradients from underflow). Verified against pytorch.org/docs/stable/notes/amp_examples.html.

scaler = torch.amp.GradScaler("cuda")            # current namespace (not torch.cuda.amp)
for x, y in train_loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        out  = model(x)                          # forward + loss ONLY under autocast
        loss = loss_fn(out, y)
    scaler.scale(loss).backward()                # backward OUTSIDE autocast, on scaled loss
    scaler.step(optimizer)                        # unscales, skips step if inf/NaN grads
    scaler.update()                               # adjusts the scale for next step

bf16 vs fp16 — the choice that matters:

dtype	Exponent range	Needs GradScaler?	Use when
bf16 (bfloat16)	same as fp32	No (range can't underflow)	Ampere+ (A100/H100/…) or CPU — prefer this, it's the robust default
fp16 (float16)	narrow → underflows	Yes — scaler is mandatory	older GPUs without bf16; squeeze extra range with the scaler

• With bf16 you can skip GradScaler entirely (GradScaler(enabled=False) or just don't use it) — its whole job is fp16 underflow. With fp16 the scaler is not optional.
• Gradient clipping under AMP: scaler.unscale_(optimizer) first, then torch.nn.utils.clip_grad_norm_(...), then scaler.step. Clipping scaled grads clips the wrong magnitude.
• Autocast wraps forward + loss only. Never wrap backward().

4. Optimizers and LR schedules

• torch.optim.AdamW is the default for most modern nets (transformers, and increasingly vision): Adam with decoupled weight decay (correct L2, unlike plain Adam(weight_decay=)). SGD + momentum + Nesterov still wins on many CNNs and generalizes well — worth trying.
• Don't weight-decay everything. Biases and Norm (LayerNorm/BatchNorm) weights should get weight_decay=0; use two param groups. Decaying them hurts with no upside.
• The LR is the hyperparameter. Too high → loss NaNs or oscillates; too low → crawls. A warmup then cosine decay is the reliable modern recipe.

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01, betas=(0.9, 0.999))

# Warmup (linear ramp) THEN cosine decay, composed with SequentialLR — verified 2.x API.
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
warmup = LinearLR(opt, start_factor=0.1, total_iters=500)          # 500 steps ramp-in
cosine = CosineAnnealingLR(opt, T_max=total_steps - 500, eta_min=3e-5)
scheduler = SequentialLR(opt, [warmup, cosine], milestones=[500])  # switch at step 500
# call scheduler.step() every STEP here (T_max/total_iters are in steps).

Alternatively OneCycleLR(opt, max_lr=..., total_steps=...) bakes warmup+anneal into one scheduler. Match step() cadence to the scheduler's units — per-step schedulers stepped once per epoch (or vice-versa) silently give the wrong curve. (pytorch.org/docs/stable/optim.html)

5. Distributed training — pick by what doesn't fit

One line each on when; mechanics and launch commands in references/distributed.md.

• DDP (torch.nn.parallel.DistributedDataParallel) — replicate the whole model on each GPU, all-reduce grads every step. Use when the model fits on one GPU and you just want more throughput / a larger effective batch. Default choice; one process per GPU via torchrun.
• FSDP2 (torch.distributed.fsdp.fully_shard) — shard parameters, gradients, and optimizer states across GPUs (gather just-in-time). Use when the model does NOT fit on one GPU. FSDP2 is the current DTensor-based, per-parameter API; apply it bottom-up (wrap layers, then the root).
• DeepSpeed ZeRO (external Microsoft library, not in core Torch) — config-driven sharding: ZeRO-1 shards optimizer states, ZeRO-2 adds gradients, ZeRO-3 adds parameters (≈ FSDP), plus CPU/NVMe offload. Use when you want offload or config-first scaling, or already live in the DeepSpeed/Accelerate ecosystem. Verify the current version — it moves independently.

Rule of thumb: fits on one GPU → DDP; too big → FSDP2 (or ZeRO-3); need to offload to CPU/NVMe → ZeRO.

6. Regularization, checkpointing, reproducibility

Regularization (reach for these when val loss diverges from train — i.e. overfitting): weight decay (AdamW), nn.Dropout, data augmentation, early stopping on a val metric, label smoothing (CrossEntropyLoss(label_smoothing=0.1)), and gradient clipping for stability (not generalization).

Checkpointing — save state_dicts, not the model object. Pickling the module ties the file to your code layout and breaks on refactor.

torch.save({"epoch": epoch,
            "model": model.state_dict(),          # DDP/FSDP: model.module.state_dict(), rank 0 only
            "optimizer": opt.state_dict(),
            "scheduler": scheduler.state_dict(),
            "scaler": scaler.state_dict()},        # AMP scale must survive a resume
           "ckpt.pt")
# Resume: load model FIRST, then optimizer/scheduler; move optimizer state to device if needed.

Save optimizer + scheduler + scaler too, or a resumed run restarts its LR schedule and loss scale.

Reproducibility — seed everything, and know the ceiling. From pytorch.org/docs/stable/notes/randomness.html, quoted: "Completely reproducible results are not guaranteed across PyTorch releases, individual commits, or different platforms. Furthermore, results may not be reproducible between CPU and GPU executions, even when using identical seeds."

import random, numpy as np, torch
seed = 42
random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)   # covers CPU + all CUDA devices
torch.use_deterministic_algorithms(True)          # errors on ops with no deterministic impl
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False            # benchmark=True picks fastest algo → nondeterministic
# For CUDA >= 10.2, deterministic matmul also needs, set BEFORE launch:
#   export CUBLAS_WORKSPACE_CONFIG=:4096:8

DataLoader workers need their own seeding (each worker forks its RNG), via worker_init_fn + a seeded generator — full snippet in references/training-recipes.md. Determinism costs speed (benchmark=False disables autotuned convs).

Anti-patterns (the silent killers)

Anti-pattern	What actually happens	Instead
Validating with the model left in train() mode, or without torch.no_grad()	Dropout keeps dropping and BatchNorm uses batch stats instead of running stats, so val numbers are wrong; the graph is still built and can OOM	model.eval() + torch.no_grad() around val/inference, model.train() after (nn.Module.eval / autograd docs, pytorch.org)
Reading a NaN loss as a modeling problem	It is almost always LR too high, fp16 with no GradScaler, unclipped exploding grads, or a log(0)/0/0 in a custom loss	Bisect: drop LR 10×, clip grads, switch fp16→bf16
fp16 with the scaler left off	Gradients underflow to zero → no learning, and no error to tell you	Keep GradScaler on fp16, or move to bf16 where underflow can't happen
A stray CPU tensor in a CUDA graph	Hard device-mismatch error mid-step, usually on the targets or a hand-built mask rather than the inputs	.to(device) inputs, targets and model
Promising bit-exact reproducibility from seeds	Not guaranteed across releases, commits, platforms, or CPU-vs-GPU even with identical seeds	Seed for debugging/regression and say so; state the caveat with the result (randomness note, pytorch.org)