Module 07: Training

Introduction

Training teaches a language model to predict the next token. The process iterates:

  1. Computing loss: How wrong are our predictions?
  2. Computing gradients: Which direction should we adjust weights?
  3. Updating weights: Take a small step in that direction
  4. Repeat: Until the model gets good at prediction

This module covers cross-entropy loss, the AdamW optimizer, learning rate scheduling, gradient accumulation, and checkpointing.

What You’ll Learn

After this module, you can:

  • Understand cross-entropy loss and perplexity for language models
  • Implement learning rate schedules (warmup + cosine decay)
  • Use gradient accumulation for effective larger batch sizes
  • Apply gradient clipping for training stability
  • Read the floating-point number line from scratch — why fp16 overflows and underflows where bf16 doesn’t, and how loss scaling rescues small gradients
  • Build both fp8 formats (e4m3 / e5m2) — why e4m3 drops IEEE’s infinity — and the per-tensor scaling that makes 8-bit training work
  • Save and load model checkpoints
  • Plan a compute-optimal run with scaling laws (the C≈6ND rule, Chinchilla’s ~20 tokens/param)
  • Fit a scaling law from scratch — sweep IsoFLOP profiles, read each valley, and recover the compute-optimal exponents (N\propto C^{0.45}, D\propto C^{0.55}) by log-log least squares, so “20 tokens/param” arrives as a measured result — then make the runs noisy and see why the valley must be fit with a parabola, not read off the single lowest run
  • Build Muon from scratch — orthogonalize the momentum matrix with a from-scratch Newton–Schulz iteration, the first serious challenger to AdamW at LLM scale

Prerequisites

This module requires familiarity with:

Note: This lesson demonstrates concepts interactively. The training.py file provides production-ready implementations of the same algorithms.

The Training Objective

Language models learn through next-token prediction:

Input:    [The, cat, sat, on, the]
Target:   [cat, sat, on, the, mat]

For each position, predict the next token.

The loss function measures how well the model predicts: Cross-entropy between predicted probabilities and actual next tokens.

\text{loss} = -\sum \log(P(\text{correct\_token}))

Lower loss means the model assigns correct tokens higher probability.

The Training Loop

Neural networks learn through the training loop:

Note: zero_grad() can be called either at the start or end of each iteration. Calling it at the start (shown above) is common because it ensures gradients are fresh before the backward pass.

Setup

import sys
import math
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# For reproducibility
torch.manual_seed(42)

# Display device info
device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
print(f"PyTorch version: {torch.__version__}")
print(f"Device: {device}")
PyTorch version: 2.9.1
Device: mps

Cross-Entropy Loss

The loss function measures prediction error. Cross-entropy loss penalizes wrong predictions more heavily when the model is confident but incorrect.

Why cross-entropy?

  1. Probabilistic interpretation: It measures the “surprise” when the true token appears
  2. Gradient properties: Gradients are proportional to the error (predicted - actual)
  3. Information theory: Minimizing cross-entropy = maximizing likelihood of data

Mathematical formulation:

\text{CrossEntropy}(p, q) = -\sum_{i} p_i \log(q_i)

For language modeling with one-hot targets (only one correct token), this simplifies to:

\text{Loss} = -\log(q_{\text{correct}})

where q_{\text{correct}} is the probability the model assigns to the correct token.

# Example: Model predicting next token
vocab_size = 10

# Model outputs logits (raw scores)
logits = torch.tensor([
    [-1.0, 0.5, 2.0, -0.5, 1.0, 0.0, -1.5, 0.3, -0.8, 0.2]  # scores for each token
])

# True next token is index 2
target = torch.tensor([2])

# Convert to probabilities
probs = F.softmax(logits, dim=-1)

print("Logits (raw model output):")
print(f"  {logits[0].tolist()}")
print(f"\nProbabilities (after softmax):")
print(f"  {[f'{p:.3f}' for p in probs[0].tolist()]}")
print(f"\nTarget token: {target.item()}")
print(f"Probability assigned to target: {probs[0, target.item()]:.4f}")
Logits (raw model output):
  [-1.0, 0.5, 2.0, -0.5, 1.0, 0.0, -1.5, 0.30000001192092896, -0.800000011920929, 0.20000000298023224]

Probabilities (after softmax):
  ['0.022', '0.097', '0.435', '0.036', '0.160', '0.059', '0.013', '0.080', '0.026', '0.072']

Target token: 2
Probability assigned to target: 0.4353
# Cross-entropy loss
loss = F.cross_entropy(logits, target)
manual_loss = -torch.log(probs[0, target.item()])

print(f"Cross-entropy loss: {loss.item():.4f}")
print(f"Manual calculation: -log({probs[0, target.item()]:.4f}) = {manual_loss.item():.4f}")

# Perplexity
perplexity = math.exp(loss.item())
print(f"\nPerplexity: {perplexity:.2f}")
Cross-entropy loss: 0.8317
Manual calculation: -log(0.4353) = 0.8317

Perplexity: 2.30

Let’s visualize how loss changes with probability:

Higher probability means lower loss means better predictions.

Cross-Entropy from Scratch

Before using F.cross_entropy, let’s understand what it does internally.

The Numerical Stability Problem

Softmax involves exp(x), which explodes for large x:

# The problem: exp() overflows easily
logits_big = np.array([1000.0, 1001.0, 1002.0])
print(f"exp(logits) = {np.exp(logits_big)}")  # [inf, inf, inf] - overflow!
exp(logits) = [inf inf inf]
/var/folders/hl/bw75m5hd5xvfyx8j9qd71vjw0000gn/T/ipykernel_65595/1953512672.py:3: RuntimeWarning: overflow encountered in exp
  print(f"exp(logits) = {np.exp(logits_big)}")  # [inf, inf, inf] - overflow!

The Fix: Log-Sum-Exp Trick

The key insight is that we can compute log-softmax stably by subtracting the maximum:

\log \text{softmax}(x_i) = x_i - \log\sum_j e^{x_j} = x_i - \underbrace{(m + \log\sum_j e^{x_j - m})}_{\text{logsumexp}}

where m = \max(x). By subtracting the max, all exponents become \leq 0, avoiding overflow.

def logsumexp(x: np.ndarray, axis: int = -1, keepdims: bool = True) -> np.ndarray:
    """
    Stable log(sum(exp(x))).

    Trick: log(sum(exp(x))) = m + log(sum(exp(x - m)))
    where m = max(x). This keeps exp() arguments <= 0.
    """
    m = x.max(axis=axis, keepdims=True)
    return m + np.log(np.exp(x - m).sum(axis=axis, keepdims=keepdims))

# Now it works!
print(f"logsumexp(logits) = {logsumexp(logits_big, keepdims=False)}")
logsumexp(logits) = [1002.40760596]

Cross-Entropy Implementation

def cross_entropy_scratch(logits: np.ndarray, targets: np.ndarray) -> float:
    """
    Cross-entropy loss from logits.

    logits: (B, C) - raw scores for each class
    targets: (B,) - integer class labels

    Formula: loss = logsumexp(logits) - logits[correct_class]

    This is equivalent to: -log(softmax(logits)[correct_class])
    but numerically stable.
    """
    B, C = logits.shape

    # log(sum(exp(logits))) for normalization
    lse = logsumexp(logits, axis=-1, keepdims=False).squeeze()  # (B,)

    # Gather correct class logits
    correct_logits = logits[np.arange(B), targets]  # (B,)

    # Loss per sample, then mean
    losses = lse - correct_logits
    return float(losses.mean())

# Test
test_logits = np.array([[2.0, 1.0, 0.1], [0.5, 2.5, 0.3]])
test_targets = np.array([0, 1])  # First sample: class 0, second: class 1
print(f"Cross-entropy loss (scratch): {cross_entropy_scratch(test_logits, test_targets):.4f}")
Cross-entropy loss (scratch): 0.3185

PyTorch Equivalent

# Compare with PyTorch
logits_pt = torch.tensor([[2.0, 1.0, 0.1], [0.5, 2.5, 0.3]])
targets_pt = torch.tensor([0, 1])
loss_pt = F.cross_entropy(logits_pt, targets_pt)
print(f"Cross-entropy loss (PyTorch): {loss_pt.item():.4f}")
Cross-entropy loss (PyTorch): 0.3185

Same result! PyTorch’s F.cross_entropy does exactly this internally, plus handles gradients automatically.

NoteKey Insight

Cross-entropy is just logsumexp(logits) - logits[correct_class]. The logsumexp trick prevents numerical overflow by subtracting the max before exponentiating.

Perplexity

Perplexity is a more intuitive measure than raw loss:

\text{Perplexity} = e^{\text{cross\_entropy\_loss}}

Interpretation: “The model is as confused as if it were choosing uniformly among N options.”

Loss Perplexity Interpretation
0.0 1.0 Perfect predictions
2.3 10 ~10 equally likely options
4.6 100 ~100 equally likely options
6.9 1000 Random guessing (vocab=1000)

For reference: - GPT-2 on WebText: ~20 perplexity - Human baseline: ~10-20 perplexity (depends on domain)

Learning Rate Schedule

We vary the learning rate over training, using warmup followed by cosine decay:

Why warmup? - Early training is unstable with large LR - Gradients are noisy before weights settle - Small LR lets model “get its bearings”

Why decay? - Large LR is good for exploration early - Small LR is good for fine-tuning later - Cosine is smooth (no sudden changes)

class CosineScheduler:
    """Learning rate scheduler with linear warmup and cosine decay."""

    def __init__(self, optimizer, warmup_steps, total_steps, min_lr=0.0):
        self.optimizer = optimizer
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.min_lr = min_lr
        self.base_lr = optimizer.param_groups[0]['lr']
        self.current_step = 0

    def get_lr(self):
        """Calculate learning rate for current step."""
        if self.current_step < self.warmup_steps:
            # Linear warmup
            return self.base_lr * self.current_step / max(1, self.warmup_steps)
        elif self.current_step >= self.total_steps:
            return self.min_lr
        else:
            # Cosine decay
            progress = (self.current_step - self.warmup_steps) / max(
                1, self.total_steps - self.warmup_steps
            )
            cosine = 0.5 * (1 + math.cos(math.pi * progress))
            return self.min_lr + (self.base_lr - self.min_lr) * cosine

    def step(self):
        """Update learning rate."""
        lr = self.get_lr()
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
        self.current_step += 1
        return lr

# Create scheduler
model = nn.Linear(10, 10)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

scheduler = CosineScheduler(
    optimizer,
    warmup_steps=100,
    total_steps=1000,
    min_lr=1e-5
)

# Collect LRs over training
lrs = []
for _ in range(1000):
    lrs.append(scheduler.get_lr())
    scheduler.step()

print(f"Initial LR: {lrs[0]:.6f}")
print(f"After warmup (step 100): {lrs[100]:.6f}")
print(f"Final LR: {lrs[-1]:.6f}")
Initial LR: 0.000000
After warmup (step 100): 0.001000
Final LR: 0.000010

The interactive visualization above shows how learning rate changes over training. Try adjusting the warmup steps and total steps sliders to see how they affect the schedule.

AdamW Optimizer

AdamW decouples weight decay from Adam (proper L2 regularization) and serves as the standard optimizer for language models.

Why AdamW over SGD or Adam?

  • SGD: Requires careful learning rate tuning per layer, slow convergence
  • Adam: Weight decay is applied to gradients (incorrect for L2 regularization)
  • AdamW: Decouples weight decay from gradient updates (mathematically correct)

Hyperparameters explained:

Parameter Default Purpose
beta1 0.9 Momentum coefficient - smooths gradient direction
beta2 0.999 Adaptive LR coefficient - smooths gradient magnitude
epsilon 1e-8 Numerical stability (prevents division by zero)
weight_decay 0.01 L2 regularization strength

Practical tip: The LLM community has converged on beta1=0.9, beta2=0.95 for large models (used by LLaMA, GPT-3). The lower beta2 adapts faster to changing gradient magnitudes.

# Creating an AdamW optimizer
model = nn.Linear(100, 10)

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,          # Learning rate
    betas=(0.9, 0.999),  # Momentum and adaptive LR
    weight_decay=0.01    # Regularization
)

print("AdamW optimizer created")
print(f"  Learning rate: {optimizer.param_groups[0]['lr']}")
print(f"  Weight decay: {optimizer.param_groups[0]['weight_decay']}")
AdamW optimizer created
  Learning rate: 0.0003
  Weight decay: 0.01

Optimizers from Scratch

Let’s build optimizers from first principles to understand what PyTorch does internally.

Plain SGD

The simplest optimizer: move parameters in the opposite direction of the gradient.

class SGD_Scratch:
    """
    Stochastic Gradient Descent.

    Update rule: theta = theta - lr * gradient
    """
    def __init__(self, params, lr=0.01):
        self.params = list(params)
        self.lr = lr

    def step(self):
        with torch.no_grad():
            for p in self.params:
                if p.grad is not None:
                    p -= self.lr * p.grad

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad = None

# Test: compare with PyTorch SGD
torch.manual_seed(42)
model_scratch = nn.Linear(10, 2)
model_pytorch = nn.Linear(10, 2)
model_pytorch.load_state_dict(model_scratch.state_dict())

opt_scratch = SGD_Scratch(model_scratch.parameters(), lr=0.1)
opt_pytorch = torch.optim.SGD(model_pytorch.parameters(), lr=0.1)

# Forward + backward
x = torch.randn(4, 10)
loss_scratch = model_scratch(x).sum()
loss_pytorch = model_pytorch(x).sum()

loss_scratch.backward()
loss_pytorch.backward()

# Update
opt_scratch.step()
opt_pytorch.step()

# Compare weights
print("After one SGD step:")
print(f"  Scratch weight[0,0]: {model_scratch.weight[0,0].item():.6f}")
print(f"  PyTorch weight[0,0]: {model_pytorch.weight[0,0].item():.6f}")
print(f"  Match: {torch.allclose(model_scratch.weight, model_pytorch.weight)}")
After one SGD step:
  Scratch weight[0,0]: 0.164868
  PyTorch weight[0,0]: 0.164868
  Match: True

SGD with Momentum

Momentum adds “velocity” to gradient descent. Instead of using the gradient directly, we accumulate a moving average of gradients:

v_t = \mu \cdot v_{t-1} + g_t \theta_t = \theta_{t-1} - \alpha \cdot v_t

This helps: - Smooth out noisy gradients - Accelerate through flat regions - Dampen oscillations in steep valleys

class SGD_Momentum_Scratch:
    """
    SGD with momentum.

    Update rule:
        v = momentum * v + gradient
        theta = theta - lr * v
    """
    def __init__(self, params, lr=0.01, momentum=0.9):
        self.params = list(params)
        self.lr = lr
        self.momentum = momentum
        # Velocity buffer for each parameter
        self.v = [torch.zeros_like(p) for p in self.params]

    def step(self):
        with torch.no_grad():
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                # Update velocity: v = momentum * v + grad
                self.v[i] = self.momentum * self.v[i] + p.grad
                # Update parameter
                p -= self.lr * self.v[i]

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad = None

# Test: compare with PyTorch SGD momentum
torch.manual_seed(42)
model_scratch = nn.Linear(10, 2)
model_pytorch = nn.Linear(10, 2)
model_pytorch.load_state_dict(model_scratch.state_dict())

opt_scratch = SGD_Momentum_Scratch(model_scratch.parameters(), lr=0.1, momentum=0.9)
opt_pytorch = torch.optim.SGD(model_pytorch.parameters(), lr=0.1, momentum=0.9)

# Multiple steps to see momentum accumulate
for step in range(3):
    x = torch.randn(4, 10)

    loss_scratch = model_scratch(x).sum()
    loss_pytorch = model_pytorch(x).sum()

    opt_scratch.zero_grad()
    opt_pytorch.zero_grad()

    loss_scratch.backward()
    loss_pytorch.backward()

    opt_scratch.step()
    opt_pytorch.step()

print("After 3 momentum SGD steps:")
print(f"  Scratch weight[0,0]: {model_scratch.weight[0,0].item():.6f}")
print(f"  PyTorch weight[0,0]: {model_pytorch.weight[0,0].item():.6f}")
print(f"  Match: {torch.allclose(model_scratch.weight, model_pytorch.weight)}")
After 3 momentum SGD steps:
  Scratch weight[0,0]: -0.793227
  PyTorch weight[0,0]: -0.793227
  Match: True
