dl-training

name

dl-training

description

Deep DL-training operational intuition — initialization, optimizer choice (AdamW/Muon/Schedule-Free), decoupled weight decay, LR schedules (cosine/WSD), μP transfer, mixed precision (bfloat16/FP8), normalization (RMSNorm, Peri-LN), NaN forensics, loss-spike rewind. Load when picking optimizer/schedule, weight-decay setup, debugging loss spikes/NaN, normalization choice, or μP transfer. Skip for architecture selection, inference/serving, or sharding topology. Pairs with `pytorch`. Triggers on: "AdamW", "Muon", "Schedule-Free", "decoupled weight decay", "WSD", "muP", "μTransfer", "GradScaler", "bfloat16", "FP8", "RMSNorm", "Peri-LN", "Fixup", "Mixup", "loss spike", "spike skip".

DL Training Operational Guide

Concise operational pointers for deep-learning training fundamentals where models tend to be shallow.

Assumes you already know backprop, MLPs, common architectures (CNN/transformer), and basic PyTorch. This skill covers the operational layer — initialization, optimizer choice, schedule, weight-decay correctness, gradient/NaN forensics, normalization, augmentation, mixed precision, μP, loss-spike forensics — current as of late 2025/early 2026. Pairs with pytorch for torch.compile/FSDP/profiler.

When to use

Load when the task is:

• Picking or tuning an optimizer for a new model (CV, transformer, diffusion)
• Choosing/configuring an LR schedule (cosine, WSD, OneCycle, Schedule-Free)
• Diagnosing loss spikes, NaN, or divergence
• Migrating a run from FP32 → BF16/FP8, or AMP-related instability
• Configuring weight-decay parameter groups (bias / norm exclusion)
• Setting gradient clipping for transformers vs RNNs
• Initializing a deep residual stack or ViT
• Adding Mixup/CutMix/label smoothing to an image pipeline
• Standing up large-batch training and choosing LR scaling rule
• Transferring HPs across model widths (μP / μTransfer)
• Configuring SWA or EMA for the final training phase or diffusion
• Memory-budget engineering: activation checkpointing, ZeRO, sequence packing
• PreLN vs PostLN vs Peri-LN, or BN vs GN vs LN vs RMSNorm
• DropPath / stochastic depth on a deep ViT/ResNet

Do NOT load for: architecture choice (attention variants, MoE routing, tokenizer), inference/serving, distributed-systems topology, RL-specific dynamics, pure data-pipeline perf.

Initialization

PyTorch's nn.Linear / nn.Conv2d default is kaiming_uniform_(a=sqrt(5)) — historical, not Kaiming-for-ReLU. It approximates a wide uniform; bias is U(−1/sqrt(fan_in), 1/sqrt(fan_in)). Don't trust it: re-init explicitly.

• Xavier/Glorot (Var = 2/(fan_in+fan_out)) — tanh/sigmoid/identity.
• He/Kaiming (Var = 2/fan_in) — ReLU/GELU.
• Orthogonal — RNNs and where preserving rotation matters.
• Residual scaling: GPT-2 scales output projections of attn and MLP residual paths by 1/sqrt(2L) (L = block count; the 2 reflects two residual writes per block). Keeps activation variance ~constant with depth. Used by virtually every modern decoder.
• Fixup (Zhang et al., 2019): trains 10k-layer ResNets without normalization by zero-init the last layer of each residual branch, scaling earlier layers by L^(−1/(2m−2)), plus per-branch scalar/bias parameters. For when you intentionally drop norms.
• Embedding init: N(0, 0.02) is the GPT-2 convention; 0.006 ≈ 1/sqrt(d_model) for d_model ≈ 40k. Tied input/output embeddings inherit init.

Optimizers

