Gradient Flow

Chapter 1: The Chain Rule is Multiplication

Chapter 0 showed the symptom. Now let's understand the cause. Backpropagation computes gradients using the chain rule: to find how a change in layer 1's weights affects the loss, you multiply the local derivatives of every layer in between.

For a network with L layers, the gradient at layer k is:

Each factor (∂a_i+1 / ∂a_i) is the Jacobian of layer i — how much the output changes per unit change in the input. For a simple linear layer y = Wx + b followed by ReLU, the Jacobian has two parts: the weight matrix W and the derivative of ReLU (which is 1 for positive inputs, 0 for negative).

Hand Calculation: 3-Layer Network

Let's trace the gradient through a tiny network with 3 layers. Each layer is y = σ(Wx) where σ is the sigmoid function.

Layer 3 (output): Gradient from loss = 1.0 (starting point).
Layer 2: Multiply by σ'(z₃) · W₃. Sigmoid's maximum derivative is 0.25 (at z=0). With a weight magnitude of ~0.5, the factor is 0.25 × 0.5 = 0.125.
Layer 1: Multiply again: 0.125 × 0.125 = 0.0156.

After just 3 layers with sigmoid, the gradient is already 64× smaller. After 10 layers: 0.125¹⁰ ≈ 10^-9. The gradient is effectively dead. This is why sigmoid networks couldn't be trained beyond ~5 layers — and why the invention of ReLU (whose derivative is 1, not 0.25) was transformative.

The Simulation

See exactly how the chain rule multiplies derivatives layer by layer. Pick the activation function and watch the gradient shrink or survive.

Chapter 2: Gradient Clipping

The telephone game showed that gradients can explode. In practice, a single bad batch — an outlier example, a corrupted sample, or an unlucky combination — can produce a gradient 100× or 1000× normal magnitude. The optimizer applies that enormous gradient as a weight update, the loss spikes, and it can take days of training to recover. In some cases, training never recovers.

Gradient clipping is the safety valve. Before the optimizer step, you check the gradient's magnitude and cap it at a maximum threshold. If the gradient is within bounds, nothing changes. If it's too large, you scale it down to the threshold.

Two Types of Clipping

Norm clipping (the standard approach): compute the global norm of all gradients — that's √(∑ g_i²) across every parameter in the entire model. If the global norm exceeds a threshold (typically 1.0), scale every gradient by (threshold / norm). This preserves the direction of the gradient while capping its magnitude.

Value clipping (less common): clamp each gradient element independently to [-threshold, +threshold]. This changes the direction of the gradient vector, which is usually undesirable. Norm clipping is preferred for this reason.

Hand Calculation: Norm Clipping

A small model has 3 parameter tensors with gradients [3.0, 4.0] and [0.0]. The global norm is √(3² + 4² + 0²) = √25 = 5.0.

With max_norm = 1.0: since 5.0 > 1.0, we scale by 1.0 / 5.0 = 0.2. The clipped gradients become [0.6, 0.8] and [0.0]. Same direction, magnitude capped at 1.0.

With max_norm = 10.0: since 5.0 < 10.0, min(1, 10/5) = min(1, 2) = 1. No clipping occurs. The gradients pass through unchanged.

The Simulation

Visualize norm clipping as a circle. Any gradient vector landing outside the circle gets pulled back to the boundary.

Code: Gradient Clipping in PyTorch

python
import torch

model = MyModel()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

for batch in dataloader:
    loss = model(batch)
    loss.backward()

    # Norm clipping — the standard approach for LLMs
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    optimizer.step()
    optimizer.zero_grad()

python
# Manual implementation — see what clip_grad_norm_ does
def clip_grad_norm(parameters, max_norm):
    total_norm = 0.0
    for p in parameters:
        if p.grad is not None:
            total_norm += p.grad.data.norm(2).item() ** 2
    total_norm = total_norm ** 0.5

    clip_coef = max_norm / (total_norm + 1e-6)
    if clip_coef < 1.0:
        for p in parameters:
            if p.grad is not None:
                p.grad.data.mul_(clip_coef)
    return total_norm