NoteKey Insight: Momentum

Momentum is like pushing a ball down a hill - it builds up speed in consistent directions and resists sudden direction changes. This makes optimization faster and more stable.

Adam from Scratch

Adam combines momentum with adaptive learning rates. It tracks two quantities:

  1. First moment m (mean of gradients) - like momentum
  2. Second moment v (mean of squared gradients) - adapts learning rate per-parameter

m_t = \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t v_t = \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2

We also need bias correction because m and v are initialized to zero:

\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}

Finally, the update:

\theta_t = \theta_{t-1} - \alpha \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

class Adam_Scratch:
    """
    Adam optimizer with optional weight decay (AdamW style).

    Tracks first moment (mean) and second moment (variance) of gradients.
    Uses bias correction to fix initialization bias.
    """
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0):
        self.params = list(params)
        self.lr = lr
        self.b1, self.b2 = betas
        self.eps = eps
        self.weight_decay = weight_decay

        # First moment (mean of gradients)
        self.m = [torch.zeros_like(p) for p in self.params]
        # Second moment (mean of squared gradients)
        self.v = [torch.zeros_like(p) for p in self.params]
        # Timestep
        self.t = 0

    def step(self):
        self.t += 1

        with torch.no_grad():
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue

                g = p.grad

                # AdamW: Weight decay applied directly to weights (decoupled)
                if self.weight_decay != 0.0:
                    p -= self.lr * self.weight_decay * p

                # Update first moment: m = beta1 * m + (1 - beta1) * g
                self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g

                # Update second moment: v = beta2 * v + (1 - beta2) * g^2
                self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (g * g)

                # Bias correction (crucial early in training!)
                mhat = self.m[i] / (1 - self.b1 ** self.t)
                vhat = self.v[i] / (1 - self.b2 ** self.t)

                # Update parameters
                p -= self.lr * mhat / (torch.sqrt(vhat) + self.eps)

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad = None

# Test: compare with PyTorch AdamW
torch.manual_seed(42)
model_scratch = nn.Linear(10, 2)
model_pytorch = nn.Linear(10, 2)
model_pytorch.load_state_dict(model_scratch.state_dict())

opt_scratch = Adam_Scratch(model_scratch.parameters(), lr=1e-3, weight_decay=0.01)
opt_pytorch = torch.optim.AdamW(model_pytorch.parameters(), lr=1e-3, weight_decay=0.01)

# Multiple steps
for step in range(5):
    x = torch.randn(4, 10)

    loss_scratch = model_scratch(x).sum()
    loss_pytorch = model_pytorch(x).sum()

    opt_scratch.zero_grad()
    opt_pytorch.zero_grad()

    loss_scratch.backward()
    loss_pytorch.backward()

    opt_scratch.step()
    opt_pytorch.step()

print("After 5 AdamW steps:")
print(f"  Scratch weight[0,0]: {model_scratch.weight[0,0].item():.6f}")
print(f"  PyTorch weight[0,0]: {model_pytorch.weight[0,0].item():.6f}")
print(f"  Close match: {torch.allclose(model_scratch.weight, model_pytorch.weight, atol=1e-6)}")
After 5 AdamW steps:
  Scratch weight[0,0]: 0.237638
  PyTorch weight[0,0]: 0.237638
  Close match: True
NoteKey Insight: Adam

Adam is “momentum + per-parameter learning rates.” The second moment v tracks how much each parameter’s gradient varies. Parameters with consistently large gradients get smaller effective learning rates (stabilizing training), while those with small gradients get larger rates (speeding up learning).

Why bias correction matters:

Without bias correction, the first few steps are biased toward zero because m and v are initialized to zero. Let’s see this:

# Demonstrate bias correction importance
m, v = 0.0, 0.0
b1, b2 = 0.9, 0.999
true_grad = 1.0  # Pretend gradient is always 1

print("Step | m (biased) | m_hat (corrected)")
print("-" * 45)
for t in range(1, 6):
    m = b1 * m + (1 - b1) * true_grad
    m_hat = m / (1 - b1 ** t)
    print(f"  {t}  |   {m:.4f}    |     {m_hat:.4f}")

print(f"\nWithout correction, m starts near 0.1 instead of 1.0!")
print(f"Bias correction fixes this, making m_hat ≈ 1.0 from the start.")
Step | m (biased) | m_hat (corrected)
---------------------------------------------
  1  |   0.1000    |     1.0000
  2  |   0.1900    |     1.0000
  3  |   0.2710    |     1.0000
  4  |   0.3439    |     1.0000
  5  |   0.4095    |     1.0000

Without correction, m starts near 0.1 instead of 1.0!
Bias correction fixes this, making m_hat ≈ 1.0 from the start.

From Inline Sketch to a Tested Module

The classes above are teaching sketches. The production versions — SGD, Adam, and AdamW, each with a shared Optimizer base and full type hints — live in optimizers.py. They are checked bit-for-bit against torch.optim: the test suite trains a twin nn.Linear under each hand-written optimizer and the matching PyTorch one and asserts the weights agree to floating-point rounding. When we say “this is what PyTorch does,” it’s a tested claim, not a slogan.

from optimizers import SGD, Adam, AdamW

# Our AdamW vs torch.optim.AdamW, five steps, same init.
torch.manual_seed(0)
a = nn.Linear(10, 3)
b = nn.Linear(10, 3)
b.load_state_dict(a.state_dict())

ours = AdamW(a.parameters(), lr=1e-2, weight_decay=0.1)
theirs = torch.optim.AdamW(b.parameters(), lr=1e-2, weight_decay=0.1)

for _ in range(5):
    x = torch.randn(5, 10)
    ours.zero_grad(); theirs.zero_grad()
    a(x).pow(2).sum().backward()
    b(x).pow(2).sum().backward()
    ours.step(); theirs.step()

print(f"max |ours - torch|: {(a.weight - b.weight).abs().max().item():.2e}")
print(f"Match: {torch.allclose(a.weight, b.weight, atol=1e-6)}")
max |ours - torch|: 2.98e-08
Match: True

Adam vs AdamW: Why Decoupling Matters

Earlier we said Adam applies weight decay “to the gradient” and AdamW “to the weights.” That one-word difference in placement is the whole reason AdamW exists. Both start from the same coefficient \lambda:

\underbrace{g_t \leftarrow g_t + \lambda\,\theta_{t-1}}_{\textbf{Adam (coupled L2)}} \qquad\text{vs.}\qquad \underbrace{\theta_{t-1} \leftarrow \theta_{t-1} - \alpha\lambda\,\theta_{t-1}}_{\textbf{AdamW (decoupled)}}

In AdamW the decay is a plain shrink: every weight loses the same fraction \alpha\lambda of its magnitude each step, no matter what its gradient is doing.

In Adam the decay is folded into the gradient, so it rides through the rest of the update — including the adaptive denominator \frac{1}{\sqrt{\hat v_t}+\epsilon}. A weight whose gradients are large (big \hat v) gets its decay divided down; a weight with tiny gradients gets its decay through at nearly full strength. The regularization you asked for silently becomes a different, per-parameter amount. Run the two with identical lr and weight_decay and they simply disagree:

from optimizers import Adam, AdamW

torch.manual_seed(1)
a = nn.Linear(8, 4)
b = nn.Linear(8, 4)
b.load_state_dict(a.state_dict())

coupled = Adam(a.parameters(), lr=1e-2, weight_decay=0.3)   # classic Adam L2
decoupled = AdamW(b.parameters(), lr=1e-2, weight_decay=0.3)  # AdamW

for _ in range(25):
    x = torch.randn(6, 8)
    coupled.zero_grad(); decoupled.zero_grad()
    a(x).pow(2).sum().backward()
    b(x).pow(2).sum().backward()
    coupled.step(); decoupled.step()

print(f"Same lr, same weight_decay, 25 steps:")
print(f"  max |Adam(L2) - AdamW| weight gap: {(a.weight - b.weight).abs().max().item():.4f}")
print(f"  identical? {torch.allclose(a.weight, b.weight, atol=1e-4)}")
Same lr, same weight_decay, 25 steps:
  max |Adam(L2) - AdamW| weight gap: 0.0150
  identical? False
NoteKey Insight: Where the Decay Lands

Set weight_decay=0 and Adam and AdamW are bit-identical — there is nothing to decouple. Turn decay on and they part ways, because coupled L2 passes through the 1/\sqrt{\hat v} denominator and decoupled decay does not. The plot below makes the damping literal: AdamW’s effective decay is flat across parameters; Adam’s falls off as gradients grow.

The next cell bridges effective_decay from optimizers.py — the steady-state decay each scheme actually applies to a weight whose gradient has a given magnitude.

TipTry This
  1. Push λ up: both lines rise, but only the orange (coupled) line stays bent — parameters with large gradients keep getting under-decayed relative to what you set.
  2. λ → 0: the two lines collapse onto the axis together. No decay, nothing to decouple — exactly why our Adam and AdamW are bit-identical at weight_decay=0.

The Optimizer Zoo, Side by Side

Weight decay aside, why reach for Adam at all? The clearest picture is a path on a hard surface. demonstrate_optimizers in optimizers.py runs SGD, SGD+momentum, and Adam on an ill-conditioned bowl f(x,y)=\tfrac12(25x^2+y^2) — 25× steeper along x than y — and records where each one goes.

Scrub the step and watch the shapes, not a winner:

  • SGD must use a step small enough for the steep x axis (here \alpha=1/25), so it snaps to the valley floor almost immediately, then crawls along the shallow y axis — the classic ill-conditioning crawl.
  • Momentum builds velocity down the valley, overshoots the bottom, and curls back.
  • Adam normalizes each axis by its own gradient history and marches at a roughly constant per-coordinate pace — and, honestly, tends to hover near the minimum on a clean quadratic rather than settle into it.
NoteKey Insight: Why AdamW Is the LLM Default

On this toy quadratic, momentum actually reaches the lowest loss — Adam is not magic. Adam earns its place in deep networks, where different parameters see gradients that differ by orders of magnitude (embeddings vs. layer norms vs. deep weights). Its per-parameter scaling makes one global learning rate work across all of them with little tuning, and AdamW adds the decoupled decay that keeps that scaling from corrupting your regularization. Robustness at scale, not toy-problem speed, is the win.

Beyond AdamW: Orthogonalized Momentum (Muon)

We just crowned AdamW the LLM default — and it has held that crown for a decade. The first optimizer to seriously challenge it doesn’t add a cleverer per-weight learning rate. It changes what a “weight” is.

Every optimizer so far treats a weight matrix as a bag of independent scalars. Adam’s 1/√v̂ rescales each scalar on its own. But a hidden weight W is a matrix, and matrices have structure that per-scalar rules are blind to: directions (its singular vectors) along which it stretches a lot or a little. When one direction dominates the gradient, per-scalar updates pour most of their energy there and starve the rest — the update is ill-conditioned.

Muon (Keller Jordan, 2024; scaled to LLM pretraining by Moonshot’s Muon is Scalable, 2025) fixes this at the level of the whole matrix. Take the momentum matrix M and factor it with the SVD, M = U\,\Sigma\,V^\top. Muon throws away the singular values \Sigma — the very thing that makes the update lopsided — and steps along

O = U V^\top \qquad(\text{all singular values set to } 1).

O is the nearest semi-orthogonal matrix to M: same directions, but every one advanced at the same rate. The catch is that an SVD every step, on every weight, would be far too slow. The whole trick is computing UV^\top without one.

The Newton–Schulz iteration

Muon approximates UV^\top with a fixed matrix polynomial, iterated about five times. First normalize M by its Frobenius norm (so its largest singular value is \le 1), then repeat

X \leftarrow a\,X + b\,(XX^\top)X + c\,(XX^\top)^2 X, \qquad (a,b,c) = (3.4445,\ -4.7750,\ 2.0315).

Because this iteration is odd in X and never touches the singular vectors, it acts on each singular value on its own through the scalar quintic p(\sigma) = a\sigma + b\sigma^3 + c\sigma^5. Here is the subtle part, and it is easy to get wrong: this is not a convergent iteration. Run it forever and the singular values oscillate. The coefficients are tuned for a fixed budget of ~5 steps — the steep slope at zero (a = 3.44 > 1) yanks even tiny singular values upward fast, and the higher-order terms bend the curve so it overshoots just past 1. After five steps every \sigma \in (0,1] lands in a loose band around 1 (roughly [0.7, 1.3]). That is approximate orthogonalization — Muon happily trades an exact \sigma = 1 for lifting the small directions quickly.

Drive the iteration below. Each line is one singular value of a real momentum matrix; watch them march out of the decaying spectrum toward the shaded band at 1.

TipTry This

Drag the step to 0: the singular values fan out — a decaying spectrum, the ill-conditioned update. Now step to 5: they collapse into the green band at 1. Notice the smallest value climbs the most and the largest barely moves — the iteration lifts starved directions without touching dominant ones. That is the entire point: equalize every direction of the update.

Orthogonalization from scratch

muon.py builds two versions: newton_schulz (the iteration above, matmuls only) and orthogonalize (the exact UV^\top via torch.linalg.svd) that we keep only as ground truth. Muon uses the first because it avoids the second.

from muon import newton_schulz, orthogonalize, singular_values
import torch

torch.manual_seed(0)
# An ill-conditioned momentum matrix: singular values span 40 : 1.
u, _ = torch.linalg.qr(torch.randn(6, 6))
v, _ = torch.linalg.qr(torch.randn(6, 6))
M = (u * torch.logspace(torch.log10(torch.tensor(40.0)), 0.0, 6)) @ v.T

raw = singular_values(M)
ns = singular_values(newton_schulz(M))
exact = singular_values(orthogonalize(M))

print("raw σ:      ", [f"{s:5.2f}" for s in raw])
print("Newton–Schulz σ:", [f"{s:5.2f}" for s in ns], "  <- loose band around 1")
print("exact UVᵀ σ:", [f"{s:5.2f}" for s in exact], "  <- all exactly 1")

# Same direction as the exact polar factor, computed with matmuls alone:
cos = torch.sum(newton_schulz(M) * orthogonalize(M)) / (newton_schulz(M).norm() * orthogonalize(M).norm())
print(f"\ncos(Newton–Schulz, exact UVᵀ) = {float(cos):.4f}")

# Scale-invariant: orthogonalizing 100·M gives the same matrix.
print("scale-invariant:", torch.allclose(newton_schulz(M), newton_schulz(100 * M), atol=1e-5))
raw σ:       ['40.00', '19.13', ' 9.15', ' 4.37', ' 2.09', ' 1.00']
Newton–Schulz σ: [' 1.13', ' 0.77', ' 0.77', ' 0.74', ' 0.74', ' 0.69']   <- loose band around 1
exact UVᵀ σ: [' 1.00', ' 1.00', ' 1.00', ' 1.00', ' 1.00', ' 1.00']   <- all exactly 1

cos(Newton–Schulz, exact UVᵀ) = 0.9840
scale-invariant: True

The next chart makes the whitening visible on a 32×32 matrix whose singular values span 40 : 1. Orange bars are the singular values of the (normalized) momentum matrix — a steep decay. Blue bars are the same matrix after Newton–Schulz: the whole tail is lifted up to ~1.

The full Muon step

The 2024 original stopped at orthogonalized momentum. Making Muon actually beat AdamW at scale (Moonshot, 2025) took two small, crucial additions — a scale so its updates are the same size as AdamW’s, and decoupled weight decay (the exact trick we built earlier). The whole step, for a weight of shape A \times B:

\begin{aligned} M_t &= \mu\, M_{t-1} + G_t & &\text{(heavy-ball momentum, } \mu = 0.95)\\ O_t &= \text{NewtonSchulz}_5\!\left(G_t + \mu M_t\right) & &\text{(Nesterov look-ahead, then orthogonalize)}\\ W_t &= W_{t-1} - \eta_t\left(\;0.2\,\sqrt{\max(A,B)}\; O_t \;+\; \lambda\, W_{t-1}\right) & &\text{(RMS-matched step + decoupled decay)} \end{aligned}

Why the 0.2\sqrt{\max(A,B)} factor? An orthogonal O has per-entry RMS \approx 1/\sqrt{\max(A,B)}; multiplying by it brings the update’s RMS to \approx 0.2, the typical size of an AdamW update — so the same learning rate you tuned for AdamW transfers straight over. Trace the pipeline:

In code, Muon updates only 2D weight matrices; embeddings, the LM head, biases, and norm scales keep AdamW. route_parameters does the split, and the optimizer itself is a few lines on top of newton_schulz:

from muon import Muon, route_parameters
import torch

# Split a model's parameters the way both Muon papers do.
params = [
    ("blocks.0.attn.qkv.weight", torch.zeros(96, 32)),   # 2D hidden weight -> Muon
    ("blocks.0.mlp.fc.weight",   torch.zeros(128, 32)),  # 2D hidden weight -> Muon
    ("wte.weight",               torch.zeros(1000, 32)), # embedding        -> AdamW
    ("blocks.0.norm.weight",     torch.zeros(32)),       # norm scale (1D)  -> AdamW
]
groups = route_parameters(params)
print("Muon (matrices):", len(groups["muon"]), " AdamW (rest):", len(groups["adamw"]))

# One Muon step, decoupled decay only (zero gradient) shrinks W by exactly lr·λ:
w = torch.ones(4, 4, requires_grad=True)
opt = Muon([w], lr=0.1, weight_decay=0.5)
w.grad = torch.zeros_like(w)
opt.step()
print("after decay-only step:", float(w[0, 0].detach()), "== 1·(1 − 0.1·0.5) = 0.95")
Muon (matrices): 2  AdamW (rest): 2
after decay-only step: 0.949999988079071 == 1·(1 − 0.1·0.5) = 0.95

Does it actually help?

demonstrate_muon runs Muon, AdamW, and plain momentum on the same mildly-ill-conditioned regression, each at a sensible learning rate for itself.

Muon’s single orthogonalized update leaves plain momentum far behind (~0.03 vs ~0.31 final loss) and lands right in AdamW’s ballpark — with no second-moment state at all. It does not beat well-tuned AdamW here, and that is the honest result: on a small convex problem, per-coordinate adaptation is very hard to top, exactly as we said for Adam itself. Muon’s real prize shows up where the loss surface is genuinely deep and non-convex — full LLM pretraining — where Moonlight reached AdamW-level quality with ~2× less compute.

NoteKey Insight

AdamW conditions the update per scalar; Muon conditions it per matrix. By replacing the momentum matrix with its nearest orthogonal matrix — every singular value 1 — Muon advances all directions of a weight equally, without ever storing a second moment. The Newton–Schulz iteration makes that orthogonalization a handful of matmuls instead of an SVD, and the 0.2\sqrt{\max(A,B)} scale lets it borrow AdamW’s learning rate wholesale.

WarningCommon Pitfalls with Muon
  • It’s not a convergent iteration. Newton–Schulz is tuned for ~5 steps and approximately orthogonalizes (singular values ~1, not exactly 1). Iterating it to “convergence” makes things worse, not better.
  • Matrices only. Applying Muon to embeddings, the LM head, biases, or norm scales hurts — those go to AdamW. Muon is for 2D hidden weights.
  • The scale and decay are not optional. The bare 2024 update (no 0.2\sqrt{\max(A,B)}, no weight decay) does not challenge AdamW at scale; both additions are what made Muon competitive in the 2025 results.
  • fp32 here, bf16 in production. We orthogonalize in float32 for clarity; real Muon runs Newton–Schulz in bfloat16, which is where its speed comes from.

Going deeper on Muon: Keller Jordan, Muon: An optimizer for hidden layers of neural networks (2024) — the original, with the Newton–Schulz coefficients and the derivation of why orthogonalized updates help. Liu et al. (Moonshot AI), Muon is Scalable for LLM Training (2025) — the RMS-match scale, decoupled weight decay, and the ~2× compute-efficiency result at 3B/16B scale.

Muon at Scale: QK-Clip and MuonClip

We just said Muon’s real prize is full LLM pretraining. There is a catch that kept “Muon at scale” a promise rather than a recipe until 2025: Muon’s larger, orthogonalized updates inflate attention logits faster than AdamW does, and at frontier scale that inflation runs away into divergence. The fix that made trillion-parameter Muon real is QK-Clip, and Muon + QK-Clip together are the MuonClip optimizer behind Moonshot AI’s Kimi K2 — a 1-trillion-parameter Mixture-of-Experts model (32B active) pretrained on 15.5T tokens with zero loss spikes.

Intuition: the logit Muon inflates

You have already met this failure mode. In Module 5’s QK-Norm section the pre-softmax logit q \cdot k / \sqrt{d_k} grew as the query/key activation norms drifted upward during training, the softmax sharpened toward one-hot, the gradient through it vanished, and the loss diverged — attention entropy collapse. There it was framed as an architectural problem and cured by normalizing the activations. Here it is the same disease with an optimizer-shaped cause: Muon’s whitened update is bigger than AdamW’s per-coordinate step, so it pushes \lVert q \rVert and \lVert k \rVert up faster, and the logits blow up sooner.

QK-Clip takes the opposite tack to QK-Norm. It changes nothing in the forward pass. Instead it watches the logits and, whenever a head’s maximum logit crosses a cap \tau, reaches into that head’s weights after the optimizer step and shrinks them just enough to pull the logit back under \tau. It is a controller, not a layer: once training settles and the logits stop crossing \tau, it stops acting entirely.

NoteTwo cures, one disease

QK-Norm (m05) normalizes q, k every forward pass so the logit is a bounded cosine \in [-g, g] — always on, and it changes the model. QK-Clip (here) leaves the model alone and clips the weights after each step — a post-hoc, per-head, self-deactivating fix. Same explosion; one works in activation space, the other in weight space.

The Math: watch, then clip

After each optimizer step, for each head h measure the largest pre-softmax logit the softmax will see over the batch \mathcal{B} — defined, like every logit in this book, after the 1/\sqrt{d_k} scaling:

S_{\max}^{h} \;=\; \frac{1}{\sqrt{d_k}}\;\max_{X \in \mathcal{B}}\;\max_{i,\,j}\; q_i^{h}\cdot k_j^{h}.

Form a per-head clip factor against a hard cap \tau, then rescale that head’s query and key projections by its square root:

\gamma_h \;=\; \min\!\left(1,\ \frac{\tau}{S_{\max}^{h}}\right), \qquad W_q^{h}\leftarrow \sqrt{\gamma_h}\,W_q^{h}, \quad W_k^{h}\leftarrow \sqrt{\gamma_h}\,W_k^{h}.

The square root is the whole trick. The logit is bilinear in (q, k): scaling W_q^{h} by \sqrt{\gamma_h} scales q^h by \sqrt{\gamma_h}, scaling W_k^{h} by \sqrt{\gamma_h} scales k^h by \sqrt{\gamma_h}, so the logit q^h\cdot k^h scales by exactly \gamma_h. A hot head’s max logit lands at precisely \tau; a cool head (S_{\max}^h \le \tau) gets \gamma_h = 1 and is left untouched. Kimi K2 used \tau = 100.

Walk the four moves one at a time — measure, compare, clip, verify:

Code: QK-Clip from scratch

The from-scratch implementation lives in qk_clip.py. It is deliberately small — four functions and a wrapper — because QK-Clip is small. Start with the two measurements: project inputs to per-head queries and keys, then take the per-head max logit.

import torch
from qk_clip import (
    project_qk, head_max_logit, qk_clip_factors,
    apply_qk_clip, clip_attention_layer, MuonClip, qk_clip_logit_growth,
)

torch.manual_seed(0)
d_model, num_heads, d = 32, 4, 8            # head-major weights: (H*d, d_model)
x = torch.randn(4, 6, d_model)              # a probe batch (B, L, d_model)
w_q = torch.randn(num_heads * d, d_model) * 0.9
w_k = torch.randn(num_heads * d, d_model) * 0.9

q, k = project_qk(x, w_q, w_k, num_heads)
print("per-head max logits:", [round(v, 2) for v in head_max_logit(q, k).tolist()])
per-head max logits: [97.84, 100.51, 78.3, 94.68]

Now the clip itself. qk_clip_factors turns those max logits into per-head factors against a cap \tau, and apply_qk_clip rescales each head’s Q/K rows by \sqrt{\gamma_h} in place:

tau = 90.0
before = head_max_logit(*project_qk(x, w_q, w_k, num_heads))
factors = qk_clip_factors(before, tau)
print("γ per head:", [round(v, 3) for v in factors.tolist()])

apply_qk_clip(w_q, w_k, factors, num_heads)
after = head_max_logit(*project_qk(x, w_q, w_k, num_heads))
print("before:", [round(v, 2) for v in before.tolist()])
print("after: ", [round(v, 2) for v in after.tolist()])
print(f"max logit now ≤ τ={tau}? ", bool((after <= tau + 1e-4).all()))
γ per head: [0.92, 0.895, 1.0, 0.951]
before: [97.84, 100.51, 78.3, 94.68]
after:  [90.0, 90.0, 78.3, 90.0]
max logit now ≤ τ=90.0?  True

Every head that was over \tau now sits at exactly \tau; every head that was already under it is unchanged. clip_attention_layer bundles the three moves (measure → factor → apply) into one call, and MuonClip wraps a Muon optimizer so each .step() runs the optimizer and then clips the registered attention layers:

torch.manual_seed(1)
w_q = torch.randn(num_heads * d, d_model, requires_grad=True)
w_k = torch.randn(num_heads * d, d_model, requires_grad=True)
from muon import Muon

opt = MuonClip(Muon([w_q, w_k], lr=0.0), tau=3.0)   # lr=0 isolates the clip phase
opt.register_attention(w_q, w_k, num_heads=num_heads, probe=x)

with torch.no_grad():                                # force the layer hot
    w_q *= 2.5; w_k *= 2.5
w_q.grad = torch.zeros_like(w_q); w_k.grad = torch.zeros_like(w_k)

hot = head_max_logit(*project_qk(x, w_q, w_k, num_heads)).max()
opt.step()
cool = head_max_logit(*project_qk(x, w_q, w_k, num_heads)).max()
print(f"max logit {hot:.1f} → {cool:.1f} after one MuonClip step (τ=3.0)")
max logit 990.7 → 3.0 after one MuonClip step (τ=3.0)
NoteKey Insight

QK-Clip is a feedback controller on the attention logits. It never enters the forward pass — it reads S_{\max}^h, and when a head crosses \tau it multiplies that head’s Q and K weights by \sqrt{\gamma_h} so the bilinear logit scales by \gamma_h and lands on \tau. Cool heads see \gamma_h = 1, so as training stabilizes and the logits stop crossing \tau, QK-Clip quietly turns itself off.

The real payoff is what happens across many inflating steps. qk_clip_logit_growth inflates a Q/K pair each step (a stand-in for repeated Muon updates) and traces the max logit with and without the clip:

rows = qk_clip_logit_growth(steps=14, tau=8.0, inflation=1.30, seed=0)
ojs_define(
    qkGrowthSteps = [r["step"] for r in rows],
    qkGrowthUncapped = [r["uncapped"] for r in rows],
    qkGrowthClipped = [r["clipped"] for r in rows],
    qkGrowthTau = 8.0,
)

The uncapped logit climbs a straight line on the log axis — geometric blow-up, exactly the runaway that ends a large Muon run. The QK-Clipped copy rises until it touches \tau, then rides it: every step still inflates the weights, and every step QK-Clip pulls them back.

Interactive Exploration

Drive the cap \tau and the per-step inflation yourself. Watch where the clipped curve pins, and how a tighter cap or a hotter optimizer changes when — and how often — QK-Clip has to fire.

TipTry This
  1. Lower the cap. Drag \tau down. The clipped curve pins lower and the clip fires on more steps — a tighter cap is a busier controller.
  2. Heat up the optimizer. Raise the inflation. The uncapped curve steepens and diverges sooner; QK-Clip still pins the clipped curve at \tau.
  3. Find the quiet regime. Lower inflation until the clip stops firing entirely (no dots). That is the self-deactivated state real runs settle into — the logits never cross \tau, so QK-Clip is a no-op.
NoteIn Kimi K2: MLA, decoupled RoPE, and MoE

Kimi K2 does not use vanilla multi-head attention — it uses MLA (Module 9) with decoupled RoPE (Module 10). The clean \sqrt{\gamma} split we built is the standard-MHA rule; Kimi K2 splits it by component: the non-RoPE content projections q^{C}, k^{C} are scaled by \sqrt{\gamma_h}, the per-head rotary query q^{R} takes the full \gamma_h, and the shared rotary key k^{R} is left untouched — clipping a shared key would perturb every other head. QK-Clip’s role in “MuonClip” is exactly this: it is what let Muon’s ~2× token efficiency scale from speedruns and a 16B MoE all the way to a 1T-parameter MoE with no loss spikes.

WarningCommon Pitfalls with QK-Clip
  • Clip after the optimizer step, not before. QK-Clip corrects what the step just did; running it first measures stale logits and clips nothing useful.
  • Only the hot heads move. \gamma_h = \min(1, \tau/S_{\max}^h) is 1 for a cool head — clipping every head uniformly would needlessly shrink healthy attention.
  • \tau too low throttles capacity. The cap is a safety rail, not a target. Set it well above the logits a healthy run produces (Kimi K2 used \tau = 100); too tight and you cap attention the model actually needs.
  • It’s a controller, not a loss term. QK-Clip has no gradient and adds nothing to the objective — it edits weights between steps. Don’t try to backprop through it.
  • The MHA rule is not the MLA rule. The \sqrt{\gamma}-on-both-projections split is for standard attention; MLA with decoupled RoPE splits it by component (see the aside above).

Going deeper on MuonClip: Kimi Team (Moonshot AI), Kimi K2: Open Agentic Intelligence (2025) — the MuonClip optimizer, the QK-Clip rule and threshold, and the zero-loss-spike 1T-parameter pretraining run.

Gradient Accumulation

Gradient accumulation increases effective batch size without adding memory.

Problem: Want batch_size=32 but only 8 fits in memory

Solution: Accumulate gradients over 4 mini-batches

# Demonstrate gradient accumulation
model = nn.Linear(10, 1)
accumulation_steps = 4

# Simulate accumulated gradients
total_loss = 0

for i in range(accumulation_steps):
    x = torch.randn(8, 10)  # Mini-batch
    y = model(x)
    loss = y.mean() / accumulation_steps  # Scale loss!
    loss.backward()  # Gradients accumulate
    total_loss += loss.item()

print(f"Accumulated loss (4 mini-batches): {total_loss:.4f}")
print(f"Gradient norm before step: {model.weight.grad.norm().item():.4f}")

# Now do one optimizer step
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer.step()
optimizer.zero_grad()

print("After optimizer.step() and zero_grad()")
Accumulated loss (4 mini-batches): 0.0813
Gradient norm before step: 0.7052
After optimizer.step() and zero_grad()

Gradient Clipping

Gradient clipping scales down gradients whose norm exceeds a threshold, preventing gradient explosion.

# Demonstrate gradient clipping
model = nn.Linear(10, 10)

# Create artificial large gradients
for p in model.parameters():
    p.grad = torch.randn_like(p) * 100  # Very large!

# Compute gradient norm before clipping
total_norm_before = 0
for p in model.parameters():
    total_norm_before += p.grad.norm().item() ** 2
total_norm_before = total_norm_before ** 0.5

print(f"Gradient norm before clipping: {total_norm_before:.2f}")

# Clip gradients
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# Compute gradient norm after
total_norm_after = 0
for p in model.parameters():
    total_norm_after += p.grad.norm().item() ** 2
total_norm_after = total_norm_after ** 0.5

print(f"Gradient norm after clipping:  {total_norm_after:.2f}")
print(f"\nGradients scaled down by {total_norm_before / total_norm_after:.1f}x")
Gradient norm before clipping: 1023.26
Gradient norm after clipping:  1.00

Gradients scaled down by 1023.3x

Gradient Clipping from Scratch

Let’s implement gradient clipping ourselves to understand the algorithm:

def clip_grad_norm_scratch(params, max_norm: float) -> float:
    """
    Clip gradients by global norm.

    Algorithm:
    1. Compute total norm: sqrt(sum of all grad^2)
    2. If total_norm > max_norm, scale all grads by (max_norm / total_norm)

    Returns the original norm (before clipping).
    """
    params = list(params)

    # Step 1: Compute total gradient norm
    total_sq = 0.0
    for p in params:
        if p.grad is not None:
            total_sq += (p.grad ** 2).sum().item()
    total_norm = total_sq ** 0.5

    # Step 2: Clip if needed
    if total_norm > max_norm:
        scale = max_norm / (total_norm + 1e-12)  # Small epsilon for numerical stability
        for p in params:
            if p.grad is not None:
                p.grad *= scale

    return total_norm