• SGD+momentum β=0.9, weight_decay 1e-4. Still SOTA on ImageNet ResNet, mostly anywhere CV throughput matters and BN is in play. Slow on transformers.
• AdamW (Loshchilov & Hutter): β1=0.9, β2=0.999 (drop to 0.95 for LLM pretrain — long-horizon variance), ε=1e-8, weight_decay 0.01–0.1. Decoupled weight decay: θ ← θ − lr·(m̂/(√v̂+ε) + λθ), NOT L2-in-loss. PyTorch's Adam(weight_decay=...) is the L2 (wrong) form for adaptive moments — use AdamW.
• Lion (Chen et al., 2023): update = sign(β1·m + (1−β1)·g), β1=0.9, β2=0.99. Sign flattens magnitudes, so LR must be 3–10× smaller than AdamW and weight_decay 3–10× larger to keep lr·λ constant. Half the optimizer-state memory of AdamW (no v).
• Sophia (Liu et al., 2023): diagonal-Hessian preconditioner estimated every k steps; element-wise clipping bounds the worst-case update. Reports ~2× wall-clock vs AdamW on GPT-2-class. Adoption mixed; sensitive to clip threshold ρ.
• Muon (Jordan, 2024 — "MomentUm Orthogonalized by Newton-Schulz"): for hidden 2D matrices only — momentum then Newton-Schulz (5 iters) orthogonalizes the update so no direction dominates. Embeddings/heads/scalars stay on AdamW. ~2× efficient for LLMs in 2025; used in Moonshot/Kimi-class runs and "Distributed Muon" (ZeRO-1 friendly). Footgun: matrix LR ≈ 0.02 is much higher than AdamW's; weight decay also retuned.
• Schedule-Free AdamW (Defazio et al., 2024): no LR schedule; replaces momentum with interpolation+averaging. Won MLCommons AlgoPerf 2024 self-tuning. Useful when training horizon is unknown or one config must serve many durations.

Weight decay correctness

L2-in-loss adds λ·‖θ‖² to objective; gradient becomes g + 2λθ. With Adam this is then divided by √v̂, so heavily-updated params get less effective decay — backwards. Decoupled (AdamW) applies θ ← θ(1 − lr·λ) outside the adaptive step.

Always exclude from decay: biases, LayerNorm/RMSNorm γ and β, BatchNorm γ/β, position embeddings (often), token embeddings (sometimes — varies). Decaying norm scales pulls them toward zero, killing the affine recovery. Build two param groups: matrices (p.dim() ≥ 2 and "norm" not in n) get λ; the rest gets 0.

Common bugs: (1) Adam with weight_decay, (2) decaying biases, (3) forgetting to scale λ when scaling LR for Lion, (4) decay applied around clipping in the wrong order.

LR schedules

• Linear warmup + cosine decay: LLM/ViT default. Warmup steps 0.5–4% of total (e.g., 2000 of 100k); cosine to 10% of peak. Why warmup matters for AdamW: at step 1, v̂ ~ g² (high-variance estimate) → updates explode. Warmup gives v̂ time to stabilize. Without it, transformers diverge; PostLN especially.
• OneCycle (Smith): triangular LR up to max_lr at ~30% of run, cosine down past initial; momentum runs anti-cyclic (high→low→high). "Super-convergence." Best on shorter CV runs with a fixed budget.
• WSD (Warmup-Stable-Decay) (MiniCPM, DeepSeek-style): warmup → long flat → short decay (last 10–20%). Decay can be linear, cosine, exponential, or inverse-sqrt. Lets you fork from the stable plateau with different decay lengths/data mixes — checkpoint reuse for continual training.
• Trapezoidal ≈ WSD with linear up + flat + linear down. Linear-decay-to-zero simple, slightly worse than cosine but optimum-time-agnostic. Inverse-sqrt (original Transformer): now uncommon.

LR-batch scaling rules

• Linear scaling (Goyal 2017): lr ∝ B, valid for SGD up to ~8k batch on ImageNet; needs ramp-up warmup to absorb the larger jump. Beyond that breaks → use LARS/LAMB layer-wise rules.
• Sqrt scaling (Krizhevsky): lr ∝ √B. Theoretically motivated for Adam/adaptive (variance argument); empirically often weaker than linear at moderate scale.
• μP / μTransfer (Yang & Hu, 2022): reparametrize so input/output/hidden layers each scale weights, init, and LR by width-dependent factors making feature-update magnitude O(1). Tune HPs on a 40M model, transfer zero-shot to 6.7B+. Increasingly standard in 2024–25 frontier runs (cost saving on HP search). PyPI mup. Footgun: only stable across width; depth/data still need re-tune. Works with AdamW; recipes for Muon emerging.

Gradient clipping

torch.nn.utils.clip_grad_norm_(params, max_norm=1.0) between loss.backward() and optimizer.step(). Transformers: max_norm=1.0. RNNs: 5–10 historically. With AMP+GradScaler, call scaler.unscale_(optimizer) first, then clip, then scaler.step(optimizer) — clipping a scaled gradient is a silent bug.

clip_grad_value_ (per-element) almost never preferred — kills direction information. error_if_nonfinite=True catches NaN/Inf early but raises in the hot path; many trainers prefer detecting via scaler skip events.

Mixed precision