Chapter 3: Gradient Accumulation

Large language models need large batch sizes to train stably. GPT-3 used an effective batch size of 3.2 million tokens. LLaMA-2 used 4 million tokens. A single H100 GPU with 80 GB of memory can barely fit a batch of 4 sequences of length 2048 — that's about 8,000 tokens. How do you bridge the gap from 8,000 to 4,000,000?

Gradient accumulation. Instead of updating weights after every forward-backward pass, you run N micro-batches, accumulate (sum) the gradients, and then do ONE optimizer step using the total. Mathematically, this is identical to running one giant batch of size N × micro_batch_size — because gradients are linear in the loss, and the loss is typically averaged over the batch.

How It Works

Hand Calculation: Why It's Mathematically Identical

True big batch (B=8): Loss = (1/8) ∑_i=1..8 L_i. Gradient = (1/8) ∑ ∇L_i.

Accumulated (2 micro-batches of 4):

• Micro-batch 1: g₁ = (1/4) ∑_i=1..4 ∇L_i
• Micro-batch 2: g₂ = (1/4) ∑_i=5..8 ∇L_i
• Accumulated: g = g₁ + g₂ = (1/4)(∑_1..4 + ∑_5..8) = (2/4) · (1/8)∑_1..8 ∇L_i
• Divide by N=2: g/2 = (1/8) ∑_1..8 ∇L_i. Same as the true big batch.

Equivalently, if you divide the loss by the accumulation steps before backward, you don't need the final division: loss = loss / accum_steps before loss.backward().

The Simulation

Watch gradients accumulate across micro-batches. Each micro-batch contributes a noisy gradient estimate; the accumulation averages out the noise.

Code: Gradient Accumulation in PyTorch

python
accum_steps = 8  # Effective batch = micro_batch x 8

for step, batch in enumerate(dataloader):
    loss = model(batch)
    loss = loss / accum_steps  # Normalize BEFORE backward
    loss.backward()            # Gradients ADD to .grad buffers

    if (step + 1) % accum_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()  # Reset .grad to zero

Chapter 4: Mixed Precision Training

A 7B-parameter model in FP32 takes 28 GB just for the weights. In FP16, it takes 14 GB. Same model, half the memory, roughly 2× the throughput on modern GPUs. The catch: FP16 can only represent numbers between ±65,504, and anything below ~0.00006 rounds to zero. Gradients are often smaller than that.

The solution is mixed precision training: use half-precision (FP16 or BF16) for the fast parts — forward pass, most of backward pass — but keep a master copy of the weights in full FP32 for the optimizer update. You get the speed of half precision and the accuracy of full precision.

But to understand why this works, we need to understand the three floating point formats and what each one sacrifices.

The Three Formats

Every floating point number is stored as three fields: a sign bit (positive or negative), exponent bits (the scale — powers of 2), and mantissa bits (the precision — significant digits). More exponent bits means wider range. More mantissa bits means finer precision.

Notice the key difference between FP16 and BF16: they're both 16 bits, but they split those bits differently. FP16 gives 5 exponent bits (limited range) and 10 mantissa bits (decent precision). BF16 gives 8 exponent bits (same range as FP32!) and only 7 mantissa bits (less precision). BF16 = FP32's range in FP16's size.

Hand Calculation: Underflow in Action

Let's trace a real gradient through all three formats. Suppose a gradient has value 0.00003 in FP32. This is typical for deep layers in a large model.

In FP32: The smallest positive normal number is ~1.2 × 10^-38. Our value 0.00003 = 3 × 10^-5 is enormous by comparison. Stored exactly (to 7 digits of precision): 0.00003000000. No problem.

In FP16: With only 5 exponent bits, the smallest positive normal number is 2^-14 ≈ 0.00006103. Our gradient is 0.00003 — that's below the smallest normal number. It enters the subnormal range, losing precision rapidly. Values below ~5.96 × 10^-8 round to exactly zero. Our 0.00003 survives as a subnormal (~0.00003), but with only ~1 significant digit. A gradient of 0.000003 would be dead.

