deep-learning · git:20260718.2408919 · 2026-07-18 · sha256 38f159902226692b
deep-learning git:20260718.2408919A
Immutable. This exact content never changes and is served at /api/v1/blob/38f159902226692b.
---
name: deep-learning
description: "Use when training or debugging a neural network in PyTorch — writing the training loop, chasing loss that won't drop or goes NaN, adding mixed precision (AMP), scaling across GPUs, or making a run reproducible. Covers tensors/autograd/nn.Module, torch.compile, the forward→loss→backward→step→zero_grad loop and its silent bugs, AdamW + cosine/warmup schedulers, DDP vs FSDP vs DeepSpeed ZeRO, checkpointing and seeds. PyTorch-first (JAX/TF exist). Triggers: 'my training loss isn't decreasing', 'loss becomes NaN after a few steps', 'how do I use torch.autocast and GradScaler', 'train this model across 4 GPUs', 'val accuracy is worse than train and I forgot model.eval()', 'gradients seem to add up across batches', 'build a CNN from scratch and train it', 'entrenar una xarxa neuronal des de zero amb PyTorch'. NOT adapting a pretrained LLM with LoRA/QLoRA (that is finetuning); NOT classic/tabular models like XGBoost or sklearn (that is machine-learning); NOT tokenization or NLP task metrics (that is nlp)."
tags: [pytorch, deep-learning, neural-networks, training-loop, autograd, mixed-precision, amp, ddp, fsdp, optimizer, lr-scheduler, reproducibility, torch-compile]
recommends: [machine-learning, finetuning, nlp, python]
origin: risco
---
# 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`](../finetuning/SKILL.md) |
| Classic/tabular — sklearn, XGBoost/LightGBM, feature engineering | [`machine-learning`](../machine-learning/SKILL.md) |
| Tokenization, NLP task modeling, task metrics (F1/BLEU/ROUGE) | [`nlp`](../nlp/SKILL.md) |
If the model is a from-scratch or vision/custom net and you're writing the loop yourself, stay.
## 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.
```python
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.
```python
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 (all in the gotchas table, worth stating once here):
- **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`.
```python
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.
```python
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_dict`s, not the model object.** Pickling the module ties the file to
your code layout and breaks on refactor.
```python
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."*
So bit-exact reproducibility is a debugging/regression tool, not a guarantee.
```python
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).
## Guardrails (the silent killers)
- **`optimizer.zero_grad()` every step.** Skipping it makes `.grad` accumulate across batches — no
error, just diverging training. (`torch.optim` / optimization tutorial, pytorch.org)
- **`model.eval()` + `torch.no_grad()` at inference/validation.** Forgetting `eval()` leaves Dropout
dropping and BatchNorm using *batch* stats instead of running stats, so val numbers are wrong;
forgetting `no_grad()` builds the autograd graph and can OOM. Flip back to `model.train()` after.
(`nn.Module.eval` / autograd docs, pytorch.org)
- **Determinism is not guaranteed across releases/platforms/CPU-vs-GPU** even with identical seeds —
seed for debugging, don't promise bit-exactness. (randomness note, pytorch.org)
- **fp16 without `GradScaler`** underflows to zero grads → no learning. bf16 doesn't need it.
- **`CrossEntropyLoss`/`BCEWithLogitsLoss` take raw logits** — no Softmax/Sigmoid before them.
- **Everything on the same device** — a stray CPU tensor in a CUDA graph is a hard error you'll hit
fast; `.to(device)` inputs, targets, and model.
- **NaN loss?** LR too high, no AMP scaler on fp16, unclipped exploding grads, or a `log(0)`/`0/0`
in a custom loss. Bisect: drop LR 10×, clip grads, switch fp16→bf16.
## Related skills
- [`finetuning`](../finetuning/SKILL.md) — adapting a **pretrained LLM** (LoRA/QLoRA, SFT via
trl/peft). If the base is a foundation model you're specializing, that's the skill; this one is
general/from-scratch net training.
- [`machine-learning`](../machine-learning/SKILL.md) — classic/**tabular** (sklearn, GBDT). No
autograd, no GPU loop.
- [`nlp`](../nlp/SKILL.md) — tokenization, NLP **task** modeling and metrics. Route text task
specifics there; route the raw training loop back here.
- [`python`](../python/SKILL.md) — env/packaging, testing, general Python hygiene under the model
code.
## Checklist
- [ ] Right skill: from-scratch/vision/custom net in PyTorch (not finetuning / ML / nlp).
- [ ] Loop has all five steps in order; `zero_grad(set_to_none=True)` every step.
- [ ] Loss takes raw logits (no double activation); running loss uses `.item()`.
- [ ] `model.train()` for training; `model.eval()` **and** `torch.no_grad()` for val/inference.
- [ ] AMP: `torch.autocast` on forward+loss only; fp16 → `GradScaler`, bf16 → skip it; clip after
`unscale_`.
- [ ] AdamW (decoupled decay) with no weight decay on bias/Norm params; LR warmup→cosine, `step()`
cadence matches scheduler units.
- [ ] Multi-GPU chosen by fit: DDP (fits) vs FSDP2 (doesn't) vs ZeRO (offload/config).
- [ ] Checkpoint saves model+optimizer+scheduler+scaler `state_dict`s; resume loads model first.
- [ ] Seeds set; determinism flags on if needed, with the "not guaranteed across releases/devices"
caveat acknowledged.
- [ ] Version claims hedged to "PyTorch 2.x (~2.13, verify)"; no hard minor pin presented as current.