• FP16 (AMP): range too narrow; needs torch.cuda.amp.GradScaler with dynamic loss scale. On overflow, scaler skips the step and halves scale; on N consecutive non-overflows, doubles. Skipped steps mean LR scheduler is now ahead of the optimizer — use scaler.step(optimizer) so PyTorch tracks correctly.
• BF16: same exponent range as FP32, lower mantissa. No GradScaler needed — usually. Default for LLM pretrain on A100/H100. Silent footgun: small accumulations (long-sequence loss reduction, layernorm reductions over big d_model) lose precision. Keep reductions / optimizer master-state in FP32; PyTorch AMP-bf16 + FP32 master via MixedPrecision in FSDP.
• TF32: Ampere+ matmul-only mode; 10-bit mantissa, FP32 range. Enable via torch.set_float32_matmul_precision('high') or torch.backends.cuda.matmul.allow_tf32 = True. ~Free quality cost, ~2× perf vs strict FP32.
• FP8 (Hopper+, NVIDIA Transformer Engine): E4M3 fwd, E5M2 bwd; per-tensor delayed scaling. By late 2025 FP8 is production for frontier pretrain (Meta, Microsoft, Google reported 30–40% throughput vs BF16) with no quality regression on most training. MXFP8 micro-scaling (Blackwell) reduces calibration footguns. Watch for: scale-factor staleness, FP8 attention Q/K outliers (need outlier protection).

Normalization

• BatchNorm: ε=1e-5, momentum=0.1 (running stat update); train/eval mode different. Eval mode uses running mean/var → forgetting to switch to eval mode re-estimates from current batch (catastrophic with B=1). Fails when B<8; B=2 has 10%+ ImageNet error vs full BN.
• LayerNorm: ε=1e-5, applies per-token across feature dim. Required for transformers (sequence-batch independence). PyTorch nn.LayerNorm includes affine γ,β; γ should not be weight-decayed.
• GroupNorm (Wu & He): default G=32; batch-independent; mainstay of detection/segmentation/video where per-image batches are small.
• RMSNorm (Zhang & Sennrich): drops mean centering, only x / RMS(x) · γ. Used by Llama, Mistral, Gemma, Qwen, DeepSeek. 7–64% faster than LayerNorm with no quality loss because residual connections keep activations roughly zero-mean already.
• PreLN (x + Sublayer(LN(x))) vs PostLN (LN(x + Sublayer(x))). PostLN is the original Transformer; gradient norms grow with depth → long warmup mandatory and divergence common past ~12 layers. PreLN is gradient-stable but suffers activation variance growth with depth → "massive activations" in 2B+ models, FP16 attention overflow. Peri-LN (LN before and after the sublayer) used in Gemma-2 / OLMo-class to bound both. NormFormer adds extra LNs after attention QK and FFN GeLU.

Gradient explosion / NaN diagnosis

torch.autograd.set_detect_anomaly(True) finds the op producing NaN (slow — debug only). Hook register_full_backward_hook on suspect modules and dump grad norms.

Track per-step: total grad norm, per-layer grad norm, max activation, attention logit max, loss-scale value.

Spike templates:

• Loss flat → spike → recover — data hot batch.
• Progressively rising grad norm before spike — PostLN drift, LR too high.
• Sudden NaN with previous skipped scale — FP16 underflow accumulating.
• Rare spike in BF16 attention with FP16-cast logits — overflow in attn softmax. Pre-softmax clamp or convert QKᵀ to FP32.

Divide-by-zero spots: softmax with −∞ rows (full-mask token), logsumexp across empty groups, RMSNorm with zero input vector, / std in custom norm.

Augmentation

