05 — Training Dynamics

Phase 2 · Months 5–6 | Estimated time: 10–14 hours across 2 weeks

Training dynamics is the discipline of understanding why a model converges, plateaus, diverges, or silently produces garbage. Most practitioners know how to call loss.backward(). Fewer know how to diagnose why their training loss hit 0.4 and stopped, or why their model performs well in the training loop but produces nonsense at inference. This module covers the mechanisms behind the tools — because tools misconfigured at the wrong moment don’t fail loudly. They fail quietly.

---

1. Learning Rate Schedules

The learning rate is the single most impactful hyperparameter. Too high: divergence or oscillation. Too low: slow convergence or plateau in a suboptimal local minimum. A fixed learning rate is almost never the right answer — the optimal rate changes as training progresses.

The following table covers the main strategies used in practice as of 2025.

Schedule	Behavior	When to Use
Constant	Fixed LR throughout	Debugging only; almost never optimal in production
Step decay	Multiply LR by γ every N epochs	Old default; arbitrary step boundaries are a hyperparameter you don’t want to tune
Cosine annealing	Smooth decay from LR_max to LR_min following a cosine curve	Standard for CNNs and anything without a warmup requirement
Warmup + cosine	Linear increase for N steps, then cosine decay	Standard for transformers in 2025; prevents early instability from random init
1-cycle policy	LR ramps up then decays; momentum inversely follows	fast.ai’s contribution; enables training at higher peak LRs; faster convergence

Why cosine annealing works: The cosine curve provides smooth exploration early (LR still high) and smooth exploitation late (LR has decayed). Unlike step decay, there are no abrupt LR drops that can destabilize training. The smoothness means the loss landscape is being navigated without sudden jumps.

import torch
import torch.optim as optim

model    = ...  # your model
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# Cosine annealing: decays LR from lr to eta_min over T_max epochs
scheduler = optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=50,       # number of epochs to complete one cycle
    eta_min=1e-6    # minimum LR at the bottom of the cosine
)

for epoch in range(50):
    train_one_epoch(model, ...)
    scheduler.step()
    print(f"Epoch {epoch}: LR = {scheduler.get_last_lr()[0]:.2e}")

# ── Warmup + Cosine (for transformers) ─────────────────────────────────────
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR

warmup    = LinearLR(optimizer, start_factor=0.1, end_factor=1.0, total_iters=10)
cosine    = CosineAnnealingLR(optimizer, T_max=90, eta_min=1e-6)
scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[10])

---

2. Weight Initialization

Initialization determines the starting point of optimization and has a larger impact than most practitioners expect. A bad initialization can cause activations to explode or vanish before training even begins — meaning the first forward pass already puts you in an unrecoverable gradient regime.

Why initialization matters — symmetry breaking: If all weights are initialized to the same value, every neuron in a layer computes the same gradient and updates identically. All neurons remain identical throughout training — the network has effectively one neuron per layer. Random initialization breaks this symmetry.

Why scale matters: Consider a linear layer with fan_in = 512. If weights are initialized from N(0, 1), the output variance is 512 × 1 = 512. After a few layers, activations are O(512^L/2) where L is depth — they explode. The initialization schemes below solve this by scaling variance to keep activations ~unit variance throughout the network.

Scheme	Variance Formula	Activation	Status
Xavier / Glorot	2 / (fan_in + fan_out)	tanh, sigmoid	Standard for older architectures
Kaiming / He	2 / fan_in	ReLU	Default for modern CNNs and MLPs
LeCun	1 / fan_in	SELU	Used with self-normalizing networks

import torch.nn as nn

# PyTorch defaults are already Kaiming uniform for Linear and Conv layers.
# Explicit initialization:
layer = nn.Linear(512, 256)

nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')
nn.init.zeros_(layer.bias)

# Xavier for tanh layers:
nn.init.xavier_uniform_(layer.weight)

# What happens with constant init (demonstrates the problem):
nn.init.constant_(layer.weight, 0.0)  # all neurons identical — symmetry NOT broken

What bad initialization looks like: Run a forward pass and print layer.output.std() at each layer. If it’s 0.01 (vanishing) or 100 (exploding) before any training has occurred, your initialization is wrong.

---

3. BatchNorm — And the Most Common PyTorch Bug

Batch Normalization normalizes the output of a layer across the mini-batch dimension, then applies learned scale γ and shift β to allow the network to undo the normalization if needed.

Mathematical definition for a mini-batch of size B:

μ_B = (1/B) Σ x_i                         # batch mean
σ²_B = (1/B) Σ (x_i - μ_B)²               # batch variance
x̂_i = (x_i - μ_B) / sqrt(σ²_B + ε)       # normalize
y_i = γ · x̂_i + β                         # scale and shift (γ, β are learned)