In BF16: With 8 exponent bits (same as FP32), the smallest positive normal number is ~1.2 × 10^-38. Our gradient 0.00003 is comfortably representable. It rounds to the nearest BF16 value: approximately 0.0000305 (only ~2.4 digits of precision, but the value is preserved). The weight update happens.

The Mixed Precision Recipe

Here's the complete recipe used by every modern training framework:

The master weights are the key insight. The optimizer (Adam, AdamW, etc.) needs full precision because weight updates are tiny: a learning rate of 10^-4 times a gradient of 10^-3 gives a delta of 10^-7. In BF16, that delta rounds away. In FP32, it accumulates over thousands of steps.

Memory breakdown for a 7B model with Adam:

Wait — that's only 13% savings on paper? The real win is throughput. BF16 matrix multiplications run 2-4× faster on Tensor Cores (A100, H100). And the activation memory drop from 40 to 20 GB means you can double your batch size.

The Simulation

Explore the representable values of each format. Pick a number and see what happens.

Code: Mixed Precision Training Loop

python
# Modern approach: BF16 with torch.autocast
import torch

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

for batch in dataloader:
    inputs, targets = batch

    # Forward pass in BF16 — 2x faster matmuls
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(inputs)
        loss = loss_fn(logits, targets)  # loss auto-cast to FP32

    # Backward in BF16 (autocast context persists for grads)
    loss.backward()

    # Optimizer step: PyTorch auto-converts grads to FP32
    # for the master weight update
    optimizer.step()
    optimizer.zero_grad()

python
# Manual approach — see what autocast does under the hood
model_bf16 = MyModel().cuda().to(torch.bfloat16)
master_weights = {n: p.float().clone() for n, p in model_bf16.named_parameters()}
optimizer = torch.optim.AdamW(master_weights.values(), lr=3e-4)

for batch in dataloader:
    # Forward in BF16
    logits = model_bf16(inputs.bfloat16())
    loss = logits.float().cross_entropy(targets)  # loss in FP32

    # Backward: gradients flow in BF16
    loss.backward()

    # Copy BF16 grads -> FP32 master params
    for n, p in model_bf16.named_parameters():
        master_weights[n].grad = p.grad.float()

    # Optimizer updates FP32 master weights
    optimizer.step()
    optimizer.zero_grad()

    # Copy FP32 master -> BF16 model
    with torch.no_grad():
        for n, p in model_bf16.named_parameters():
            p.copy_(master_weights[n].bfloat16())

Chapter 5: Loss Scaling

Mixed precision training in FP16 still has a problem: many gradients live in the range 10^-5 to 10^-8, all of which underflow to zero in FP16. You might lose 50-80% of your gradient values — the network trains, but slowly and poorly. Loss scaling is the fix.

The idea is beautifully simple. Loss scaling multiplies the loss by a large constant S (say 1024) before backprop. Because backprop is linear in the loss, this scales every single gradient by the same factor S. All those tiny 10^-7 gradients become 10^-4 gradients — safely inside FP16's representable range. After backprop, you divide every gradient by S to get the true values back.

Same math. No underflow. That's the whole trick.

How It Works

Hand Calculation: Rescuing a Gradient

True gradient: g = 0.00002. In FP16, the smallest normal number is ~0.00006. Our gradient enters the subnormal range and loses most of its precision. A value like 0.000003 would underflow to exactly zero.

With loss scaling (S = 1024):

• Scaled gradient: g_scaled = 0.00002 × 1024 = 0.02048
• 0.02048 is well within FP16's normal range — represented precisely to ~3.3 digits
• After backprop, unscale: g = 0.02048 / 1024 = 0.00002
• Back to the true value, but now safely in FP32 for the optimizer step

The gradient survived. Without scaling, it would have been truncated or zeroed. Over thousands of training steps, those lost gradients compound into a slower, worse model.

Dynamic Loss Scaling