• Image: RandAugment (n=2, m=9 typical), AutoAugment policies, RandomErasing.
• Mixup: x' = λ x_a + (1−λ) x_b, y' = λ y_a + (1−λ) y_b, λ ~ Beta(α,α), α=0.2 (CV) to 0.8 (ViT).
• CutMix: pastes a rectangle from b into a, λ = pasted-area / total → label is interpolated by area.
• Train-only: never apply to val/test.
• Text: synonym swap, back-translation, span deletion. Be careful with classification labels surviving the swap.
• Audio: SpecAugment — time/frequency masks plus time warp on spectrogram.
• Mixup label math: with cross-entropy and one-hot, expand to loss = λ·CE(logits, y_a) + (1−λ)·CE(logits, y_b) (don't mix tensor and call cross_entropy — lose precision).

Label smoothing

y_smooth = (1−ε)·one_hot + ε/K with ε=0.1 typical. Improves calibration (ECE lower), slightly hurts top-1 by ~0.1–0.3% — accepted tradeoff in production. Compose with Mixup carefully (don't double-smooth: usually drop LS when Mixup is on, or use nn.CrossEntropyLoss(label_smoothing=0.1) only).

SWA / EMA

• SWA (Izmailov 2018): cycle LR high+constant during last 25% of training, average snapshots, recompute BN running stats post-hoc with one data pass (update_bn). Wider basin → better generalization. PyTorch torch.optim.swa_utils.AveragedModel.
• EMA: throughout training, θ_ema ← β θ_ema + (1−β) θ, β=0.999 (image classif), 0.9999 (large LM, diffusion). Stable Diffusion, Imagen ship EMA weights, not live ones. EMA ≈ low-pass filter; reduces noise from late high-LR updates. Cheaper than SWA; near-identical quality.

Stochastic Depth / DropPath

Per-block Bernoulli drop with linear schedule p_l = 1 − (l/L)(1 − p_L), p_L = 0.5 (deep ResNet) to p_L = 0.9 (small ViT). At inference, scale residual by survival probability. Critical for ViT-Large, ConvNeXt; absent or tiny in ViT-Base. timm drop_path_rate arg.

Memory-efficient training

• Activation checkpointing (torch.utils.checkpoint): saves ~70% activation memory, ~25% slower. Selective AC: checkpoint attn but not MLP (or vice-versa).
• ZeRO (DeepSpeed) / FSDP (PyTorch): ZeRO-1 shards optimizer state (Muon needs Distributed Muon — vanilla ZeRO-1 incompatible because orthogonalization needs the full matrix), ZeRO-2 + grads, ZeRO-3 + params. Offload to CPU/NVMe if VRAM-bound but PCIe-fed.
• FlashAttention / SDPA: PyTorch 2.2+ ships FlashAttention-2 inside F.scaled_dot_product_attention; 2.5+ adds CuDNN SDPA tuned for Hopper. torch.nn.attention.sdpa_kernel to force a backend; math kernel is the slow fallback.
• Sequence packing: concatenate short docs into max-length sequences with per-doc attention masks (block-diagonal). 2–5× tokens-per-step on web pretrain; needs cross-doc-attention masking to avoid leakage.

Loss-spike forensics

Causes ranked by frequency: (1) data hot batch (unusual token co-occurrence; logit saturation), (2) LR too high after warmup, (3) FP16 underflow → scaler-skip → optimizer/state desync, (4) optimizer-state corruption (NaN propagated into v̂), (5) data poisoning / undetected bad doc.

Skip-and-rewind protocol (PaLM-style): if loss > k · rolling_mean, restore from N-step-old checkpoint, skip the last N batches, replay. Common k=1.5, N=200. Cheap insurance; standard in 2024+ pretrain stacks. Pair with batch logging for forensic root-cause.

Recent changes (late 2025 / 2026)

• Muon mainstreamed for LLM pretrain matrix params (Distributed Muon for ZeRO-1 cluster training).
• μP standard for HP transfer in frontier runs; cuts HP search costs > 90%.
• FP8 is production for BF16-class quality with ~30–40% throughput; MXFP8 (Blackwell) eases calibration.
• Schedule-Free AdamW displacing cosine where horizon is unknown.
• WSD preferred for continual pretrain / data-mix forks (DeepSeek, MiniCPM lineage).
• Peri-LN seen alongside PreLN at >1B scale to control activation variance.

Authoritative references

• AdamW
• Lion (Symbolic Discovery)
• Sophia
• Muon (Jordan), Distributed Muon
• Schedule-Free
• μP (Tensor Programs V), microsoft/mup
• Fixup init
• Group Normalization
• RMSNorm
• Linear scaling rule (1-Hour ImageNet)
• WSD / MiniCPM
• PyTorch SDPA / FlashAttention-2
• NVIDIA Transformer Engine (FP8)
• Spike No More (LLM stability)
• Peri-LN

Guardrails

Before recommending a non-trivial training change (optimizer swap, schedule change, weight-decay rewiring, FP8 enable, normalization swap):

1. Quote the parameter and any standard default (β, ε, λ, warmup %)
2. Cite the originating paper or PyTorch doc
3. Make the recommendation conditional on observed evidence (loss-curve, grad-norm trace, scaler skip rate, validation delta) — never blanket-tune
4. Verify hardware/library generation. Many features gate (FP8 on Hopper+, MXFP8 on Blackwell, FlashAttention-2 in PT 2.2+)

Tuning without measurement is worse than defaults.