γ and β are learnable parameters initialized to 1 and 0. The network can learn γ = σ and β = μ to effectively pass through unchanged — BatchNorm can learn to do nothing if that’s optimal.

Train mode vs eval mode — the most common PyTorch bug:

During training, BatchNorm uses the current mini-batch statistics (μ_B, σ²_B). It also maintains exponential moving averages of these statistics as running_mean and running_var.

During evaluation (inference), BatchNorm must use the running statistics — not the batch statistics — because:

1. You might be predicting on a single sample (batch size = 1, making batch statistics meaningless)

2. You want deterministic, reproducible predictions independent of what other samples happen to be in the batch

If you forget model.eval() at inference, BatchNorm uses batch statistics → wildly wrong predictions on small batches or single samples.

import torch
import torch.nn as nn

# ── Demonstrating the BatchNorm train/eval bug ──────────────────────────────
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.BatchNorm1d(64),
    nn.ReLU(),
    nn.Linear(64, 1),
)

x_single = torch.randn(1, 10)   # single sample at inference

# ❌ WRONG: model is in train mode, BatchNorm uses batch stats
# With B=1, variance = 0, division by near-zero → NaN or garbage
model.train()
out_wrong = model(x_single)
print(f"Train mode output (B=1): {out_wrong.item():.4f}")  # often NaN or extreme value

# ✅ CORRECT: model.eval() switches BatchNorm to use running stats
model.eval()
with torch.no_grad():
    out_correct = model(x_single)
print(f"Eval mode output  (B=1): {out_correct.item():.4f}")  # stable, reproducible

# ── Verifying running stats are populated ──────────────────────────────────
# Train on some data first, then evaluate:
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model.train()
for _ in range(100):
    x_batch = torch.randn(32, 10)
    loss = model(x_batch).mean()
    optimizer.zero_grad(); loss.backward(); optimizer.step()

# Now eval mode uses the running stats accumulated over 100 batches
model.eval()
with torch.no_grad():
    out = model(x_single)
print(f"After training, eval mode (B=1): {out.item():.4f}")  # correct

The three-line fix you must memorize:

model.eval()          # switch to eval mode before any inference
with torch.no_grad(): # disable gradient computation (also saves memory)
    output = model(input)

---

4. Dropout

Dropout randomly zeros a fraction p of neurons during each forward pass in training. Each neuron is zeroed independently with probability p.

Mathematical definition: For a neuron with output x, dropout produces:

x_dropped = x · Bernoulli(1 - p) / (1 - p)

The / (1 - p) scaling keeps the expected value of x_dropped equal to x — this is inverted dropout, which is what PyTorch implements.

Why it works — ensemble interpretation: With n neurons, dropout samples a different sub-network at each training step. There are 2^n possible sub-networks. Training with dropout is equivalent to training and averaging over an exponentially large ensemble — at inference, all neurons are active (scaled appropriately), approximating the ensemble mean.

import torch.nn as nn

# Dropout for classic MLPs: p=0.5 is standard
dropout_heavy = nn.Dropout(p=0.5)

# Dropout for transformers: p=0.1 is standard (attention already regularizes)
dropout_light = nn.Dropout(p=0.1)

# CRITICAL: Dropout is automatically disabled in eval mode
model.train()  # Dropout active
model.eval()   # Dropout disabled — all neurons contribute at full scale

p values by context:

• Classic MLP, dense layers: p=0.3–0.5

• Transformers (attention, feedforward): p=0.1

• CNNs: often p=0.0–0.25 (BatchNorm provides regularization already; too much dropout hurts)

---

5. Gradient Clipping

Gradient clipping caps the L2 norm of all gradients before the optimizer step. It prevents exploding gradients — particularly relevant in RNNs (BPTT), deep transformers during unstable phases, and any model trained with high learning rates.

import torch

# Standard pattern — call AFTER loss.backward(), BEFORE optimizer.step()
loss.backward()

# Clip the global gradient norm to max_norm=1.0
# Returns the norm before clipping — useful for monitoring
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

optimizer.step()

# Log grad_norm to detect instability:
# Consistently high (>10): LR too high or bad initialization
# Suddenly spikes then drops: gradient explosion that clipping caught
# Consistently near 0: vanishing gradients, model isn't learning
print(f"Gradient norm: {grad_norm:.4f}")

When to clip: Always for RNNs/LSTMs. For transformers, standard practice is max_norm=1.0. For CNNs, only if you observe NaN losses or loss spikes. The overhead is negligible — it’s a single norm computation.

---

6. Mixed Precision Training

Modern GPUs have dedicated hardware for FP16 and BF16 operations (Tensor Cores on NVIDIA). Mixed precision trains in lower precision for forward/backward passes and accumulates gradients in FP32, reducing VRAM by ~40–50% and increasing throughput.

FP16 vs BF16:

Format	Exponent bits	Mantissa bits	Range	Notes
FP32	8	23	±3.4×10³⁸	Full precision baseline
FP16	5	10	±65,504	Overflow risk with large activations
BF16	8	7	±3.4×10³⁸	Same range as FP32; preferred for LLMs

Why BF16 is preferred for LLMs: BF16 has the same 8-bit exponent as FP32, so it can represent the same dynamic range without overflow. FP16’s 5-bit exponent caps at ~65,504 — large activations overflow to Inf, causing NaN losses. On NVIDIA Ampere/Hopper and newer GPUs, BF16 Tensor Cores are available and BF16 is the default choice.

import torch
from torch.cuda.amp import autocast, GradScaler

model     = YourModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# GradScaler prevents FP16 underflow by scaling the loss up before backward
# (Only needed for FP16; BF16 doesn't require scaling)
scaler = GradScaler()   # use with FP16; omit or set enabled=False for BF16

for batch in dataloader:
    inputs, targets = batch

    # autocast selects the lower-precision dtype automatically
    with autocast(device_type='cuda', dtype=torch.bfloat16):  # or torch.float16
        outputs = model(inputs)
        loss    = criterion(outputs, targets)

    # FP16 path: scale loss to prevent gradient underflow
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()

    # BF16 path (simpler — no scaler needed):
    # loss.backward()
    # torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    # optimizer.step()
    # optimizer.zero_grad()

Memory savings in practice: A ResNet-50 in FP32 uses ~97MB for weights. In BF16: ~48MB. Activations and gradients scale similarly. For a 7B parameter LLM: FP32 ≈ 28GB, BF16 ≈ 14GB — the difference between fitting on one A100 and needing two.

---

7. Debugging Training: A Diagnostic Checklist

Training failures have a finite set of root causes. Most of them are identifiable from the loss curve and gradient norm alone if you know what to look for.

Before debugging, instrument your training loop with these three metrics: train_loss, val_loss, and grad_norm. Without them, you are debugging blind.

Symptom	Root Causes	Diagnostic	Fix
Loss → NaN	LR too high; log(0) in loss; bad init; exploding gradients	Check grad_norm spike; print logits.min() for log inputs	Reduce LR 10×; add gradient clipping; check for log(0+ε)
Loss plateaus early	LR too low; dying ReLUs; wrong architecture	Check if gradients near 0; inspect activation histograms	Increase LR; switch to LeakyReLU; add BatchNorm
Loss oscillates	LR too high; batch size too small	Loss curve looks like high-frequency noise	Reduce LR; increase batch size; add gradient clipping
Train↓ Val↑ (overfitting)	Model too large; insufficient data; no regularization	Large train/val gap by epoch 10	Add Dropout; augmentation; weight decay; early stopping
Train↓ Val stays flat	Data distribution mismatch; val set is harder	Consistent gap from epoch 1	Check val preprocessing matches train; check class imbalance
Model learns nothing	Wrong loss function; wrong labels; frozen layers accidentally	Loss = log(num_classes) and stays there	Verify labels; check requires_grad; test with tiny batch overfit

The single most useful debugging technique: Before training the full dataset, try to overfit a single batch of 32 samples to near-zero loss. If you can’t do that, there’s a bug in your model or training loop — not a data problem, not a hyperparameter problem. Fix the bug first.

# ── Single-Batch Overfit Test ───────────────────────────────────────────────
model.train()
x_debug, y_debug = next(iter(train_loader))
x_debug, y_debug = x_debug.to(DEVICE), y_debug.to(DEVICE)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)  # high LR to overfit fast

for step in range(200):
    optimizer.zero_grad()
    logits = model(x_debug)
    loss   = criterion(logits, y_debug)
    loss.backward()
    optimizer.step()
    if step % 20 == 0:
        print(f"Step {step:3d} | loss={loss.item():.6f}")

# Expected: loss should reach <0.001 within 200 steps
# If it doesn't: model architecture bug, loss function mismatch, or frozen layers

---

What Most People Get Wrong

The BatchNorm train() vs eval() distinction. This single bug has burned more debugging sessions than almost any other PyTorch pitfall.

The mechanism: model.train() sets BatchNorm (and Dropout) to training behavior. model.eval() switches BatchNorm to use its running statistics and disables Dropout. If you run inference in train() mode, every prediction depends on the other samples in the batch — your model’s output for sample A changes if you change sample B. On batch size 1, variance = 0, and the normalization collapses.

The fix is two lines: model.eval() and with torch.no_grad(). The fact that this is two lines that must both be remembered is why this bug persists. Write a helper:

@torch.no_grad()
def predict(model, x):
    model.eval()
    return model(x)

A secondary subtlety: model.eval() only affects BatchNorm and Dropout. It does not disable gradient computation — that’s what torch.no_grad() does. You need both for correct, memory-efficient inference.

---

Return to README.md · Next: 06_phase_projects.md