What scale factor should you use? Too small, and gradients still underflow. Too large, and gradients overflow (exceed FP16's max of 65,504), producing NaN or Inf. The answer: dynamic loss scaling.

Start with a large scale factor, like S = 2¹⁶ = 65,536. Every training step, check if any gradient overflowed (is NaN or Inf). If yes: skip that optimizer step and halve S. If no overflow for N consecutive steps (typically N = 2000): double S.

This automatically finds the sweet spot — the largest scale factor that doesn't cause overflow. PyTorch's GradScaler does exactly this.

The Simulation

See how loss scaling shifts the gradient histogram from the underflow danger zone into the representable zone.

Code: Loss Scaling in PyTorch

python
# FP16 with GradScaler — the classic approach (pre-BF16 hardware)
import torch
from torch.amp import GradScaler

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = GradScaler()  # Starts with scale = 65536

for batch in dataloader:
    optimizer.zero_grad()

    # Forward in FP16
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(inputs)
        loss = loss_fn(logits, targets)

    # Scale loss -> backward -> scaled gradients in FP16
    scaler.scale(loss).backward()

    # Unscale grads -> clip -> step (skips if overflow detected)
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)   # No-op if grads overflowed
    scaler.update()          # Adjust scale factor

python
# Manual loss scaling — what GradScaler does internally
scale = 65536.0
growth_interval = 2000
steps_since_overflow = 0

for batch in dataloader:
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        logits = model(inputs)
        loss = loss_fn(logits, targets)

    # Scale and backprop
    scaled_loss = loss * scale
    scaled_loss.backward()

    # Check for overflow
    has_overflow = any(
        torch.isinf(p.grad).any() or torch.isnan(p.grad).any()
        for p in model.parameters() if p.grad is not None
    )

    if has_overflow:
        scale /= 2       # Halve on overflow
        steps_since_overflow = 0
        optimizer.zero_grad()
        continue        # Skip this step entirely

    # Unscale gradients
    for p in model.parameters():
        if p.grad is not None:
            p.grad /= scale

    optimizer.step()
    optimizer.zero_grad()

    steps_since_overflow += 1
    if steps_since_overflow >= growth_interval:
        scale *= 2       # Double if stable
        steps_since_overflow = 0

Chapter 6: Gradient Checkpointing

During backprop, you need the activations from the forward pass. Each layer's backward step uses the output of the previous layer to compute gradients. Normally, you store all of these in memory during the forward pass and consume them during the backward pass.

For a 96-layer transformer with a batch of sequences, the stored activations can easily exceed 60 GB — more than the weights, optimizer states, and gradients combined. For very long sequences, activations dominate everything.

Gradient checkpointing (also called activation checkpointing or rematerialization) solves this with a simple tradeoff: store activations at only every K^th layer, and recompute the rest during backprop. You trade a small amount of extra compute for a massive reduction in memory.

The Core Idea

In standard backprop through N layers, the forward pass stores N activations, and the backward pass consumes them in reverse. Memory = O(N).

With checkpointing every K layers, you divide the network into N/K segments of K layers each. During the forward pass, you only save the activation at the boundary of each segment — that's N/K saved activations. During backprop, when you need a layer's activation inside a segment, you re-run the forward pass from the nearest saved checkpoint.

The saved checkpoints cost N/K memory. The recomputation within each segment needs at most K activations at a time (just the current segment). Total memory: N/K + K. The extra compute cost: one additional forward pass per segment = N/K forward passes of K layers each = N extra layer-forward-passes = roughly one full extra forward pass through the entire network.

Hand Calculation: 64-Layer Model

Standard backprop (K=1, save everything):

• Stored activations: 64
• Extra forward passes: 0
• Total forward compute: 1× (just the normal forward pass)

Checkpoint every K=8 layers:

• Segments: 64 / 8 = 8 segments
• Saved checkpoints: 8 (one per segment boundary)
• During backprop, recompute within each segment: need at most 8 activations at a time
• Peak memory: 8 (checkpoints) + 8 (current segment) = 16 activations
• Extra forward compute: 8 segments × 8 layers/segment = 64 layer passes = 1 extra full forward pass
• Total compute: ~1.33× (forward + backward + recompute ≈ 1F + 2F + 1F = 4F vs standard 3F)
• Memory savings: 64 → 16 = 4× reduction

