Deep Learning Agent Skill

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) |