Training in
Deep Learning

Training Loop — End-to-End Mental Model

What actually happens in one training step, tensor by tensor.

A training step is not magic. It is a precise sequence of tensor operations. Every step: data enters, activations flow forward, loss is computed, gradients flow backward, weights change. Repeat ~millions of times.

What Lives in Memory During a Step

Forward Pass Deep Dive

What actually happens layer by layer — and why activations eat your memory.

The forward pass is deterministic given weights and inputs. It is a composed function: output = f_L(f_{L-1}(...f_1(input))). Each f_i is a layer. Each intermediate result is called an activation.

Linear Layers

Most fundamental operation. Given input x: [B, in], weight W: [out, in], bias b: [out]:

During backprop, you need: the input x to compute ∂L/∂W = xᵀ ∂L/∂y. This is why x must be stored.

Attention (Transformer)

Self-attention is the expensive operation. For each layer:

Memory cost: The attention matrix QKᵀ is [B, H, S, S] — quadratic in sequence length S. For S=8192: that's 8192² = 67M floats per layer per head. FlashAttention avoids materializing this by fusing the kernel.

Why Activations Are Stored

Loss Functions

Choosing the right error signal is as important as the architecture.

Logits vs Probabilities

Logits = raw unnormalized scores from the model's last linear layer. They can be any real number. Probabilities = after softmax, sum to 1.

Backpropagation

The engine of learning — how error signals flow backward through a network.

Chain Rule Intuition

Backprop is just the chain rule applied to composed functions. If L = f(g(h(x))), then:

In a neural network: loss depends on outputs which depend on layer n-1 which depends on layer n-2... The gradient of the loss with respect to each weight is a product of local gradients chained together.

Layer-by-Layer Gradient Computation

Vanishing & Exploding Gradients

The gradient at layer k involves a product of N weight matrices and activation derivatives. This product can shrink to near-zero (vanishing) or blow up (exploding).

Residual Connections are the Key Insight

Optimizers

How gradients get translated into weight updates — and the hidden state they carry.

SGD (Stochastic Gradient Descent)

The simplest update rule. One gradient step in the negative gradient direction. Problem: sensitive to gradient noise, oscillates in narrow valleys. When to use: CNNs with careful tuning, when you have very large batches with low gradient noise. SGD with momentum is more commonly used.

SGD + Momentum

Accumulates a velocity vector in the direction of consistent gradients. Dampens oscillations. β=0.9 is standard — means each step uses 90% of previous velocity + 10% new gradient. Physical intuition: ball rolling downhill with inertia.

Adam (Adaptive Moment Estimation)

Per-parameter adaptive learning rates. Dividing by √v̂ normalizes the step size by the gradient's recent variance. Effect: parameters with noisy gradients take smaller steps; parameters with consistent gradients take larger steps. Defaults: β₁=0.9, β₂=0.999, ε=1e-8.

AdamW

Optimizer State Memory

Learning Rate & Schedulers

The single most important hyperparameter. Get this wrong and nothing else matters.

Why LR Is #1

Learning rate controls how large a step you take in the loss landscape. Too high: overshoot minima, loss diverges or oscillates. Too low: trains too slowly, gets stuck in poor local minima. There is a "Goldilocks zone" — and it changes during training.

LR Warmup

Start with very low LR, linearly increase to target LR over N steps. Why: at training start, parameters are random, gradients are large and noisy. High LR on random init → divergence. Warmup lets the model stabilize before taking aggressive steps.

Rule of thumb: 1–5% of total steps as warmup. For LLMs: 1000–2000 warmup steps common.

Schedulers

Cosine Schedule with Warmup (Modern Default)

def get_lr(step, warmup_steps, max_steps, max_lr, min_lr=1e-5):
    # Linear warmup
    if step < warmup_steps:
        return max_lr * step / warmup_steps
    # Cosine decay
    progress = (step - warmup_steps) / (max_steps - warmup_steps)
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + cos(math.pi * progress))

Gradient Mechanics

Controlling how gradients flow — the difference between a stable run and NaN hell.

Gradient Accumulation

When your GPU can't fit a large batch, split it into micro-batches and accumulate gradients before updating weights:

for i, (x, y) in enumerate(loader):
    loss = model(x, y) / accum_steps   # divide to normalize
    loss.backward()                         # grads accumulate in .grad
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Effective batch size = micro_batch_size × accum_steps × n_gpus. Behaves identically to a large batch if you also scale LR proportionally.

Gradient Clipping

Cap the gradient norm to prevent exploding gradients from destroying your training:

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Activations

Non-linearity is what makes neural networks universal approximators.

Without activation functions, stacking linear layers gives you a single linear transformation — no matter how many layers. Activations introduce non-linearity, enabling the network to learn complex functions.