Checkpoint every K=16 layers:

• Segments: 64 / 16 = 4 segments
• Saved checkpoints: 4
• Peak memory: 4 + 16 = 20 activations
• Extra forward compute: 4 × 16 = 64 layer passes = 1 extra full forward pass

Extreme: K=N (save nothing, checkpoint only input):

• Saved: 1 (just the input)
• Peak memory: 1 + 64 = 65 activations — wait, that's worse! The recomputation within the single segment still needs all 64 layer activations. This extreme only helps if you recompute layer-by-layer. In practice, "save none" reduces to 1 + N peak with naive implementation, or O(N) with recursive checkpointing — but at 2× the compute cost.

Code: Checkpointing in PyTorch

python
import torch
from torch.utils.checkpoint import checkpoint

class TransformerWithCheckpointing(torch.nn.Module):
    def __init__(self, n_layers=64, d_model=4096):
        super().__init__()
        self.layers = torch.nn.ModuleList([
            TransformerBlock(d_model) for _ in range(n_layers)
        ])
        self.ckpt_every = int(n_layers ** 0.5)  # sqrt(N)

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            if i % self.ckpt_every == 0 and self.training:
                # Checkpoint this layer — don't save activation,
                # recompute during backward pass
                x = checkpoint(layer, x, use_reentrant=False)
            else:
                x = layer(x)
        return x

python
# Even simpler: checkpoint_sequential for a block of layers
from torch.utils.checkpoint import checkpoint_sequential

class SimpleCheckpointModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(*[
            TransformerBlock(4096) for _ in range(64)
        ])

    def forward(self, x):
        # Split into 8 segments, checkpoint each
        return checkpoint_sequential(self.layers, segments=8, input=x)

The Showcase

Watch checkpointing in action. Configure the model, set the checkpoint interval, and step through forward and backward passes to see exactly which activations are saved, which are recomputed, and how memory compares.

Chapter 7: The Modern Training Pipeline

You've learned each technique in isolation. Now let's see how they all fit together in a single training step of a real LLM. Every one of these techniques is used simultaneously — they're not alternatives, they're a stack.

Here is the complete pipeline for one optimizer step in a modern LLM training run (like LLaMA, GPT-4, or Gemma). Every box is a technique from this lesson.

The Complete Pipeline

Why Every Step Matters

Remove any one of these techniques and something breaks at scale:

The Complete Code

python
# Complete modern training loop — all techniques combined
import torch
from torch.utils.checkpoint import checkpoint

# -- Setup --
model = LLaMA_7B().cuda().bfloat16()  # BF16 model
master = {n: p.float().clone() for n, p in model.named_parameters()}
optimizer = torch.optim.AdamW(master.values(), lr=3e-4, betas=(0.9, 0.95))
scheduler = CosineWarmupScheduler(optimizer, warmup=2000, total=100000)

accum_steps = 8      # 8 micro-batches per optimizer step
max_norm = 1.0       # gradient clip threshold

# -- Training Loop --
for step, batch in enumerate(dataloader):
    micro_batches = split(batch, accum_steps)

    # -- Accumulate gradients over micro-batches --
    for i, micro in enumerate(micro_batches):
        # 1. Forward in BF16 (with checkpointing inside model)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            logits = model(micro.input_ids)    # checkpointed internally
            loss = cross_entropy(logits, micro.labels)
            loss = loss / accum_steps          # normalize for accumulation

        # 3. Backward — grads accumulate in .grad buffers
        loss.backward()

    # 4. Copy accumulated BF16 grads -> FP32 master params
    for n, p in model.named_parameters():
        master[n].grad = p.grad.float()

    # 5. Clip gradient norm
    torch.nn.utils.clip_grad_norm_(master.values(), max_norm=max_norm)

    # 6. Optimizer step (FP32)
    optimizer.step()
    optimizer.zero_grad()

    # 7. LR schedule step
    scheduler.step()

    # 8. Copy FP32 master -> BF16 model
    with torch.no_grad():
        for n, p in model.named_parameters():
            p.copy_(master[n].bfloat16())

    # Zero model grads for next accumulation
    model.zero_grad()