# Test: compare with PyTorch
model_scratch = nn.Linear(10, 10)
model_pytorch = nn.Linear(10, 10)

# Set same large gradients
torch.manual_seed(42)
for p in model_scratch.parameters():
    p.grad = torch.randn_like(p) * 100
for ps, pp in zip(model_scratch.parameters(), model_pytorch.parameters()):
    pp.grad = ps.grad.clone()

# Clip with both
norm_scratch = clip_grad_norm_scratch(model_scratch.parameters(), max_norm=1.0)
norm_pytorch = torch.nn.utils.clip_grad_norm_(model_pytorch.parameters(), max_norm=1.0)

print(f"Original norm (scratch): {norm_scratch:.4f}")
print(f"Original norm (PyTorch): {norm_pytorch.item():.4f}")

# Check gradients match after clipping
grads_match = all(
    torch.allclose(ps.grad, pp.grad)
    for ps, pp in zip(model_scratch.parameters(), model_pytorch.parameters())
)
print(f"Gradients match after clipping: {grads_match}")
Original norm (scratch): 1037.5083
Original norm (PyTorch): 1037.5083
Gradients match after clipping: True
NoteKey Insight: Gradient Clipping

Gradient clipping scales ALL gradients by the same factor to preserve their relative magnitudes. This is different from clipping each gradient independently - we want to maintain the direction of the overall update while limiting its magnitude.

When to use gradient clipping:

  • Always for transformer training (standard practice)
  • max_norm=1.0 is a good default
  • Monitor gradient norms during training - consistently high norms suggest instability

Batch Size Considerations

Batch size affects both training dynamics and memory usage:

Tradeoffs:

Aspect Small Batch Large Batch
Memory Less More
Gradient noise More (regularization effect) Less (stable gradients)
Convergence May generalize better Faster convergence
LR needed Lower Higher (linear scaling rule)

The Linear Scaling Rule: When you double the batch size, you can double the learning rate. This maintains similar training dynamics.

Effective batch size = batch_size x gradient_accumulation_steps

# Batch size vs memory example (conceptual)
print("Memory usage scales linearly with batch size:")
print()
for batch_size in [8, 16, 32, 64]:
    # Simulated memory calculation
    tokens_per_batch = batch_size * 512  # sequence length
    memory_mb = batch_size * 50  # ~50MB per sample for a small model
    print(f"  Batch size {batch_size:2d}: ~{tokens_per_batch:,} tokens/batch, ~{memory_mb}MB")
Memory usage scales linearly with batch size:

  Batch size  8: ~4,096 tokens/batch, ~400MB
  Batch size 16: ~8,192 tokens/batch, ~800MB
  Batch size 32: ~16,384 tokens/batch, ~1600MB
  Batch size 64: ~32,768 tokens/batch, ~3200MB

Mixed Precision & Numerics

Everything so far assumed a comfortable 32-bit float (fp32). But modern models train in 16-bit — and increasingly 8-bit — because the arithmetic is 2–8× faster and every tensor is half the size (or less). The catch is that a narrower float represents fewer numbers, and two failure modes follow:

  • Overflow — a value too large for the format becomes inf, and any math touching it turns to NaN. One NaN poisons the whole update.
  • Underflow — a value too small rounds all the way to 0, so a real gradient silently vanishes and that weight never learns.

To see exactly where those walls are, we build the floating-point number line from scratch in precision.py — a small IEEE-754-style encoder that rounds any Python float into a chosen (exponent, mantissa) layout, the same rounding a GPU does when it stores a number in 16 bits. It reproduces fp32, fp16, and bf16 exactly (the tests check it against struct and PyTorch).

Anatomy of a float

A binary float spends its bits on two jobs: exponent bits buy range (how big and how small), mantissa bits buy precision (how many significant digits). A 16-bit budget forces a choice, and the two 16-bit formats split it oppositely:

from precision import FORMATS, format_spec

for name in ("fp32", "fp16", "bf16"):
    s = format_spec(FORMATS[name])
    print(f"{s['name']:5}  exp={s['exp_bits']:>2}  mantissa={s['mantissa_bits']:>2}  "
          f"max={s['max_normal']:.3e}  min_normal={s['min_normal']:.3e}  "
          f"eps={s['eps']:.2e}  (~{s['decimal_digits']:.1f} decimal digits)")
fp32   exp= 8  mantissa=23  max=3.403e+38  min_normal=1.175e-38  eps=1.19e-07  (~6.9 decimal digits)
fp16   exp= 5  mantissa=10  max=6.550e+04  min_normal=6.104e-05  eps=9.77e-04  (~3.0 decimal digits)
bf16   exp= 8  mantissa= 7  max=3.390e+38  min_normal=1.175e-38  eps=7.81e-03  (~2.1 decimal digits)

Read the two 16-bit rows against each other: fp16 spends 10 bits on the mantissa (fine precision) but only 5 on the exponent (a tiny range — it caps at 65504). bf16 keeps fp32’s full 8 exponent bits (so it barely ever overflows or underflows) at the cost of just 7 mantissa bits (coarse precision). That one trade is the whole story of mixed-precision training.

Step through some values and watch how each format stores them — where the bits go, what value comes back, and when a value hits a wall:

from precision import FORMATS, encode, round_to_format, rounding_error, overflows, underflows

def _bit_row(x, fmt):
    import math as _m
    bits = encode(x, fmt)
    s = format(bits, f"0{fmt.total_bits}b")
    r = round_to_format(x, fmt)
    return {
        "sign": s[0],
        "exp": s[1:1 + fmt.exp_bits],
        "mant": s[1 + fmt.exp_bits:],
        "stored": r if _m.isfinite(r) else None,   # null in JS; overflow flag drives display
        "error": min(rounding_error(x, fmt), 9.99),
        "overflow": overflows(x, fmt),
        "underflow": underflows(x, fmt),
    }

_demo = [0.1, 1.0, 3.1415927, 0.15625, 65504.0, 70000.0, 6.1e-5, 1e-8]
precisionData = [
    {"value": x,
     "fp16": _bit_row(x, FORMATS["fp16"]),
     "bf16": _bit_row(x, FORMATS["bf16"])}
    for x in _demo
]
ojs_define(precisionData = precisionData)
NoteKey Insight

Same 16 bits, opposite bets. fp16 gives the mantissa 10 bits and the exponent 5 — precise, but it overflows at 65504 and underflows around 6\times10^{-5}. bf16 gives the exponent 8 bits (fp32’s whole range) and the mantissa only 7 — it almost never over/underflows, but every stored number is coarser. Step to 70000 and watch fp16 flip to inf while bf16 holds it; step to 1e-8 and watch fp16 collapse to 0.

Failure mode 1: overflow

70000 is a perfectly ordinary number, but it is past fp16’s largest value (65504). Storing it there gives inf — and inf - inf = NaN a few operations later. bf16, with fp32’s exponent range, does not blink:

from precision import round_to_format, FORMATS

for name in ("fp16", "bf16"):
    print(f"70000 in {name}: {round_to_format(70000.0, FORMATS[name])}")
70000 in fp16: inf
70000 in bf16: 70144.0

Failure mode 2: underflow (and the loss-scaling fix)

The subtler killer is underflow. Late in training, gradients get small — and a gradient like 1e-8 is below fp16’s smallest representable value, so it rounds to 0 and that weight simply stops updating. The classic fix is loss scaling: multiply the loss by a constant S before backprop, which multiplies every gradient by S, lifting it out of the underflow hole. The optimizer then divides the update back out by S, so the math is unchanged — only the representation was rescued.

from precision import loss_scale_gradient, FORMATS

out = loss_scale_gradient(grad=1e-8, scale=1024.0, fmt=FORMATS["fp16"])
print(f"true gradient:            {out['grad']:g}")
print(f"stored directly in fp16:  {out['naive_stored']}   underflowed={out['naive_underflowed']}")
print(f"x{out['scale']:g}, store, unscale:   {out['recovered']:g}   rescued={out['rescued']}")
true gradient:            1e-08
stored directly in fp16:  0.0   underflowed=True
x1024, store, unscale:   1.00117e-08   rescued=True

Multiplying by S = 1024 turned an unrepresentable 1e-8 into ~1e-5 (safely inside fp16), and unscaling recovered the gradient almost exactly. This is what PyTorch’s GradScaler automates — and it is why bf16, which underflows far less, usually needs no loss scaling at all.

The whole number line at a glance

The range ladder below shows why the choice comes out the way it does: each bar runs from a format’s smallest subnormal to its largest value (log scale). fp16 is a narrow window; bf16 spans essentially all of fp32. The two markers are the problem values above — 70000 sits past fp16’s right wall, 1e-8 past its left.

from precision import FORMATS, format_spec

_ranges = []
for name in ("fp32", "fp16", "bf16"):
    s = format_spec(FORMATS[name])
    _ranges.append({"format": name, "lo": s["min_subnormal"], "hi": s["max_normal"]})
ojs_define(precisionRanges = _ranges)
TipTry This

In the bit stepper, compare 0.1 across the two formats: neither is exact (0.1 is not a finite binary fraction), but fp16’s error is ~10× smaller than bf16’s — that is the extra 3 mantissa bits. Then jump to 70000 and 1e-8: fp16 hits a wall at both ends, bf16 at neither. That asymmetry — bf16 trades precision you rarely miss for range you can’t afford to lose — is why bf16 became the default 16-bit training format on A100/H100/TPU hardware.

Putting it together: the mixed-precision recipe

“Mixed” precision means using both: fast low-precision math for the heavy forward/backward pass, full-precision fp32 for the parts that must be accurate.

  1. Keep a master copy of the weights in fp32.
  2. Cast to fp16/bf16 for the forward and backward pass (the expensive part, now 2× faster).
  3. If using fp16, apply loss scaling so small gradients survive.
  4. Apply the update to the fp32 master weights (accurate accumulation), then re-cast.

In PyTorch this is a few lines of autocast + GradScaler:

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()                       # manages loss scaling automatically
for batch in dataloader:
    optimizer.zero_grad()
    with autocast(dtype=torch.bfloat16):    # fp16/bf16 forward+backward
        loss = F.cross_entropy(model(input_ids), targets)
    scaler.scale(loss).backward()           # scale up -> gradients survive fp16
    scaler.step(optimizer)                  # unscale -> fp32 master-weight update
    scaler.update()                         # adapt the scale factor

Practical advice:

  • Prefer bf16 where the hardware supports it (A100/H100/TPU) — its fp32-sized range means no loss scaling and far fewer NaN surprises.
  • Use fp16 + loss scaling on older GPUs (e.g. V100) that lack bf16.
  • fp8 (e4m3 / e5m2) is the frontier — H100-class hardware trains parts of the network in 8 bits, with per-tensor scaling doing the loss-scaling job the number line above makes necessary. The next section builds both fp8 formats and that scaling from scratch.

Going to 8 bits: fp8 and per-tensor scaling

The same encoder that gave us fp16 and bf16 also gives us fp8 — the frontier of low-precision training, run natively on H100/H200 tensor cores and used to train DeepSeek-V3 (671B parameters, 2024) end to end. Eight bits is so few that the range-vs-precision trade can no longer be settled with a single format, so fp8 ships as two:

  • e4m3 — 4 exponent, 3 mantissa bits. More precision, less range (caps at 448). Used for the forward pass: weights and activations, which sit in a fairly narrow band once normalized.
  • e5m2 — 5 exponent, 2 mantissa bits. fp16’s exponent, so fp16’s range (caps at 57344) but coarser. Used for gradients, which span many orders of magnitude — the same reason bf16 exists, taken to 8 bits.

Both are already in precision.py as FP8_FORMATS:

from precision import FP8_FORMATS, format_spec

for name in ("e4m3", "e5m2"):
    s = format_spec(FP8_FORMATS[name])
    print(f"{s['name']:5}  exp={s['exp_bits']}  mantissa={s['mantissa_bits']}  "
          f"max={s['max_normal']:>7.0f}  min_sub={s['min_subnormal']:.2e}  "
          f"eps={s['eps']:.3f}  inf={s['has_inf']}")
e4m3   exp=4  mantissa=3  max=    448  min_sub=1.95e-03  eps=0.125  inf=False
e5m2   exp=5  mantissa=2  max=  57344  min_sub=1.53e-05  eps=0.250  inf=True

e4m3 breaks the IEEE rules

Notice inf=False for e4m3. This is the one genuinely new idea in fp8, and it is worth pausing on. Every format so far — fp32, fp16, bf16, e5m2 — is IEEE-754: the all-ones exponent is reserved to mean inf (mantissa 0) or NaN (mantissa nonzero), and no normal number is allowed to use it. That costs a whole binade of range at the top.

With only 4 exponent bits, e4m3 cannot afford that. So it reclaims the all-ones exponent for normal numbers and keeps exactly one bit pattern for NaN — S.1111.111 — and no infinity at all (this is what PyTorch calls float8_e4m3fn, the “fn” = finite). Reclaiming that exponent is why e4m3 reaches 448 instead of the 240 an IEEE-style e4m3 would cap at. The trade: a value past the top no longer rounds to a harmless inf — it becomes NaN.

Our encoder handles this with a single has_inf flag; step through some values and watch e4m3 saturate to 448 and then go NaN while e5m2 rounds to inf like fp16 does:

from precision import FP8_FORMATS, encode, round_to_format, rounding_error, overflows, underflows

def _fp8_row(x, fmt):
    import math as _m
    bits = encode(x, fmt)
    s = format(bits, f"0{fmt.total_bits}b")
    r = round_to_format(x, fmt)
    return {
        "sign": s[0],
        "exp": s[1:1 + fmt.exp_bits],
        "mant": s[1 + fmt.exp_bits:],
        "stored": r if _m.isfinite(r) else None,
        "isnan": _m.isnan(r),
        "isinf": _m.isinf(r),
        "error": min(rounding_error(x, fmt), 9.99),
        "overflow": overflows(x, fmt),
        "underflow": underflows(x, fmt),
    }

_fp8_demo = [1.0, 0.1, 0.0018, 6.1e-4, 448.0, 500.0, 57344.0, 1e5]
fp8Data = [
    {"value": x,
     "e4m3": _fp8_row(x, FP8_FORMATS["e4m3"]),
     "e5m2": _fp8_row(x, FP8_FORMATS["e5m2"])}
    for x in _fp8_demo
]
ojs_define(fp8Data = fp8Data)
NoteKey Insight

e4m3 has just 4 exponent bits, so it drops IEEE’s reserved-infinity rule to buy range: no inf, one NaN, and a top exponent that holds normal numbers up to 448. Step to 500 and watch e4m3 flip to NaN (there is no inf to catch it) while e5m2 — with fp16’s 5 exponent bits — still holds 512. Step to 57344 and only e5m2 survives. This is why e4m3 does the forward pass (narrow, precise) and e5m2 does the gradients (wide, coarse).

Per-tensor scaling: the fp8 loss-scaling

There is a second problem 8 bits creates. e4m3’s smallest normal value is about 0.016, and its smallest subnormal about 0.002 — but a trained network’s weights are routinely ~0.001–0.05. Quantize such a tensor to e4m3 directly and much of it underflows to zero, exactly the failure loss scaling fixed for fp16 gradients. The fix is the same idea, applied per tensor: compute one scale that maps the tensor’s largest magnitude (amax) onto max_normal, multiply the whole tensor by it so it fills the 8-bit range, quantize, and divide the scale back out in the matmul.

from precision import quantize_per_tensor, per_tensor_scale, FP8_FORMATS

weights = [0.031, -0.012, 0.004, 0.0018, -0.0009, 0.021, -0.006, 0.0003]
e4m3 = FP8_FORMATS["e4m3"]

scale = per_tensor_scale(weights, e4m3)      # 448 / amax(weights)
out = quantize_per_tensor(weights, e4m3)
print(f"per-tensor scale:        {scale:,.0f}  (maps amax {max(map(abs, weights))} -> 448)")
print(f"naive fp8:   {out['naive_zeroed']}/{len(weights)} entries underflow to 0, "
      f"mean rel error {out['naive_rel_error']:.1%}")