Why GELU Won in Transformers

GELU = Gaussian Error Linear Unit. Unlike ReLU's hard gate (0 or pass), GELU has a smooth probabilistic gate: inputs are weighted by how likely they are to be positive under a Gaussian. This smooth gradient behavior empirically outperforms ReLU in transformers. LLaMA uses SiLU (nearly identical) in SwiGLU: SiLU(W₁x) ⊙ W₃x.

Precision & Training Efficiency

The difference between FP32 and BF16 is not just speed — it fundamentally changes what you can train.

FP16 vs BF16 — The Critical Distinction

Mixed Precision Training

Store model weights in FP16/BF16 for speed, but compute loss and keep master copy of weights in FP32 for precision during optimizer step.

Loss Scaling (for FP16)

Multiply loss by a large scalar S before backward → gradients are S× larger → less underflow. After backward, divide gradients by S before optimizer step. PyTorch GradScaler does this automatically with dynamic scaling — if overflow detected, it halves S and skips the step.

scaler = GradScaler()

with autocast('cuda'):
    output = model(x)
    loss = criterion(output, y)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

Memory Optimization Techniques

How to train models that don't fit in GPU memory — the trade-offs that define production LLM training.

Gradient Checkpointing Deep Dive

from torch.utils.checkpoint import checkpoint

output = checkpoint(block, x)  # recomputes block during backward

ZeRO Optimizer (DeepSpeed)

ZeRO = Zero Redundancy Optimizer. In standard DDP, every GPU holds a full copy of optimizer states. For 7B param model with Adam: 7B×4 (FP32 weights) + 7B×4 (m) + 7B×4 (v) = 84GB just for optimizer. ZeRO shards this across N GPUs. With 8 GPUs, that's 84/8 = 10.5GB per GPU.

Batch Size & Scaling

Batch size is not just a throughput knob — it fundamentally changes training dynamics.

What Batch Size Controls

Each gradient update is computed on a sample of the data. Larger batches = more accurate gradient estimate (less noise), but also fewer updates per epoch. The gradient noise from small batches acts as a regularizer and helps escape sharp minima.

Linear Scaling Rule

Effective Batch Size

If micro_batch=4, accum_steps=8, n_gpus=8: effective_batch = 4×8×8 = 256. This is what matters for training dynamics. You can trade any of these three factors against each other based on hardware constraints.

Regularization

Overfitting detection and prevention — keeping your model generalizable.

Detecting Overfitting

The signature: training loss continues to decrease but validation loss stagnates or increases. Monitor both curves throughout training. The gap between train and val loss = degree of overfitting.

Dropout in Transformers

Modern LLMs (LLaMA, GPT-4) often use very low or zero dropout during pre-training (large dataset + short training = less overfitting risk). Dropout rates of 0.0–0.1 are common. Fine-tuning on small datasets: increase to 0.1–0.3.

Weight Decay Intuition

Without weight decay, nothing prevents weights from growing arbitrarily large (as long as loss decreases). Large weights make the model sensitive to small input perturbations → poor generalization. Weight decay pulls weights back toward zero at each step, keeping the model in a well-behaved regime.

Training Stability & Debugging

The practical guide to why training broke — and how to fix it.

Debugging Checklist

Distributed Training

Scaling training across multiple GPUs and machines — the architecture of modern LLM training.

Data Parallelism (DDP) — The Standard

Each GPU has a full model copy. Each GPU processes a subset of the batch. After backward, gradients are all-reduced (summed + divided) across GPUs — every GPU ends up with the same average gradient. Effectively multiplies batch size by N GPUs.

# PyTorch DDP setup
from torch.nn.parallel import DistributedDataParallel as DDP
model = DDP(model.to(rank), device_ids=[rank])
# All-reduce happens automatically in .backward()

Communication Overhead

The bottleneck in multi-node training is often not compute — it's gradient synchronization. NVLink (within a node): 600 GB/s. InfiniBand (across nodes): 200–400 Gb/s. For a 7B model, syncing gradients (28GB @ FP16) = ~0.1s over IB. At 10-step gradient accumulation, this is amortized. Overlap compute and communication wherever possible.

Training LLMs vs CNNs vs VLMs

Key differences in training dynamics across model families.

Sequence Length Impact in LLMs

Going from S=2048 to S=8192: attention memory scales as O(S²) → 16× more attention memory. KV cache during inference = 2 × n_layers × n_heads × d_head × S × dtype_bytes. For LLaMA-7B at S=4096 in FP16: ~2GB KV cache per request. This is why long-context is expensive.

VLM Training Specifics