The Pipeline Visualizer

Click any stage to see its details. Toggle options to see how the pipeline changes.

Memory Budget for a 7B Model on H100 (80 GB)

Even with every optimization, a 7B model's training state exceeds 80 GB. That's why real LLM training uses model parallelism (split the model across multiple GPUs) and ZeRO (shard optimizer states). But that's a distributed systems lesson, not a gradient flow lesson.

Chapter 8: The Gradient Flow Arena

Let's race them all.

You've learned six techniques for taming gradients: no intervention, gradient clipping alone, mixed precision (BF16), gradient accumulation, gradient checkpointing, and the full modern pipeline. Each handles different failure modes — but reading about them is one thing. Watching them train side by side is another.

This simulation trains six configurations on the same loss landscape. Drag the sliders to create the conditions where each one fails. You'll discover that no single technique is enough — you need the full stack.

What to Discover

Experiment 1: Depth = 48 layers. The baseline (no clipping, no mixed precision) explodes almost immediately — the gradient product through 48 layers overflows. With clipping alone, it survives but trains slowly. The full stack trains smoothly because BF16 prevents the intermediate overflow and clipping catches spikes.

Experiment 2: Learning rate = 1.0. Drag the LR slider all the way right. Everything except the full stack diverges. The full stack's gradient clipping caps the step size, and the optimizer's adaptive moments (Adam) further stabilize. Clipping alone helps but isn't sufficient without the momentum buffer.

Experiment 3: Batch size = 1. With a single sample per batch, gradient variance is enormous. The accumulation config (which simulates a larger effective batch) produces a much smoother loss curve. This is the main benefit of gradient accumulation: variance reduction, not just memory savings.

Experiment 4: Compare +BF16 vs +Clip. At moderate depth and LR, both work well. But crank the depth to 48: BF16 prevents intermediate overflow during the forward pass (smaller numbers), while clipping only fixes the backward pass. You need both.

Failure Mode Summary

Chapter 9: Cheat Sheet & Connections

You now understand the complete gradient flow toolkit — from the chain rule to the full modern training pipeline. This chapter is your practical reference. No new concepts. Just the techniques, the decision guide, and the connections to where you go next.

Every Technique at a Glance

The Decision Flowchart

Follow the path that matches your situation:

Symbol Glossary

PyTorch One-Liner Reference

python
import torch
from torch.amp import GradScaler
from torch.utils.checkpoint import checkpoint

# Gradient clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# Mixed precision (BF16 — preferred)
with torch.autocast("cuda", dtype=torch.bfloat16): ...

# Mixed precision (FP16 + loss scaling)
scaler = GradScaler()
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
scaler.step(optimizer)
scaler.update()

# Gradient accumulation
loss = loss / accum_steps
loss.backward()  # grads accumulate in .grad
if step % accum_steps == 0: optimizer.step(); optimizer.zero_grad()

# Gradient checkpointing
x = checkpoint(layer, x, use_reentrant=False)

Summary of Everything

Connections

Gradient flow doesn't exist in isolation. Here's where to go next:

• Optimizers — Adam's adaptive learning rate per-parameter is the complement to gradient clipping. Together they stabilize training dynamics.
• Normalization — LayerNorm and RMSNorm stabilize activations in the forward pass, which stabilizes gradients in the backward pass. Pre-LN placement creates a gradient highway.
• Loss Functions — The loss is where gradients begin. Cross-entropy's gradient behavior directly affects whether clipping is needed and at what threshold.
• GPT — GPT's training at scale uses every technique from this lesson. See them in context.
• Learning Rate Schedules — Warmup prevents gradient explosion in early training. Cosine decay prevents overshooting in late training. The LR schedule and gradient clipping work in concert.

Key Papers