PyTorch Training skill for Claude Code

A training script is correct when it overfits a single small batch to near-zero loss, resumes from a checkpoint to the same loss curve, keeps the GPU busy (data loading is not the bottleneck), and logs train and validation metrics you can compare across runs. Scale up only after those hold on a tiny run.

Ground rules

• Debug small: tiny model or subset, a few hundred steps, seconds per iteration. Scale only once the small run behaves.
• Always overfit one batch first. If the model cannot memorize 8 examples, there is a bug.
• Log the configuration, seed, git commit, and library versions with every run.
• Checkpoint everything needed to resume exactly: model, optimizer, scheduler, scaler, step, epoch, RNG states.
• Validate on a held-out split on a fixed schedule; never select checkpoints by train loss.

Workflow

- [ ] 1. Inspect data: shapes, dtypes, label distribution, a few decoded samples
- [ ] 2. Build Dataset/DataLoader; time an epoch of loading alone
- [ ] 3. Write the loop; overfit one batch
- [ ] 4. Add mixed precision, accumulation, clipping, scheduler
- [ ] 5. Add checkpoint save/resume; verify resume
- [ ] 6. Profile throughput; then torch.compile / DDP
- [ ] 7. Full run with validation, early stopping, experiment tracking

2. Data loading

from torch.utils.data import DataLoader

train_loader = DataLoader(
    train_ds,
    batch_size=64,
    shuffle=True,
    num_workers=8,              # start near CPU cores per GPU; measure
    pin_memory=True,            # faster host-to-GPU copies
    persistent_workers=True,    # avoid re-forking each epoch
    prefetch_factor=4,
    drop_last=True,             # stable batch size for BatchNorm/compile
    worker_init_fn=seed_worker, generator=g,   # reproducible augmentation
)

• Do heavy decoding and augmentation in workers, not in the main loop.
• Move batches with x.to(device, non_blocking=True) (pairs with pin_memory).
• Measure: iterate the loader alone for 200 batches. If that is slower than a training step, data loading is the bottleneck; fix it before anything else.
• With num_workers > 0, each worker gets a copy of the dataset object; open files and DB handles lazily in __getitem__ or worker_init_fn, not in __init__.

3. The loop

import torch
from torch.nn.utils import clip_grad_norm_

torch.set_float32_matmul_precision("high")      # TF32 on Ampere+ for fp32 matmuls
device = "cuda"
model = Model().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1, fused=True)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=3e-4, total_steps=total_steps, pct_start=0.05)
accum = 4

for epoch in range(start_epoch, epochs):
    model.train()
    for i, (x, y) in enumerate(train_loader):
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            loss = criterion(model(x), y) / accum
        loss.backward()
        if (i + 1) % accum == 0:
            grad_norm = clip_grad_norm_(model.parameters(), max_norm=1.0)
            opt.step()
            sched.step()
            opt.zero_grad(set_to_none=True)
            step += 1
            if step % log_every == 0:
                log(step=step, loss=loss.item() * accum, grad_norm=grad_norm.item(),
                    lr=sched.get_last_lr()[0])
    evaluate(model, val_loader)
    save_checkpoint(...)

@torch.no_grad()
def evaluate(model, loader):
    model.eval()                      # disables dropout, uses BatchNorm running stats
    ...
    model.train()

Loop rules:

• model.train() / model.eval() at the right places; forgetting eval() makes validation noisy and wrong with dropout or BatchNorm.
• Wrap evaluation in torch.no_grad() (or torch.inference_mode()).
• Call loss.item() only when logging; it forces a GPU sync every time.
• Scheduler steps per optimizer step (not per micro-batch) for step-based schedules.

4. Precision

Hardware	Setting
Ampere / Hopper / newer NVIDIA, recent AMD, TPU	torch.autocast(..., dtype=torch.bfloat16), no GradScaler
Older NVIDIA (V100, T4)	fp16 autocast + torch.amp.GradScaler("cuda")
Apple MPS	Start in fp32; try torch.autocast("mps", dtype=torch.float16) if supported and stable

fp16 with GradScaler:

scaler = torch.amp.GradScaler("cuda")
with torch.autocast(device_type="cuda", dtype=torch.float16):
    loss = criterion(model(x), y)
scaler.scale(loss).backward()
scaler.unscale_(opt)                       # before clipping
clip_grad_norm_(model.parameters(), 1.0)
scaler.step(opt)
scaler.update()
opt.zero_grad(set_to_none=True)

Keep the loss computation in fp32 when it involves large reductions or log-softmax outside the model (autocast handles common ops; check custom ones).

5. Checkpointing

def save_checkpoint(path):
    state = {
        "model": model.state_dict(), "optimizer": opt.state_dict(),
        "scheduler": sched.state_dict(), "scaler": scaler.state_dict() if scaler else None,
        "step": step, "epoch": epoch, "config": cfg,
        "rng": {"torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state_all()},
    }
    tmp = path + ".tmp"
    torch.save(state, tmp)
    os.replace(tmp, path)              # atomic: a crash never leaves a corrupt checkpoint

state = torch.load(path, map_location="cpu", weights_only=False)  # trusted files only

• torch.load defaults to weights_only=True in current PyTorch; full training state with config objects needs weights_only=False, which unpickles arbitrary code. Only do that for files you created.
• With torch.compile or DDP, save model._orig_mod.state_dict() / model.module.state_dict() so keys load into a plain model.
• Verify resume: train 200 steps, checkpoint at 100, resume, and compare the loss at steps 101-200. They should match closely (exactly with deterministic settings).
• Keep the last N and the best-by-validation checkpoint; delete the rest automatically.

6. Speed

1. Profile one window with torch.profiler or simply time data vs. forward/backward.
2. Fix data loading first (see above).
3. Raise batch size until memory is ~85-90% used; scale learning rate accordingly and re-check stability.
4. model = torch.compile(model); expect a slow first iteration and recompiles if shapes vary (pad or bucket sequence lengths).
5. Out of memory: lower micro-batch and raise accumulation, enable activation checkpointing, use bf16, or shard with FSDP.

Multi-GPU (single node):

torchrun --standalone --nproc_per_node=8 train.py --config cfg.yaml

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
model = DDP(model.to(local_rank), device_ids=[local_rank])
sampler = DistributedSampler(train_ds, shuffle=True)
loader = DataLoader(train_ds, batch_size=per_gpu_bs, sampler=sampler, num_workers=8, pin_memory=True)
for epoch in range(epochs):
    sampler.set_epoch(epoch)           # otherwise every epoch has the same order
    ...
# log, evaluate, and save only on rank 0: if dist.get_rank() == 0: ...

Use model.no_sync() on non-final accumulation micro-steps to skip redundant gradient all-reduces. Models that do not fit on one GPU need FSDP; that is a separate design.

Reproducibility

Seed Python, NumPy, and torch; seed DataLoader workers via worker_init_fn and a torch.Generator. For bit-exact debugging, torch.use_deterministic_algorithms(True) (slower; some ops raise or need CUBLAS_WORKSPACE_CONFIG=:4096:8).

Loss looks wrong

Quick triage; for the full decision tree read references/loss-debugging.md.

Symptom	First checks
NaN/Inf after some steps	Learning rate too high, fp16 overflow (switch to bf16 or check scaler), log of zero, division by zero in custom loss
Loss flat from step 0	Gradients not flowing (requires_grad, detached tensors), LR far too low, labels misaligned with inputs
Train down, val up	Overfitting: more data/augmentation, regularization, early stopping
Val far better than train	Dropout/augmentation only in train (normal), or leakage between splits
Spiky loss	LR too high or no warmup, bad samples, batch too small

Output contract

**Run:** <name> commit <sha>, config <path>, seed <n>, torch <version>
**Sanity:** overfit-one-batch loss <x>; resume check matched: <yes/no>
**Throughput:** <samples/s>, GPU util <x%>, data-loader time share <y%>
**Results:** best val <metric> at step <n>; checkpoint <path>
**Issues fixed:** <bug -> evidence -> fix>