print(f"scaled fp8:  {out['scaled_zeroed']}/{len(weights)} entries underflow to 0, "
      f"mean rel error {out['scaled_rel_error']:.1%}")
per-tensor scale:        14,452  (maps amax 0.031 -> 448)
naive fp8:   2/8 entries underflow to 0, mean rel error 27.3%
scaled fp8:  0/8 entries underflow to 0, mean rel error 1.9%

Without scaling, the two smallest weights vanish and the average entry is off by a quarter of its value. Scale the tensor first and nothing underflows — the whole vector survives at true 8-bit precision. The bars below show each weight before (target), after naive quantization (some collapse to 0), and after scaled quantization (all recovered):

from precision import quantize_per_tensor, per_tensor_scale, FP8_FORMATS

_w = [0.031, -0.012, 0.004, 0.0018, -0.0009, 0.021, -0.006, 0.0003]
_e4m3 = FP8_FORMATS["e4m3"]
_out = quantize_per_tensor(_w, _e4m3)
fp8Scaling = [
    {"i": i, "target": _w[i], "naive": _out["naive"][i], "scaled": _out["recovered"][i]}
    for i in range(len(_w))
]
ojs_define(fp8Scaling = fp8Scaling)
ojs_define(fp8Scale = per_tensor_scale(_w, _e4m3))
TipTry This

In the bit stepper, step to 0.0018 and 6.1e-4: e4m3 stores the first as a tiny subnormal but rounds the second to 0 — that is the underflow the scaling above rescues. Real fp8 training goes further still: DeepSeek-V3 uses fine-grained scaling (a separate scale per 1×128 tile of activations and per 128×128 block of weights) so one giant outlier can’t force the whole tensor’s scale down and drown everything else. Same idea, finer grain.

Going deeper on fp8: FP8 Formats for Deep Learning (Micikevicius et al., 2022) defines e4m3/e5m2 and the forward-vs-gradient split; the OCP Microscaling (MX) spec (2023) standardizes them; and the DeepSeek-V3 technical report (2024) shows fp8 mixed-precision training of a 671B model with fine-grained per-tile/per-block scaling.

Distributed Training Basics

Large models require multiple GPUs. A brief overview:

Data Parallel (DP/DDP):

  • Same model copied to all GPUs
  • Each GPU processes different data
  • Gradients are averaged across GPUs
  • Memory per GPU = full model size

Fully Sharded Data Parallel (FSDP):

  • Model is sharded across GPUs
  • Each GPU holds a fraction of parameters
  • Memory per GPU = model_size / num_gpus
  • Enables training models larger than single GPU memory
# Distributed training concepts
print("Distributed Training Strategies:")
print()
print("1. Data Parallel (DDP):")
print("   - Best for: Models that fit in one GPU")
print("   - Scales: Batch size (effective_batch = batch * num_gpus)")
print()
print("2. Fully Sharded Data Parallel (FSDP):")
print("   - Best for: Large models (>10B parameters)")
print("   - Scales: Model size and batch size")
print()
print("3. Pipeline Parallel:")
print("   - Best for: Very deep models")
print("   - Splits model layers across GPUs")
print()
print("4. Tensor Parallel:")
print("   - Best for: Models with large layers")
print("   - Splits individual layers across GPUs")
Distributed Training Strategies:

1. Data Parallel (DDP):
   - Best for: Models that fit in one GPU
   - Scales: Batch size (effective_batch = batch * num_gpus)

2. Fully Sharded Data Parallel (FSDP):
   - Best for: Large models (>10B parameters)
   - Scales: Model size and batch size

3. Pipeline Parallel:
   - Best for: Very deep models
   - Splits model layers across GPUs

4. Tensor Parallel:
   - Best for: Models with large layers
   - Splits individual layers across GPUs

Training Stability and Failure Modes

Understanding common failure modes helps you debug training issues:

Loss = NaN or Inf

Causes: - Learning rate too high - Gradient explosion - Numerical overflow in fp16

Solutions: - Reduce learning rate (try 10x smaller) - Add gradient clipping - Use bf16 instead of fp16 or add gradient scaling

Loss stuck at high value

Causes: - Learning rate too low - Poor weight initialization - Data loading bug (same batch every time)

Solutions: - Increase learning rate - Check data loader with small sample - Verify model architecture

Loss oscillates or increases

Causes: - Learning rate too high - Batch size too small - Bug in loss computation

Solutions: - Add warmup period - Reduce learning rate - Use gradient accumulation

Debugging checklist:

  1. Check initial loss - should be ~log(vocab_size) for untrained model
  2. Verify data is being loaded correctly (print a few samples)
  3. Monitor gradient norms - should be stable, not growing
  4. Check learning rate schedule is working (print LR each step)
  5. Test with a tiny dataset first to verify overfitting capability

Text Dataset

Let’s create a simple dataset for language modeling:

from torch.utils.data import Dataset, DataLoader

class TextDataset(Dataset):
    """Simple text dataset for language modeling."""

    def __init__(self, tokens, seq_len):
        self.tokens = tokens
        self.seq_len = seq_len

    def __len__(self):
        return max(0, len(self.tokens) - self.seq_len)

    def __getitem__(self, idx):
        input_ids = self.tokens[idx:idx + self.seq_len]
        targets = self.tokens[idx + 1:idx + self.seq_len + 1]
        return input_ids, targets

# Create a simple dataset
tokens = torch.arange(100)  # Token IDs 0-99
seq_len = 8

dataset = TextDataset(tokens, seq_len=seq_len)

print(f"Token IDs: {tokens[:20].tolist()}...")
print(f"Sequence length: {seq_len}")
print(f"Number of samples: {len(dataset)}")
Token IDs: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]...
Sequence length: 8
Number of samples: 92
# Look at a sample
input_ids, targets = dataset[0]

print("Sample 0:")
print(f"  Input:  {input_ids.tolist()}")
print(f"  Target: {targets.tolist()}")
print(f"\n  Target is input shifted by 1 position!")

# Another sample
input_ids, targets = dataset[50]
print(f"\nSample 50:")
print(f"  Input:  {input_ids.tolist()}")
print(f"  Target: {targets.tolist()}")
Sample 0:
  Input:  [0, 1, 2, 3, 4, 5, 6, 7]
  Target: [1, 2, 3, 4, 5, 6, 7, 8]

  Target is input shifted by 1 position!

Sample 50:
  Input:  [50, 51, 52, 53, 54, 55, 56, 57]
  Target: [51, 52, 53, 54, 55, 56, 57, 58]

Training a Model

Now let’s put it all together and train a tiny model:

import sys
sys.path.insert(0, '..')
from m06_transformer.transformer import create_gpt_tiny

# Create model and data
torch.manual_seed(42)

vocab_size = 100
model = create_gpt_tiny(vocab_size=vocab_size)

# Random "training data"
tokens = torch.randint(0, vocab_size, (5000,))

print(f"Model: {model.num_params:,} parameters")
print(f"Training data: {len(tokens):,} tokens")
Model: 838,912 parameters
Training data: 5,000 tokens
# Check initial loss (should be ~log(vocab_size) for random predictions)
dataset = TextDataset(tokens, seq_len=32)
input_ids, targets = dataset[0]
input_ids = input_ids.unsqueeze(0)  # Add batch dimension
targets = targets.unsqueeze(0)

model.eval()
with torch.no_grad():
    logits = model(input_ids)
    # Reshape for loss computation
    B, T, V = logits.shape
    initial_loss = F.cross_entropy(logits.view(B*T, V), targets.view(B*T))

print(f"Initial loss: {initial_loss.item():.4f}")
print(f"Initial perplexity: {math.exp(initial_loss.item()):.2f}")
print(f"\nExpected for random guessing: loss ~ {np.log(vocab_size):.2f}, ppl ~ {vocab_size}")
Initial loss: 4.6686
Initial perplexity: 106.55

Expected for random guessing: loss ~ 4.61, ppl ~ 100
def train_model(model, tokens, num_steps=100, batch_size=16, seq_len=32, learning_rate=3e-4):
    """Simple training loop."""
    dataset = TextDataset(tokens, seq_len)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
    scheduler = CosineScheduler(optimizer, warmup_steps=10, total_steps=num_steps, min_lr=1e-5)

    model.train()
    losses = []
    step = 0

    while step < num_steps:
        for input_ids, targets in dataloader:
            if step >= num_steps:
                break

            # Forward pass
            logits = model(input_ids)
            B, T, V = logits.shape
            loss = F.cross_entropy(logits.view(B*T, V), targets.view(B*T))

            # Backward pass
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

            # Update
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()

            losses.append(loss.item())

            if step % 10 == 0:
                lr = optimizer.param_groups[0]['lr']
                ppl = math.exp(loss.item())
                print(f"Step {step:3d} | Loss: {loss.item():.4f} | PPL: {ppl:.2f} | LR: {lr:.2e}")

            step += 1

    return losses

# Train!
print("Starting training...\n")
losses = train_model(model, tokens, num_steps=100)
print(f"\nFinal loss: {losses[-1]:.4f}")
print(f"Final perplexity: {math.exp(losses[-1]):.2f}")
Starting training...

Step   0 | Loss: 4.6194 | PPL: 101.44 | LR: 0.00e+00
Step  10 | Loss: 4.6242 | PPL: 101.92 | LR: 3.00e-04
Step  20 | Loss: 4.6074 | PPL: 100.23 | LR: 2.91e-04
Step  30 | Loss: 4.5950 | PPL: 98.99 | LR: 2.66e-04
Step  40 | Loss: 4.5934 | PPL: 98.83 | LR: 2.27e-04
Step  50 | Loss: 4.6067 | PPL: 100.16 | LR: 1.80e-04
Step  60 | Loss: 4.5851 | PPL: 98.02 | LR: 1.30e-04
Step  70 | Loss: 4.5777 | PPL: 97.29 | LR: 8.25e-05
Step  80 | Loss: 4.6026 | PPL: 99.74 | LR: 4.39e-05
Step  90 | Loss: 4.5959 | PPL: 99.08 | LR: 1.87e-05

Final loss: 4.5968
Final perplexity: 99.17
# Pass losses to OJS for visualization
import json
ojs_define(training_losses = losses, vocab_size_val = vocab_size)

Effect of Learning Rate

Learning rate is crucial - too high causes instability, too low is slow:

# Train with different learning rates
learning_rates = [1e-5, 1e-4, 3e-4, 1e-3, 3e-3]
all_losses = {}

for lr in learning_rates:
    torch.manual_seed(42)
    model = create_gpt_tiny(vocab_size=100)
    tokens = torch.randint(0, 100, (3000,))

    # Train silently
    dataset = TextDataset(tokens, seq_len=32)
    dataloader = DataLoader(dataset, batch_size=8, shuffle=True, drop_last=True)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)

    model.train()
    losses = []
    step = 0

    while step < 50:
        for input_ids, targets in dataloader:
            if step >= 50:
                break
            logits = model(input_ids)
            B, T, V = logits.shape
            loss = F.cross_entropy(logits.view(B*T, V), targets.view(B*T))
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            losses.append(loss.item())
            step += 1

    all_losses[lr] = losses
    print(f"LR={lr:.0e}: final_loss={losses[-1]:.3f}, final_ppl={math.exp(losses[-1]):.1f}")
LR=1e-05: final_loss=4.637, final_ppl=103.2
LR=1e-04: final_loss=4.615, final_ppl=101.0
LR=3e-04: final_loss=4.613, final_ppl=100.8
LR=1e-03: final_loss=4.639, final_ppl=103.4
LR=3e-03: final_loss=4.634, final_ppl=103.0
# Pass learning rate comparison data to OJS
# Convert dict with float keys to list of dicts for JSON serialization
lr_comparison_data = [{"lr": lr, "losses": losses} for lr, losses in all_losses.items()]
ojs_define(lr_comparison = lr_comparison_data)

Observations:

  • Too low (1e-5): Training is very slow
  • Just right (3e-4): Smooth, fast convergence
  • Too high (3e-3): Unstable, loss may spike or diverge

Checkpointing

Save regularly! Training can crash. Here’s what to save:

# Demonstrate checkpointing
import json
from pathlib import Path

def save_checkpoint(model, optimizer, step, loss, path):
    """Save a training checkpoint."""
    checkpoint = {
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'step': step,
        'loss': loss,
    }
    torch.save(checkpoint, path)
    print(f"Checkpoint saved to {path}")

def load_checkpoint(model, optimizer, path):
    """Load a training checkpoint."""
    checkpoint = torch.load(path, weights_only=False)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    print(f"Checkpoint loaded from {path}")
    print(f"  Step: {checkpoint['step']}, Loss: {checkpoint['loss']:.4f}")
    return checkpoint['step'], checkpoint['loss']

# Save example
model = create_gpt_tiny(vocab_size=100)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
save_checkpoint(model, optimizer, step=50, loss=2.5, path="demo_checkpoint.pt")

# Load example
model2 = create_gpt_tiny(vocab_size=100)
optimizer2 = torch.optim.AdamW(model2.parameters(), lr=3e-4)
step, loss = load_checkpoint(model2, optimizer2, "demo_checkpoint.pt")

# Clean up
Path("demo_checkpoint.pt").unlink()
Checkpoint saved to demo_checkpoint.pt
Checkpoint loaded from demo_checkpoint.pt
  Step: 50, Loss: 2.5000

Validation and Early Stopping

Monitor validation loss to detect overfitting:

Tips: - Monitor validation loss, not just training loss - Save the model with the best validation loss - Consider early stopping if validation loss increases consistently

Scaling Laws: Compute-Optimal Training

Everything so far answers how to train. This section answers the question that comes before the first step: given a fixed compute budget, how big should the model be, and how many tokens should it see? The surprising answer is that this is predictable — model loss follows smooth power laws in size, data, and compute, so you can plan the run instead of guessing.

The Compute Rule: C ≈ 6ND

The compute (in FLOPs) of one training run is captured by a famously simple approximation:

C \approx 6 \cdot N \cdot D

where N is the number of parameters and D is the number of training tokens. The 6 counts roughly 2 FLOPs per parameter per token in the forward pass and 4 in the backward pass (the backward pass computes gradients w.r.t. both activations and weights). This one relation is the budget line every planning decision moves along: for a fixed C, buying more parameters means buying fewer tokens, and vice versa.

from scaling import compute_flops, chinchilla_loss

# A 1B-parameter model trained on 20B tokens.
flops = compute_flops(n_params=1e9, n_tokens=20e9)
print(f"Compute: {flops:.2e} FLOPs")
print(f"Predicted loss: {chinchilla_loss(1e9, 20e9):.3f} nats/token")
Compute: 1.20e+20 FLOPs
Predicted loss: 2.580 nats/token

From Kaplan to Chinchilla

Two landmark papers shaped how the field spends compute:

  • Kaplan et al. (2020) measured that test loss falls as a power law in N, D, and C, and concluded that most extra compute should go into bigger models — data could lag behind.
  • Hoffmann et al. (2022), “Chinchilla” re-fit the curves more carefully and found Kaplan’s recipe left models badly under-trained. Compute-optimal training grows parameters and tokens together — a rule of thumb of about 20 tokens per parameter. Their 70B-parameter Chinchilla, trained on 1.4T tokens, beat the 280B-parameter Gopher trained on only 300B tokens — using the same compute.

Chinchilla summarized the loss surface with a parametric fit:

L(N, D) = E + \frac{A}{N^{\alpha}} + \frac{B}{D^{\beta}}

E is the irreducible loss (the entropy of natural text you can never beat); the two power-law terms shrink as you add parameters or data. scaling.py implements this fit and finds the loss-minimizing split of any compute budget by a ternary search along the C = 6ND constraint.

from scaling import compute_optimal_allocation, chinchilla_rule_allocation

C = compute_flops(7e10, 1.4e12)          # Chinchilla's own compute budget
rule = chinchilla_rule_allocation(C, tokens_per_param=20)
print(f"20x rule of thumb -> {rule['n_params']/1e9:.0f}B params, "
      f"{rule['n_tokens']/1e12:.1f}T tokens")   # recovers Chinchilla exactly

opt = compute_optimal_allocation(C)
print(f"Parametric optimum -> {opt['n_params']/1e9:.0f}B params, "
      f"{opt['n_tokens']/1e12:.1f}T tokens, "
      f"{opt['tokens_per_param']:.0f} tokens/param")