VLMs (e.g., LLaVA, Flamingo, CLIP) typically use contrastive loss. Key challenge: vision and text encoders are pre-trained separately and must be aligned. Training is often staged: (1) freeze both encoders, train projection layer; (2) unfreeze text encoder, joint fine-tuning; (3) optionally unfreeze vision encoder. Large negatives in contrastive learning are critical — CLIP used batch sizes of 32,768.

Practical Training Pipeline

Production-ready PyTorch training loop with all the essentials.

import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler
from torch.utils.data import DataLoader

# ─── Setup ───────────────────────────────────────────────────
device = torch.device("cuda")
model = MyModel().to(device)
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
    weight_decay=0.1,
    betas=(0.9, 0.95),    # β2=0.95 is better for LLMs
)
scaler = GradScaler()   # only if using FP16 (not needed for BF16)
criterion = nn.CrossEntropyLoss()
loader = DataLoader(dataset, batch_size=32, num_workers=4,
                    pin_memory=True, persistent_workers=True)

ACCUM_STEPS = 4
MAX_GRAD_NORM = 1.0

# ─── Training Loop ───────────────────────────────────────────
model.train()
optimizer.zero_grad(set_to_none=True)

for step, (x, y) in enumerate(loader):
    x, y = x.to(device), y.to(device)

    # Forward pass with mixed precision
    with autocast("cuda", dtype=torch.bfloat16):
        logits = model(x)
        loss = criterion(logits, y) / ACCUM_STEPS

    # Backward
    scaler.scale(loss).backward()

    if (step + 1) % ACCUM_STEPS == 0:
        # Unscale before clipping!
        scaler.unscale_(optimizer)
        grad_norm = torch.nn.utils.clip_grad_norm_(
            model.parameters(), MAX_GRAD_NORM
        )

        scaler.step(optimizer)
        scaler.update()
        scheduler.step()
        optimizer.zero_grad(set_to_none=True)

        # Logging
        print(f"step={step}, loss={loss.item()*ACCUM_STEPS:.4f}, "
              f"grad_norm={grad_norm:.3f}, lr={scheduler.get_last_lr()[0]:.6f}")

DataLoader Best Practices

loader = DataLoader(
    dataset,
    batch_size=128,
    shuffle=True,
    num_workers=8,           # parallel data loading
    pin_memory=True,         # faster CPU→GPU transfer
    persistent_workers=True,# avoid worker re-init overhead
    prefetch_factor=2,       # pre-load batches ahead
    drop_last=True,          # avoid partial batch issues
)

Engineering Tips

High-value production insights — what separates a struggling run from a clean one.

Speed-Up Tricks

Mental Models & Intuition

The way to think about training that makes everything else click.

More Key Intuitions

Code Snippets Reference

Drop-in code for the most common training operations.

Mixed Precision (BF16 — Modern Best Practice)

import torch

# BF16 — no loss scaling needed (A100+, H100)
with torch.autocast("cuda", dtype=torch.bfloat16):
    output = model(x)
    loss = criterion(output, y)
loss.backward()
optimizer.step()

Gradient Clipping

loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(
    model.parameters(),
    max_norm=1.0   # standard for LLMs; try 0.5–2.0
)
optimizer.step()
# grad_norm tells you the norm BEFORE clipping

AdamW Optimizer Setup

# Standard LLM setup
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,
    betas=(0.9, 0.95),    # β2=0.95 for LLMs (faster adaptation)
    eps=1e-8,
    weight_decay=0.1,
    fused=True,            # fused CUDA kernel — 20% faster
)

# Or: parameter groups (no weight decay on embeddings/norms)
no_decay = ["bias", "LayerNorm.weight", "embedding"]
params = [
    {"params": [p for n,p in model.named_parameters()
               if not any(nd in n for nd in no_decay)], "weight_decay": 0.1},
    {"params": [p for n,p in model.named_parameters()
               if any(nd in n for nd in no_decay)], "weight_decay": 0.0},
]
optimizer = torch.optim.AdamW(params, lr=3e-4)

Gradient Checkpointing

from torch.utils.checkpoint import checkpoint, checkpoint_sequential

# Option 1: wrap individual expensive blocks
class TransformerBlock(nn.Module):
    def forward(self, x):
        return checkpoint(self._forward, x, use_reentrant=False)

# Option 2: enable via HuggingFace
model.gradient_checkpointing_enable()

# Option 3: FSDP with activation checkpointing
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

Checkpoint Saving/Loading

# Save
torch.save({
    "step": step,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "scaler": scaler.state_dict(),   # FP16 only
    "loss": loss.item(),
}, f"ckpt_step_{step}.pt")

# Load
ckpt = torch.load("ckpt_step_1000.pt", map_location=device)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
start_step = ckpt["step"]