20x rule of thumb -> 70B params, 1.4T tokens
Parametric optimum -> 32B params, 3.0T tokens, 93 tokens/param

Notice the parametric optimum’s ratio is above 20 — the fit implies the compute-optimal tokens-per-parameter ratio drifts upward with scale. That is one reason modern models are trained far past 20× (Llama-3’s 8B model saw ~15T tokens — about 1875 tokens per parameter): once training compute is spent, a smaller-but-longer-trained model is cheaper to serve.

The Undertraining Story, Quantified

The clearest way to see scaling laws at work is to place real models against the compute-optimal frontier — the lowest loss achievable at each compute budget. A model with too many parameters for its token budget sits above the frontier: it wasted FLOPs it could have spent on data.

from scaling import demonstrate_scaling

_ = demonstrate_scaling()
========================================================================
REAL MODELS vs THE COMPUTE-OPTIMAL FRONTIER (Chinchilla fit)
========================================================================
model                tok/param     loss  frontier   excess
GPT-3 (2020)               1.7    2.002     1.954    0.048
Gopher (2021)              1.1    1.993     1.936    0.058
Chinchilla (2022)         20.0    1.937     1.930    0.007
Llama-3 8B (2024)       1875.0    1.949     1.923    0.026

Kaplan (2020): spend compute on bigger models.
Chinchilla (2022): grow params and tokens together (~20 tok/param as
a rule of thumb) — the parameter-heavy models above were undertrained.

The excess column is how much loss each model gives up versus a compute-optimal model using identical FLOPs. Gopher and GPT-3 — parameter-heavy, ~1 token per parameter — give up the most; Chinchilla sits almost on the frontier.

How the Frontier Is Found: IsoFLOP Profiles

Everything above used the coefficients (E, A, B, \alpha, \beta) — but where do they come from? CHINCHILLA_COEFFS is a fitted result, not a law of nature. This book earns its keep by building things, so let’s build the measurement behind the most-quoted number in the field. Chinchilla found the frontier three ways; the cleanest to reproduce is Approach 2 — IsoFLOP profiles:

  1. Pick a compute budget C.
  2. Train models of many sizes N at that budget — the tokens D = C / 6N are forced by the budget line, so bigger models see fewer tokens.
  3. Plot final loss against N. The curve is a valley: too-small models waste compute on surplus data, too-large models starve for it. The bottom is the compute-optimal model for that budget.
  4. Read N_{\text{opt}}(C) off each valley, then fit how it grows with C.

The valley has a closed form we can check against. Substitute D = C/6N into the loss and differentiate along the budget line:

L(N) = E + \frac{A}{N^{\alpha}} + B\left(\frac{6N}{C}\right)^{\beta}, \qquad \frac{dL}{dN} = 0 \;\Longrightarrow\; N_{\text{opt}} \propto C^{\,\beta/(\alpha+\beta)}.

So the compute-optimal exponent is a slope you can measure: a = \beta/(\alpha+\beta). With Chinchilla’s \alpha=0.34,\ \beta=0.28 that is a \approx 0.45, and because N \cdot D is pinned by C, tokens scale as D_{\text{opt}} \propto C^{\,1-a} \approx C^{0.55}. Two exponents near \tfrac12, summing to one — that is “grow parameters and tokens together.”

We recover both from data without touching the closed form. isoflop_optimum reads a valley off a simulated profile by grid argmin, and fit_power_law fits a straight line in log–log space (a power law is a line once you take logs):

from scaling import fit_power_law, isoflop_optimum, compute_flops

# A power law is a line in log-log space; recover y = 3·x^2 from clean points.
fit = fit_power_law([1, 2, 4, 8], [3, 12, 48, 192])
print(f"fit_power_law -> y = {fit['coefficient']:.1f}·x^{fit['exponent']:.2f}")

# The measured valley matches the analytic optimum — no closed form used.
print(f"IsoFLOP valley at 1e21 FLOPs: "
      f"{isoflop_optimum(1e21)['n_params']/1e6:.0f}M params")
fit_power_law -> y = 3.0·x^2.00
IsoFLOP valley at 1e21 FLOPs: 1840M params

Now run the full Approach-2 pipeline — profile nine budgets, read each valley, fit the two exponents — and compare to Chinchilla’s own reported numbers:

from scaling import demonstrate_isoflop

_ = demonstrate_isoflop()
========================================================================
FITTING THE SCALING LAW FROM SCRATCH — IsoFLOP profiles (Approach 2)
========================================================================
Measured from 9 IsoFLOP profiles:
  N_opt ∝ C^a   with a = 0.452
  D_opt ∝ C^b   with b = 0.548
  a + b = 1.000   (the C = 6ND budget line: should be 1)

Chinchilla's own reported exponents (Hoffmann et al., 2022):
  Approach 1 (fixed model, vary tokens)    a=0.50  b=0.50
  Approach 2 (IsoFLOP profiles)            a=0.49  b=0.51
  Approach 3 (parametric loss fit)         a=0.46  b=0.54

Our profiles were generated by the Approach-3 loss law, so we
recover Approach-3's exponents (a≈0.46, b≈0.54) — measured, not quoted.

Fitting profiles generated by the Approach-3 loss law recovers that approach’s exponents — a \approx 0.45, b \approx 0.55 — measured, not quoted, with a + b = 1 falling out to three decimals because the C = 6ND budget line forces D \propto C/N.

Interactive: The IsoFLOP Valley

Drag the compute budget. The left panel recomputes the loss-vs-N profile live and drops a marker at its valley; the right panel places that valley on the fitted scaling line. Watch the bottom slide rightward along a straight line whose slope is the exponent a \approx 0.45 — the scaling law, drawn as you move.

NoteKey Insight

The compute-optimal exponent is a slope you can measure. “20 tokens per parameter” is not an axiom handed down from a paper — it is a consequence of a \approx b \approx \tfrac12, which itself is read off the valleys of IsoFLOP profiles. Fitting the minima of curves generated by the loss law recovers its own exponents (a\approx0.45, b\approx0.55), and a + b = 1 is just the C = 6ND budget line reasserting itself. Discover the law and the rule of thumb falls out for free.

WarningSimulated Runs Are a Teaching Stand-In

Real IsoFLOP profiles come from actually training hundreds of models, each with its own noisy final loss; the valley is estimated through that scatter, and the fit is only as trustworthy as the spread of sizes around each minimum. Here the fitted loss law plays the role of those runs so the whole method runs in milliseconds — the procedure (sweep, find the valley, fit the slope) is identical, but a from-scratch sweep of tiny real models would wobble where this one is exact.

Interactive: Compute-Optimal Explorer

Slide the compute budget and read off the compute-optimal model (via the 20× rule). The chart plots the loss frontier against compute on a log-compute axis, with four real models placed where their own compute and loss land — the ones floating above the line are the under-trained ones.

NoteKey Insight

Under-training is spending compute on parameters instead of data. A model that is too large for its token budget sits above the compute-optimal frontier — it would have reached a lower loss, at the same FLOPs, as a smaller model trained on more tokens. Chinchilla’s headline result was that most large models of its era (Gopher, GPT-3, MT-NLG) were on the wrong side of that line.

Warning“20 Tokens per Parameter” Is a Rule of Thumb, Not a Law

The 20× ratio comes from Chinchilla’s compute-optimal fit at their scale. The parametric loss law implies the optimal ratio grows with compute, and — just as important — compute-optimal training only minimizes training cost. When a model will be served to many users, it is often worth over-training a smaller model far past 20× (Llama-3 8B: ~1875 tokens/param) to cut inference cost. Use scaling laws to plan, but pick the ratio for your deployment, not a magic number.

Fitting Through Noise: Fit the Curve, Not the Lowest Dot

The valley-reader above (isoflop_optimum) took the single lowest point of a profile — a plain argmin. That is honest here only because our losses are exact: they come from a formula, so the lowest dot really is the minimum. Real IsoFLOP profiles are measured — each dot is one model you actually trained, with its own noisy final loss (the seed, the data order, exactly where you stopped). The “simulated runs are a teaching stand-in” warning above flagged this; now we pay it off.

Here is the trap. Near the bottom the valley is nearly flat — that is what “optimal” means, the loss barely changes as you nudge the model size. So a whisker of measurement noise is enough to hand “lowest” to a neighbour. Run the same experiment again with a new seed and the argmin jumps to a different model size. The estimate is chasing the noise, not the valley.

This is exactly why Hoffmann et al. did not take the argmin. In their own words: “We fit a parabola to each IsoFLOPs curve to directly estimate at what model size the minimum loss is achieved.” A parabola uses every point in a window around the bottom and reads off the vertex — it pools the measurements and leans on the curvature, which noise cannot easily fake.

The Math: Why a Flat Valley Punishes the Argmin

Near its minimum any smooth loss is locally a parabola in x = \log_{10} N:

L(x) \;\approx\; a\,(x - x^\star)^2 + c ,

where x^\star is the true compute-optimal size and a > 0 is the valley’s curvature — small a means a flat, forgiving valley. Add independent measurement noise of scale \sigma to each observed loss. The argmin is fooled whenever noise overcomes the real loss gap between two sizes; on a flat valley that gap grows only like a\,\Delta x^2, so the argmin’s spread blows up roughly like \sqrt{\sigma / a} — the flatter the valley, the worse it gets.

The parabola fit instead solves an ordinary least squares for y = a x^2 + b x + c over a window of points and returns the vertex x^\star = -b / 2a. Because it averages many points and estimates the curvature a directly, its spread shrinks like \sigma / \sqrt{K} in the number of points K — it barely moves where the argmin lurches. Same runs, same noise; the fit just uses more of the information in them.

parabola vertex: x* = 1.000, y* = 0.000  (want 1.000, 0.000)
true optimum :   1824.2M params
argmin  guess:   1534.9M params  ( 15.9% off)
parabola fit :   1948.0M params  (  6.8% off)

Now the honest headline — not one lucky seed but the spread over many. The demonstrate_noisy_isoflop driver draws 200 independent noisy profiles at the same budget, estimates the valley both ways on each, and reports how far each estimator wanders from the analytic truth:

========================================================================
FITTING THROUGH NOISE — argmin vs. parabola on the IsoFLOP valley
========================================================================
Budget C = 1e21 FLOPs · per-run loss noise σ = 0.02 nats · 200 trials
True compute-optimal log₁₀N = 9.261
estimator      mean log₁₀N       std      rmse
argmin               9.269     0.178     0.178
parabola             9.261     0.062     0.062

The bare argmin wobbles 2.9× wider than the parabola fit.
Same runs, same noise — fitting the curve pools the points and holds still.

At a realistic \sigma = 0.02 nats of per-run noise the bare argmin scatters almost three times wider than the parabola fit — and, crucially, both are roughly unbiased, so the argmin isn’t wrong on average, it is just noisy. That extra variance is compute you paid for and threw away by reading one dot instead of the curve.

Interactive: One Noisy Profile

Turn up the measurement noise and reseed the runs. The grey curve is the true (noise-free) valley; the dots are what you actually measured. Watch the argmin marker (which dot is lowest) teleport between sizes as you reseed, while the fitted parabola’s vertex stays pinned near the true optimum. At \sigma = 0 they agree; crank \sigma and the argmin comes apart first.

TipTry This

Drag \sigma to 0: the argmin and the parabola vertex land on the same size — with no noise, the lowest dot is the minimum. Now push \sigma up and hold down the reseed slider. The argmin marker jumps from size to size on almost every seed, while the parabola line barely twitches. Keep going: find the \sigma where even the fit starts to wobble — that is the point where your runs are too noisy for the number of models you trained, and the cure is more sizes around the valley, not a lower dot.

Interactive: The Estimator Race

One profile is an anecdote. Run the whole experiment 120 times at the current noise level and histogram where each estimator lands. The true optimum is the dashed line; a good estimator makes a tight pile on it. The argmin cloud is wide and grainy (it can only ever land on a grid size); the parabola cloud is tight and continuous.

NoteKey Insight

A flat minimum is where you must fit, not pick. The very property that makes a model compute-optimal — the loss is insensitive to size right there — is what makes the single lowest run a terrible pointer to it. The estimator’s whole job is to recover the valley’s location from its curvature, and a parabola fit does exactly that while an argmin throws the curvature away and reads one noisy dot. This is not a scaling-laws quirk; it is the general rule for reading a minimum off noisy measurements, from learning-rate sweeps to hyperparameter search.

WarningPitfall: Never Argmin a Noisy Sweep

Taking the single best run of a sweep — the lowest loss, the top eval score, the fastest step — is argmin, and on a flat landscape it reports the noise, not the signal. Fit a smooth curve through the points and read its optimum. The opposite mistake also bites: a too-wide parabola window drags in the profile’s tails, which are not quadratic, and biases the vertex. Fit a window snug around the minimum — wide enough to average the noise, narrow enough to stay in the bowl.

Going Deeper

  • Hoffmann, Borgeaud, Mensch, et al., Training Compute-Optimal Large Language Models (Chinchilla, 2022) — the three estimation approaches; §3.2 is the IsoFLOP method we reproduced, and Table 2 reports the exponents (N\propto C^a, D\propto C^b): Approach 1 (0.50/0.50), Approach 2 (0.49/0.51), Approach 3 (0.46/0.54, the parametric fit whose coefficients scaling.py carries).
  • Kaplan, McCandlish, Henighan, et al., Scaling Laws for Neural Language Models (2020) — the original power-law framing Chinchilla corrected.
  • Besiroglu et al., Chinchilla Scaling: A Replication Attempt (2024) — how sensitive the fitted exponents are to the fitting procedure, a good sanity check on the numbers above.

Zero-Shot HP Transfer: Maximal Update Parametrization (μP)

Scaling laws told us how big to make the model. They said nothing about its hyperparameters — the learning rate, the initialization scale, the output multiplier. And here is the trap: under the parametrization everyone uses by default (call it standard parametrization, SP), the optimal learning rate shifts every time you make the model wider. The value you carefully tuned on a 100M-parameter model is simply wrong for a 10B one — but you cannot afford to sweep learning rates on the 10B model to find out.

Maximal Update Parametrization (μP) removes the shift. Under μP the optimal hyperparameters are (nearly) invariant to width, so you can tune once on a tiny, cheap proxy and copy the numbers straight to the giant target with no further tuning. This is μTransfer (Yang et al., 2022). Cerebras-GPT tuned its hyperparameters on a ~40M-parameter proxy and reused the same learning rate all the way up to 2.7B — for about 7% of the target’s compute.

Intuition: keep every layer’s update Θ(1)

Why does the optimal learning rate drift under SP? Watch what one Adam step does to a hidden preactivation. A hidden unit is a sum over the width:

h_i = \sum_{j=1}^{n} W_{ij}\, x_j .

Adam normalizes each gradient to roughly unit scale, so after a step every weight entry moves by about \eta (the learning rate). The preactivation therefore moves by

\Delta h_i = \sum_{j=1}^{n} \Delta W_{ij}\, x_j \;\approx\; \eta \cdot \Theta(n),

because after the first gradient step the update aligns with the input and the n terms add coherently. Read that again: the size of the update grows with the width n. SP holds \eta fixed as you scale, so a wider model takes a bigger step in feature space — it overshoots, and you have to lower the learning rate to compensate. That is the drift.

μP’s fix is embarrassingly simple: scale the hidden learning rate down in proportion to width, \eta \propto 1/n, so that \Delta h_i = \Theta(1) at every width. Each layer then updates maximally — a full \Theta(1) change, never vanishing — yet stably — never exploding. That is the “maximal update” the name refers to, and it is exactly what pins the loss landscape (and its optimal learning rate) in place as the model grows.

NoteKey Insight

μP is not a new optimizer or architecture. It is a rescaling of three things — initialization, forward multipliers, and per-layer learning rates — as a function of width, chosen so that the training dynamics have a well-defined infinite-width limit. Because the dynamics stop depending on width, so do the best hyperparameters. Tune small, transfer big.

The rules (SP vs μP, for Adam)

μP is an abc-parametrization: for each weight tensor it specifies how the init variance (b), a forward multiplier (a), and the learning rate (c) scale with that tensor’s fan_in. Weights fall into three categories — input (embeddings; fan_in is fixed), hidden (both dimensions scale with width), and output (the readout/unembedding). Here is the table we implement (the “Table 8” form used by the mup package), expressed relative to a base width n_0 with width multiplier m = n/n_0:

Weight Init variance Forward multiplier Adam LR
Input (embeddings) 1/\text{fan\_in} 1 \eta (constant)
Hidden (n\times n) 1/\text{fan\_in} 1 \eta \cdot n_0/n
Output (readout) 1/\text{fan\_in} n_0/n \eta (constant)

Every weight is still initialized \mathcal{N}(0,\,1/\text{fan\_in}) — the usual Kaiming init. The only two things μP changes versus SP are the two shaded cells: the hidden learning rate (\eta \cdot n_0/n instead of \eta) and the output multiplier (n_0/n instead of 1). At the base width m=1, both revert to 1 — so μP is identical to SP at the base width, which is precisely the anchor you tune at.

WarningTransformers need one more change: 1/d attention

For an MLP the table above is complete. A Transformer maps the width down to a finite number of attention heads, and every such “infinite → finite” map picks up a 1/\text{fan\_in} multiplier. Concretely, μP scales attention logits by 1/d_k (the head dimension) instead of the familiar 1/\sqrt{d_k}: \text{logit} = q^\top k / d_k. The mup package uses 8/d_k in practice purely so it coincides with 1/\sqrt{d_k} at the common d_k=64.

Code: an MLP under both parametrizations

The implementation lives in mup.py. The scaling rules are three one-line pure functions:

from mup import (
    width_multiplier, mup_hidden_lr, mup_output_multiplier,
    MuMLP, make_task, coordinate_check, coord_blowup_ratio,
    lr_transfer_sweep, optimal_log_lr,
)

# At the base width everything is 1x; a 4x-wider model quarters the hidden LR.
print("width mult (1024 vs base 256):", width_multiplier(1024, 256))
print("hidden LR  (base 0.01):        ", mup_hidden_lr(0.01, 1024, 256))
print("output multiplier:             ", mup_output_multiplier(1024, 256))
width mult (1024 vs base 256): 4.0
hidden LR  (base 0.01):         0.0025
output multiplier:              0.25

MuMLP builds the network under either parametrization. Watch the anchor property — at the base width, μP and SP are the same model:

task = make_task()
sp = MuMLP(task.d_in, 256, task.d_out, parametrization="sp",  base_width=256, seed=7)
mp = MuMLP(task.d_in, 256, task.d_out, parametrization="mup", base_width=256, seed=7)

import torch
print("identical outputs at base width:", torch.allclose(sp(task.X), mp(task.X)))
print("μP output multiplier at 256:    ", mp.output_multiplier)      # 1.0

# Make it 4x wider and the two μP knobs engage:
mp_wide = MuMLP(task.d_in, 1024, task.d_out, parametrization="mup", base_width=256, seed=7)
print("μP output multiplier at 1024:   ", mp_wide.output_multiplier) # 0.25
print("μP hidden LR at 1024 (base 1e-2):", mp_wide.param_groups(1e-2)[1]["lr"])
identical outputs at base width: True
μP output multiplier at 256:     1.0
μP output multiplier at 1024:    0.25
μP hidden LR at 1024 (base 1e-2): 0.0025

The coordinate check

The coordinate check is the diagnostic for a correct μP implementation, and it makes the whole idea visible in one picture. The recipe: build the network at a range of widths, take a few Adam steps, and plot the average absolute activation coordinate of a hidden layer against width — one line per training step.

  • Correct μP: the lines are flat across width — activations are \Theta(1) no matter how wide the model is.
  • Broken SP: after the very first gradient step, the activations explode with width — exactly the overshoot our intuition predicted.
widths = [64, 128, 256, 512, 1024]
sp_cc = coordinate_check(widths, parametrization="sp")   # base_lr=0.03, 2 steps
mp_cc = coordinate_check(widths, parametrization="mup")

print("hidden coord size after 1 Adam step (t=1):")
print(f"{'width':>6} | {'SP':>9} | {'μP':>7}")
for w in widths:
    s = next(r["coord"] for r in sp_cc if r["width"] == w and r["step"] == 1)
    m = next(r["coord"] for r in mp_cc if r["width"] == w and r["step"] == 1)
    print(f"{w:>6} | {s:>9.3f} | {m:>7.3f}")

print(f"\nSP blow-up 64→1024: {coord_blowup_ratio(sp_cc, 1):.0f}x")
print(f"μP blow-up 64→1024: {coord_blowup_ratio(mp_cc, 1):.1f}x")
hidden coord size after 1 Adam step (t=1):
 width |        SP |      μP
    64 |     0.484 |   0.484
   128 |     1.435 |   0.551
   256 |     4.309 |   0.604
   512 |    20.968 |   0.614
  1024 |    69.405 |   0.596

SP blow-up 64→1024: 143x
μP blow-up 64→1024: 1.2x

Drive the slider below through the training steps. At t = 0 (before any update) both panels are flat — the difference is invisible at initialization. Push to t = 1 and the left (SP) panel fans out across two orders of magnitude while the right (μP) panel stays pinned flat. That single frame is μP.

TipTry This

Hold the slider at t = 1 and read the left panel: each doubling of width roughly triples the SP activation size — the coherent \Theta(n) growth from the intuition, made literal. Now look at t = 0: SP and μP are indistinguishable. The pathology is invisible at init and only appears once gradients flow. That is why you cannot catch it by inspecting a freshly initialized model — you have to take a step.

μTransfer in action

The payoff: if the dynamics are width-invariant, so is the optimal learning rate. Here we sweep the learning rate at four widths and plot the final loss. Under SP the U-curves slide leftward as the model grows — the optimum keeps moving. Under μP the minima line up in a vertical column: the best learning rate found on the smallest model is (nearly) the best on all of them.

grid = [round(-3.0 + 0.25 * i, 2) for i in range(13)]   # log10(lr): -3.0 … 0.0
for name in ("sp", "mup"):
    sweep = lr_transfer_sweep([64, 128, 256, 512], grid,
                              parametrization=name, base_width=64, n_steps=25)
    best = {w: optimal_log_lr(sweep, w) for w in [64, 128, 256, 512]}
    spread = max(best.values()) - min(best.values())
    row = "  ".join(f"n={w}: {best[w]:+.2f}" for w in best)
    print(f"{name:>3}  best log₁₀(lr) →  {row}   (spread {spread:.2f} dex)")
 sp  best log₁₀(lr) →  n=64: -2.00  n=128: -2.25  n=256: -2.50  n=512: -3.00   (spread 1.00 dex)
mup  best log₁₀(lr) →  n=64: -2.00  n=128: -1.75  n=256: -1.50  n=512: -1.50   (spread 0.50 dex)
TipTry This

Flip the toggle between μP and Standard. Under Standard, the coloured minima march to the left as width grows — a full order of magnitude of drift from the smallest to the largest model. Under μP they collapse into a narrow band: the learning rate transfers. (A small residual wobble remains because μTransfer is exact only in the infinite-width limit — at widths of 64–512 there are honest finite-size corrections. μP shrinks the drift from ~10× to under ~3×.)

Common Pitfalls with μP

  • Forgetting the base shape. μP scalings are relative to a base width. You must record the proxy’s per-tensor fan_in (“base shapes”) and scale everything against it — a μP model with no base shape is just SP.
  • Zero-init the readout (and query). The paper recommends initializing the output layer (and the attention query projection) to zero, so their activations are exactly 0 at t=0 regardless of width. This removes a proxy-vs-target mismatch at initialization; they become nonzero after the first gradient step.
  • Weight decay and dropout do not transfer. μTransfer covers the optimization hyperparameters (LR, init, multipliers, Adam betas). Regularization strength is width-dependent and must be set at the target scale.
  • 1/d, not 1/√d, attention. The single easiest thing to get wrong when moving from an MLP to a Transformer — see the warning above.

What μP buys you

NoteKey Insight

The economics are the whole point. Sweeping hyperparameters on a frontier model is often as expensive as training it. μTransfer moves that sweep onto a proxy that is orders of magnitude cheaper, and the result carries over. Combined with scaling laws — which pick the compute-optimal size — μP is what makes training a model you have never trained before a predictable engineering task rather than a gamble.

Going Deeper

Multi-Token Prediction: Training n Futures at Once

Every objective so far has the model predict one token — the next one — from the trunk representation z_t at position t. But z_t carries far more than “what comes next”: to predict the next token well it must already anticipate the words after it. Multi-token prediction (MTP) makes that latent planning an explicit target. Keep the same shared trunk, but attach n output heads: head k predicts the token k positions ahead. Head 1 is the ordinary next-token head; heads 2..n look further out.

Two things fall out for free:

  • A denser training signal. Forcing z_t to also predict x_{t+2},\dots,x_{t+n} pushes the trunk to plan ahead. The gains grow with scale — worse than the baseline for small models, but a 13B model trained this way solves 12% more HumanEval and 17% more MBPP problems than a next-token baseline, at no extra training time.
  • A built-in speculative drafter. The extra heads are a draft model: one forward pass proposes n tokens, so you decode in bigger jumps — up to 3.0× faster on code, 2.7× on natural language, 6.4× for a byte-level model. This is the same draft/verify machinery you built in Module 16, except the drafter costs nothing extra to train.

The from-scratch implementation lives in mtp.py.

Intuition: One Trunk, Many Futures

Picture the sequence laid out left to right. A standard model, standing at position t, draws a single arrow to t+1. MTP fans n arrows out of the same spot — one per head — reaching t+1, t+2, \dots, t+n. The heads share everything below them; only the final linear read-out differs. Step through the fan-out:

The Math: Shifted Targets and a Summed Loss

The paper writes the objective as one factorized cross-entropy over the next n tokens (their Eq. 2), factorized across the heads on the shared trunk output z_t (Eq. 3):

\mathcal{L}_n \;=\; -\sum_t \log P_\theta\!\left(x_{t+n:t+1}\mid x_{t:1}\right) \;=\; -\sum_t \sum_{k=1}^{n} \log P_\theta\!\left(x_{t+k}\mid z_t\right).

So expanded it is a plain sum over positions and heads of per-token cross-entropy. The only new mechanic is the target: head k is scored against x_{t+k}, not x_{t+1}. Building those targets is a shift — and the last k positions of a length-T sequence have no k-ahead token, so they are masked out (ignore_index = -100, skipped by cross-entropy). The grid below shows, for a toy sequence, which token each head is trained to predict at each position:

NoteKey Insight

Read the first column: at position t=0 the four heads are trained toward x_1, x_2, x_3, x_4 — the next four tokens — all from the single representation z_0. Standard training only ever asks for x_1. Set n=1 and the grid collapses to that one row: multi-token prediction with one head is exactly next-token prediction.

Code: The Heads and the Loss

MultiTokenHeads is just n bias-free linear read-outs sharing the trunk; mtp_loss sums their per-head cross-entropies. The load-bearing check is the n=1 identity — the new objective must reduce to the old one:

import torch
from mtp import MultiTokenHeads, mtp_loss, next_token_loss

torch.manual_seed(0)
ids = torch.randint(0, 50, (4, 16))          # a batch of token sequences
trunk = torch.randn(4, 16, 32)               # a stand-in trunk output z

# n = 1: MTP must equal ordinary next-token cross-entropy, bit for bit.
one_head = MultiTokenHeads(d_model=32, vocab_size=50, n_future=1)
logits1 = one_head(trunk)                    # (B, T, 1, V)
total1, _ = mtp_loss(logits1, ids)
baseline = next_token_loss(logits1[:, :, 0, :], ids)
print(f"n=1 MTP loss == next-token loss: {torch.allclose(total1, baseline)}")

# n = 4: four heads, one summed loss. Per-head losses decompose the total.
four_heads = MultiTokenHeads(d_model=32, vocab_size=50, n_future=4)
logits4 = four_heads(trunk)                   # (B, T, 4, V)
total4, per_head = mtp_loss(logits4, ids)
print(f"per-head losses: {[round(x, 3) for x in per_head.tolist()]}")
print(f"total == sum(per-head): {torch.allclose(total4, per_head.sum())}")
n=1 MTP loss == next-token loss: True
per-head losses: [4.027, 4.108, 4.057, 3.95]
total == sum(per-head): True

To match DeepSeek-V3’s recipe — which keeps the next-token head at full weight but down-weights the extra heads (its \lambda) — pass head_weights:

# DeepSeek-V3 uses lambda = 0.3 (then 0.1) on the extra-token loss.
weighted, _ = mtp_loss(logits4, ids, head_weights=[1.0, 0.3, 0.3, 0.3])
print(f"weighted total: {weighted.item():.3f}")
weighted total: 7.661

The Memory Trick: n Heads, One Head’s Memory

There is an obvious objection: an output head produces a logit for every token in the vocabulary, and V is enormous (often 100k+). Materializing all n heads’ logits at once costs O(nV) floats per position — the vocabulary axis, n times over. That would make MTP a memory hog.

The paper’s fix (their Fig. 2) is to compute the forward and backward pass of each head sequentially, accumulate its gradient into the d-dimensional trunk gradient, and free the head’s logits before moving to the next head. Peak memory drops from O(nV + d) to O(V + d) — independent of n. Because the loss \sum_k w_k \mathcal{L}_k is separable across heads, this streaming computation yields the identical loss and the identical trunk gradient:

from mtp import mtp_loss_streaming

heads = MultiTokenHeads(d_model=32, vocab_size=50, n_future=4)

# Batched path: all four heads' logits live at once.
trunk_a = torch.randn(4, 16, 32, requires_grad=True)
loss_batched, _ = mtp_loss(heads(trunk_a), ids)
loss_batched.backward()

# Streaming path: one head's logits in memory at a time (the paper's Fig. 2).
trunk_b = trunk_a.detach().clone().requires_grad_(True)
loss_stream, _ = mtp_loss_streaming(trunk_b, heads, ids)
loss_stream.backward()

print(f"same loss:          {torch.allclose(loss_batched, loss_stream)}")
print(f"same trunk gradient: {torch.allclose(trunk_a.grad, trunk_b.grad, atol=1e-6)}")
same loss:          True
same trunk gradient: True

That gradient equality is the memory argument: you never pay for n vocab-sized logit tensors, yet you train exactly as if you had.

Free Speedup: Self-Speculative Decoding

At inference the heads become a drafter. From the last position, head k proposes the token k steps ahead — so one forward pass yields an n-token draft. Feed the draft back through the model and verify it with head 1 (the trusted next-token head): accept the longest prefix head 1 agrees with, plus its correction at the first disagreement. Because that correction is conditioned only on the accepted prefix, the output is bit-for-bit the greedy head-1 string — exactly the draft/verify guarantee from Module 16, now powered by MTP’s own heads:

from mtp import MTPModel, plain_generate, self_speculative_generate

torch.manual_seed(0)
model = MTPModel(vocab_size=32, d_model=64, n_future=4, max_len=32).eval()
prompt = torch.randint(0, 32, (1, 5))

plain = plain_generate(model, prompt, max_new_tokens=24)          # 1 token / forward
spec, accepts = self_speculative_generate(model, prompt, max_new_tokens=24)

same = torch.equal(plain[:, :spec.shape[1]], spec[:, :plain.shape[1]])
print(f"self-speculative == greedy: {same}")
print(f"tokens committed per round: {accepts}  (each round is one draft+verify)")
self-speculative == greedy: True
tokens committed per round: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]  (each round is one draft+verify)

The output is identical to plain greedy decoding; the win is that several rounds commit more than one token, so you reach the same string in fewer forward passes. On real text the far heads are progressively harder, so acceptance is partial (the paper reports ~2.5 of 3 drafted code tokens accepted); on a perfectly predictable stream, whole drafts commit.

Interactive Exploration

Train the tiny MTP model on a fully predictable stream for several values of n, then read off two things: how accurately each head predicts its k-ahead token, and how many tokens self-speculative decoding commits per forward pass (its speedup over one-token-at-a-time decoding). Increase n and watch the drafter reach further:

TipTry This
  1. Set n=1: one head, exactly one token per forward pass — the ordinary decoder, no speedup. This is the baseline every larger n improves on.
  2. Climb to n=8: on this fully-learnable stream every head hits ~100% and the drafter commits all eight tokens per round. Real language won’t be this kind, but the shape — a free drafter that costs nothing to train — is real.
  3. Weight the heads: in mtp_loss, try head_weights=[1.0, 0.3, 0.3, ...] to mirror DeepSeek-V3’s \lambda, protecting next-token quality while still training the drafter.
WarningMTP Helps at Scale, Not on Toys

The headline gains are a large-model phenomenon. Gloeckle et al. are explicit that multi-token prediction is worse than a next-token baseline for small models and only overtakes it as models grow (their wins are clearest at 3B–13B, and on code). The tiny model here exists to make the mechanics — shifted targets, the summed loss, exact self-speculative decoding — visible and provable; do not read its toy-stream accuracy as evidence MTP will lift a small model on real data.

Training Tips

Quick Reference Table

Symptom Likely Cause Solution
Loss = NaN LR too high Reduce LR by 10x
Loss stuck LR too low Increase LR by 2-5x
Loss oscillates Batch too small Use gradient accumulation
Overfitting Not enough data More data, more dropout
Underfitting Model too small More layers/heads/dims
Slow training No GPU/MPS Use hardware acceleration
OOM errors Batch too large Reduce batch size, use accumulation
Training crash No checkpoints Save every N steps

Hyperparameter Recommendations

Based on published research and common practices:

Hyperparameter Small Models (<1B) Large Models (>1B)
Learning rate 1e-4 to 6e-4 1e-4 to 3e-4
Warmup 1-2% of steps 0.1-1% of steps
Weight decay 0.01 - 0.1 0.01 - 0.1
Beta1 0.9 0.9
Beta2 0.999 0.95
Batch size 256 - 1024 tokens 1M - 4M tokens
Gradient clip 1.0 1.0

Memory Optimization Strategies

  1. Gradient accumulation: Simulate larger batches
  2. Mixed precision (fp16/bf16): ~50% memory reduction
  3. Gradient checkpointing: Trade compute for memory
  4. FSDP/DeepSpeed: Shard model across GPUs

Interactive Exploration

Experiment with learning rate schedules in real-time. Adjust the hyperparameters to see how warmup and cosine decay shape the learning rate curve.

TipTry This

  1. Effect of warmup: Set warmup to 0, then gradually increase to 200. Notice how the curve changes from immediate peak to gradual ramp-up.

  2. Long vs short training: Compare total_steps=500 vs total_steps=2000 with the same warmup. See how the decay rate changes.

  3. Min LR matters: Set min_lr to 0, then to 1e-5. The floor prevents the model from completely stopping learning.

  4. Warmup ratio: Try warmup_steps = 1-2% of total_steps (common in practice). For 1000 steps, that’s 10-20 warmup steps.

  5. Drag the current step slider to see the exact LR at any point in training.

Exercises

Exercise 1: Learning Rate Finder

Implement a learning rate finder that trains for a few iterations at exponentially increasing learning rates and plots loss vs learning rate.

# Your implementation here
def lr_finder(model, tokens, start_lr=1e-7, end_lr=1e-1, num_steps=100):
    """Find optimal learning rate by training with exponentially increasing LR."""
    # TODO: Implement this
    pass

Exercise 2: Custom Scheduler

Implement a linear warmup + linear decay scheduler (instead of cosine decay).

# Your implementation here
class LinearScheduler:
    def __init__(self, optimizer, warmup_steps, total_steps, min_lr=0.0):
        # TODO: Implement this
        pass

    def step(self):
        pass

Exercise 3: Training with Validation

Modify the training loop to: 1. Compute validation loss every N steps 2. Save the best model (lowest validation loss) 3. Implement early stopping if validation loss doesn’t improve for M steps

Exercise 4: Prove the Decoupling to Yourself

Using the from-scratch Adam and AdamW from optimizers.py, show that coupled and decoupled weight decay are not the same thing. Train two identically-initialized layers with the same lr and weight_decay, then confirm (a) they diverge with decay on, and (b) they are bit-identical with weight_decay=0.

# Your implementation here
from optimizers import Adam, AdamW

def decoupling_gap(weight_decay, steps=25, seed=0):
    """Return the max weight difference between Adam(L2) and AdamW after training."""
    # TODO: build twin nn.Linear layers, run Adam vs AdamW, return the gap
    pass

# Expect: decoupling_gap(0.3) > 0  and  decoupling_gap(0.0) == 0

Exercise 5: Why fp8 Needs Per-Tensor Scaling

Take a small weight vector and quantize it to e4m3 without scaling, then with per-tensor scaling, and measure how much the scaling recovers. Confirm that the scale is exactly max_normal / amax, and that shrinking every weight by 10× (so more of them underflow) makes the naive path worse while the scaled path is unchanged — because scaling is scale-invariant by construction.

# Your implementation here
from precision import FP8_FORMATS, quantize_per_tensor, per_tensor_scale

def scaling_benefit(weights, shrink=1.0):
    """Return (naive_rel_error, scaled_rel_error) for e4m3 after shrinking weights."""
    w = [x * shrink for x in weights]
    # TODO: quantize_per_tensor(w, e4m3); return the two mean relative errors
    pass

# Expect: naive error grows as shrink -> 0.1, scaled error stays ~constant.

Exercise 6: Orthogonalization Whitens the Update

Using newton_schulz and orthogonalize from muon.py, confirm that Muon’s update is (approximately) the polar factor of the momentum matrix — and that this whitens the spectrum. Build a deliberately ill-conditioned matrix, orthogonalize it both ways, and compare.

# Your implementation here
from muon import newton_schulz, orthogonalize, singular_values

def whitening_check(condition=50.0, seed=0):
    """Return (raw_condition_number, orthogonalized_condition_number)."""
    import torch
    torch.manual_seed(seed)
    u, _ = torch.linalg.qr(torch.randn(32, 32))
    v, _ = torch.linalg.qr(torch.randn(32, 32))
    m = (u * torch.logspace(0, torch.log10(torch.tensor(condition)), 32)) @ v.T
    # TODO: singular_values(m) vs singular_values(newton_schulz(m));
    #       return max/min for each (the raw one is ~condition, the ortho one ~1)
    pass

# Expect: raw condition number ~= `condition`, orthogonalized ~= 1.
# Bonus: check newton_schulz(m) and orthogonalize(m) point the same way
# (cosine > 0.97), but only orthogonalize() sets every singular value to exactly 1.

Exercise 7: Break μP and Watch the Coordinate Check Fail

μP is a package deal — its two knobs work together. Break just one and the coordinate check should light up. Using MuMLP and coordinate_check from mup.py, build a “half-μP” model that scales the hidden learning rate but forgets the output multiplier, and confirm its coordinate check is no longer flat.

# Your implementation here
from mup import MuMLP, coordinate_check, coord_blowup_ratio

def half_mup_check():
    """Compare full-μP vs SP blow-up ratios at t=1 over an 8x width sweep."""
    widths = [64, 128, 256, 512]
    mp = coordinate_check(widths, parametrization="mup")
    sp = coordinate_check(widths, parametrization="sp")
    # TODO: return (coord_blowup_ratio(mp, 1), coord_blowup_ratio(sp, 1))
    #       Expect μP ~1.x (flat), SP >10x (exploding).
    pass

# Bonus: monkey-patch a MuMLP so output_multiplier stays 1.0 while the hidden LR
# still scales, then run the check — the readout logits should now drift with
# width even though the hidden activations stay put.

Exercise 8: Find the Cap That Just Barely Never Fires

QK-Clip is a self-deactivating controller: for a given optimizer “heat”, there is a cap \tau high enough that the logits never cross it, so the clip is a perpetual no-op. Using qk_clip_logit_growth from qk_clip.py, find the smallest integer \tau (for a fixed inflation) at which no step fires the clip.

# Your implementation here
from qk_clip import qk_clip_logit_growth

def smallest_quiet_tau(inflation=1.15, steps=20):
    """Return the smallest integer tau at which no clip event fires."""
    # A run never fires iff the *uncapped* curve never exceeds tau. Trace once
    # with a huge tau to read the uncapped max, then round up.
    # TODO: rows = qk_clip_logit_growth(steps=steps, tau=1e9, inflation=inflation)
    #       return math.ceil(max(r["uncapped"] for r in rows))
    pass

# Bonus: lower `inflation` toward 1.0 and watch the quiet tau shrink — a gentler
# optimizer needs a looser rail. This is why healthy runs settle below tau=100.

Exercise 9: Fit the Data Exponent Yourself

You saw fit_isoflop_scaling recover a\approx0.45 for parameters. The tokens exponent b should come out near 0.55 so that a+b=1. Without calling fit_isoflop_scaling, reproduce b from scratch: read the valley of each IsoFLOP profile, collect D_{\text{opt}}, and fit D_{\text{opt}}\propto C^b.

# Your implementation here
from scaling import isoflop_optimum, fit_power_law

def fit_data_exponent(flops_grid=None):
    """Return the fitted exponent b in D_opt ∝ C^b, measured off IsoFLOP valleys."""
    # TODO: budgets = flops_grid or [10**e for e in range(19, 28)]
    #       d_opt = [isoflop_optimum(c)["n_tokens"] for c in budgets]
    #       return fit_power_law(budgets, d_opt)["exponent"]
    pass

# Bonus: change CHINCHILLA_COEFFS' alpha/beta and confirm a = beta/(alpha+beta)
# still predicts what the fit recovers — the exponent is not magic, it's the ratio.

Exercise 10: The Argmin’s Variance, and the Window’s Bias

The noise section claimed the parabola fit beats the bare argmin and that too wide a fit window re-biases the vertex. Measure both yourself with noise_robustness and estimate_valley. First confirm the headline: at \sigma=0.02 the argmin’s spread is ~3× the parabola’s. Then sweep the fit window from very narrow (say 0.2 decades) to very wide (the whole profile) and watch the parabola estimate’s bias (mean − truth) grow as the window swallows the non-quadratic tails — the classic bias–variance trade, live.

# Your implementation here
from scaling import noise_robustness, isoflop_profile, estimate_valley, compute_optimal_allocation
import math

# Part 1: reproduce the ~3× spread ratio.
r = noise_robustness(flops=1e21, noise=0.02, trials=200, seed=0)
print(f"argmin std {r['argmin']['std']:.3f}  vs  parabola std {r['parabola']['std']:.3f}  "
      f"({r['std_ratio']:.1f}×)")

# Part 2: TODO — for window in (0.2, 0.4, 0.6, 1.0, 1.5):
#   over many seeds, estimate_valley(prof, method="parabola", window=window),
#   and print mean(estimate) − truth (the bias) alongside its std (the variance).
#   You should see the bias stay ~0 for snug windows, then grow as it widens.
truth = math.log10(compute_optimal_allocation(1e21)["n_params"])
argmin std 0.178  vs  parabola std 0.062  (2.9×)

Summary

This module covered:

  1. Cross-entropy loss measures prediction quality (lower = better), with mathematical foundations in information theory
  2. Perplexity provides an intuitive metric: exp(loss) - “choosing among N equally likely options”
  3. Learning rate scheduling with warmup + cosine decay prevents early instability and enables fine-tuning
  4. Optimizers from scratch — built SGD, Adam, and AdamW in optimizers.py, each matching torch.optim bit-for-bit. The Adam-vs-AdamW difference is one term’s placement: coupled L2 rides the gradient through the 1/√v̂ denominator (per-parameter, unintended), while AdamW’s decoupled decay shrinks every weight uniformly — identical when weight_decay=0, divergent otherwise
  5. Gradient accumulation increases effective batch size without adding memory
  6. Gradient clipping (max_norm=1.0) prevents exploding gradients, essential for transformers
  7. Batch size tradeoffs affect memory, training dynamics, and generalization
  8. Mixed precision & numerics — built the float number line from scratch: fp16 spends bits on precision (overflows at 65504, underflows near 6e-5), bf16 keeps fp32’s range at coarser precision, and loss scaling lifts tiny gradients out of fp16’s underflow hole (2× speed, 50% memory). At the frontier, fp8 uses two formats — e4m3 (precise, no infinity, max 448) for the forward pass and e5m2 (fp16’s range) for gradients — with per-tensor scaling (map amax→max_normal) rescuing tensors that would otherwise underflow in 8 bits
  9. Distributed training (DDP, FSDP) scales training to multiple GPUs
  10. Common failure modes (NaN loss, stuck training, oscillation) and their solutions
  11. Checkpointing strategies ensure you never lose training progress
  12. Scaling laws predict loss from compute: the C≈6ND rule plus Chinchilla’s ~20 tokens/param let you size a model and its data budget instead of guessing — and reveal that parameter-heavy models like GPT-3 and Gopher were under-trained. And the law isn’t handed down — you fit it: IsoFLOP profiles (built in scaling.py as isoflop_profile/isoflop_optimum/fit_power_law/fit_isoflop_scaling) sweep model size at each compute budget, read the loss valley, and recover the compute-optimal exponents (N\propto C^{0.45}, D\propto C^{0.55}, summing to one because C=6ND) by log-log least squares — so “20 tokens/param” is a measured consequence of a\approx b\approx\tfrac12, not an axiom. And because real runs are noisy, the valley is fit, not argmin-ed: near a flat minimum the single lowest run is a high-variance pointer, so — as Chinchilla actually did — you fit a parabola to each profile and read its vertex (fit_parabola/estimate_valley/noise_robustness), which scatters ~3× less than a bare argmin at the same measurement noise
  13. Muon treats each weight as a matrix, not a bag of scalars: it orthogonalizes the momentum matrix (all singular values → 1) with a from-scratch Newton–Schulz iteration — no SVD, no second-moment state — scales the result by 0.2·√max(A,B) to match AdamW’s update RMS, and adds decoupled weight decay. Built in muon.py, it clears plain momentum on a toy and, at scale (Moonlight), reaches AdamW quality with ~2× less compute — the first real challenger to a decade of Adam
  14. Maximal Update Parametrization (μP) makes hyperparameters transfer across width: because SP takes a \Theta(\text{width})-sized feature step per Adam update, its optimal learning rate drifts as you scale — μP rescales the hidden LR (\propto 1/\text{width}) and the readout multiplier (\propto 1/\text{width}) so every layer’s update stays \Theta(1). Built in mup.py, the coordinate check shows SP activations exploding ~140× across an 8× width sweep while μP stays flat, and the μTransfer sweep shows the optimal LR pinned in place. Tune on a tiny proxy, copy the numbers to the giant model (Yang et al., 2022; Cerebras-GPT tuned a 40M proxy for a 2.7B target)
  15. QK-Clip & MuonClip are what let Muon scale to the frontier. Muon’s larger updates inflate the attention logits faster than AdamW — the same explosion m05’s QK-Norm cures in activation space — so QK-Clip cures it in weight space: after each step, any head whose max logit crosses a cap \tau has its Q/K projections rescaled by \sqrt{\min(1, \tau/S_{\max}^h)}, and because the logit is bilinear it lands back at exactly \tau. Built in qk_clip.py with the MuonClip wrapper (step, then clip); it is a self-deactivating controller — once the logits settle it stops firing — and it is the recipe behind Kimi K2’s zero-loss-spike 1T-parameter run (\tau = 100)
  16. Multi-token prediction changes the objective: keep the shared trunk but add n output heads, head k predicting the token k steps ahead, and sum their cross-entropies. Built in mtp.py, it reduces to plain next-token training at n=1 (proven bit-for-bit); the paper’s memory trick (compute each head sequentially, free its logits) keeps peak memory at O(V+d) regardless of n — the same loss and trunk gradient. It buys a denser signal (a 13B model solves 12%/17% more HumanEval/MBPP) and a built-in speculative drafter whose greedy self-speculative decoding is exact (up to 3× faster on code). Helps at scale, hurts small models

Key Takeaways

  • Always use warmup (at least 1% of steps) to stabilize early training
  • Monitor gradient norms alongside loss - they tell you about training stability
  • Start with standard hyperparameters (lr=3e-4, wd=0.01, clip=1.0), then adjust
  • Test your training loop on a tiny dataset first - verify it can overfit
  • Tuning a model far larger than any you’ve trained? Parametrize it in μP and transfer the learning rate from a small proxy instead of sweeping at full scale
  • Reading a minimum off noisy measurements — an IsoFLOP valley, an LR sweep, a hyperparameter search? Fit a curve and take its optimum; never argmin the raw points. On a flat landscape the single best run is chasing the noise

What’s Next

Module 08: Generation uses the trained model to generate text with various decoding strategies: greedy, sampling, and top-k/top-p.