Optimization, Scaling, and Checkpointing

Derive optimizer updates, then measure PyTorch schedules, clipping, and token budgets.

The pretraining loop from Chapter 05 exposed a gradient and a learning rate, but it treated the optimizer as one line. This chapter opens that line. The NumPy cells preserve update equations on a fixed loss surface; the reusable PyTorch path then applies warmup, cosine decay, global clipping, AdamW parameter groups, and token accounting to the ProofLM decoder.

The experiments are intentionally small. Their purpose is to make update equations and accounting rules testable before they are applied to a larger model.

Optimization algorithms

Use a positive-definite quadratic

f(\theta)=\frac{1}{2}\theta^\top A\theta+b^\top\theta, \qquad \nabla f(\theta)=A\theta+b.

Its minimizer is \theta^*=-A^{-1}b, so optimization error can be measured independently of the update rule. Plain gradient descent uses \theta_{t+1}=\theta_t-\eta g_t. Momentum keeps a velocity v_t=\beta v_{t-1}+g_t and updates with that velocity. AdamW keeps first and second moments, corrects their initialization bias, and applies weight decay directly to the parameter rather than adding it to the gradient.

Implement each state transition explicitly and compare the trajectories from one common starting point.

import numpy as np


SEED = 61
A = np.array([[4.0, 1.0], [1.0, 2.0]])
b = np.array([-2.0, 1.0])
theta_star = -np.linalg.solve(A, b)


def quadratic(theta):
    value = 0.5 * theta @ A @ theta + b @ theta
    gradient = A @ theta + b
    return float(value), gradient


def gd_update(theta, gradient, learning_rate):
    return theta - learning_rate * gradient


def momentum_update(theta, gradient, velocity, learning_rate, beta=0.9):
    velocity = beta * velocity + gradient
    return theta - learning_rate * velocity, velocity


def adamw_update(theta, gradient, first_moment, second_moment, step, learning_rate,
                 beta1=0.9, beta2=0.99, epsilon=1e-8, weight_decay=0.01):
    first_moment = beta1 * first_moment + (1.0 - beta1) * gradient
    second_moment = beta2 * second_moment + (1.0 - beta2) * gradient ** 2
    first_hat = first_moment / (1.0 - beta1 ** step)
    second_hat = second_moment / (1.0 - beta2 ** step)
    adaptive_step = learning_rate * first_hat / (np.sqrt(second_hat) + epsilon)
    theta = (1.0 - learning_rate * weight_decay) * theta - adaptive_step
    return theta, first_moment, second_moment


def run_optimizer(name, steps=80):
    theta = np.array([3.0, -3.0])
    losses = []
    velocity = np.zeros_like(theta)
    first_moment = np.zeros_like(theta)
    second_moment = np.zeros_like(theta)
    for step in range(1, steps + 1):
        value, gradient = quadratic(theta)
        losses.append(value)
        if name == "gd":
            theta = gd_update(theta, gradient, learning_rate=0.18)
        elif name == "momentum":
            theta, velocity = momentum_update(theta, gradient, velocity, learning_rate=0.08)
        elif name == "adamw":
            theta, first_moment, second_moment = adamw_update(
                theta, gradient, first_moment, second_moment, step, learning_rate=0.08
            )
        else:
            raise ValueError(name)
    return theta, np.asarray(losses)


for optimizer_name in ("gd", "momentum", "adamw"):
    final_theta, losses = run_optimizer(optimizer_name)
    print(
        f"{optimizer_name:>8}: loss {losses[0]:7.3f} -> {losses[-1]:9.6f}, "
        f"distance to optimum={np.linalg.norm(final_theta - theta_star):.5f}"
    )
    assert losses[-1] < losses[0]
print("quadratic optimum:", theta_star)

      gd: loss   9.000 -> -1.142857, distance to optimum=0.00000
momentum: loss   9.000 -> -1.141568, distance to optimum=0.03134
   adamw: loss   9.000 -> -1.141892, distance to optimum=0.03144
quadratic optimum: [ 0.71428571 -0.85714286]

The quadratic makes the comparison independent of data sampling and model architecture. Plain gradient descent follows the local slope, momentum smooths successive slopes through its velocity, and AdamW rescales coordinates using their estimated second moments while shrinking the parameter separately. Their final losses need not rank the same way for every learning rate; the useful invariant is that each stateful implementation is evaluated on the same known surface.

Learning-rate warmup

A warmup schedule increases the learning rate from a small value to a peak over W updates. Cosine decay then moves it toward a floor over the remaining S-W updates:

\eta_t=\eta_{\min}+\frac{1}{2}(\eta_{\max}-\eta_{\min}) \left(1+\cos\frac{\pi(t-W)}{S-W}\right).

The schedule is indexed by optimizer updates, not epochs. This matters when batch size changes because one epoch then contains a different number of updates and a different number of processed tokens.

def warmup_cosine(step, warmup_steps, total_steps, maximum_rate, minimum_rate=0.0):
    if not 0 <= step < total_steps:
        raise ValueError("step must lie in [0, total_steps)")
    if warmup_steps < 1 or warmup_steps >= total_steps:
        raise ValueError("warmup_steps must be smaller than total_steps")
    if step < warmup_steps:
        return maximum_rate * (step + 1) / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps - 1)
    cosine = 0.5 * (1.0 + np.cos(np.pi * progress))
    return minimum_rate + (maximum_rate - minimum_rate) * cosine


schedule = np.asarray([
    warmup_cosine(step, warmup_steps=4, total_steps=16, maximum_rate=0.2, minimum_rate=0.02)
    for step in range(16)
])
print("warmup+cosine schedule:", np.round(schedule, 4))
assert np.all(np.diff(schedule[:4]) > 0.0)
assert np.isclose(schedule[3], 0.2)
assert np.isclose(schedule[-1], 0.02)
assert np.all(np.diff(schedule[4:]) <= 1e-12)

warmup+cosine schedule: [0.05   0.1    0.15   0.2    0.2    0.1964 0.1857 0.1689 0.1474 0.1228
 0.0972 0.0726 0.0511 0.0343 0.0236 0.02  ]

The first four rates rise to the configured maximum, then the cosine segment decreases smoothly to the floor. The step + 1 in warmup makes update zero a nonzero but reduced learning rate; the final cosine denominator maps the last valid update exactly to the minimum. A schedule saved in a checkpoint must retain its update position, just as the random state did in Chapter 05.

Gradient clipping

For a model with parameter gradients g_1,\ldots,g_m, define the global norm

\lVert g\rVert_2=\sqrt{\sum_j\lVert g_j\rVert_2^2}.

Global clipping uses one scale for all arrays. It prevents an occasional large batch from producing an update larger than the configured trust threshold while preserving the relative direction among layers. The clipping operation should be logged separately from the raw norm so frequent clipping is visible as a training-dynamics change.

def global_norm(gradients):
    if isinstance(gradients, dict):
        values = gradients.values()
    else:
        values = gradients
    return float(np.sqrt(sum(np.sum(np.asarray(value) ** 2) for value in values)))


def clip_by_global_norm(gradients, maximum_norm):
    norm = global_norm(gradients)
    scale = min(1.0, maximum_norm / (norm + 1e-12))
    if isinstance(gradients, dict):
        clipped = {name: np.asarray(value) * scale for name, value in gradients.items()}
    else:
        clipped = [np.asarray(value) * scale for value in gradients]
    return clipped, norm


gradients = {"embedding": np.array([[3.0, 4.0]]), "projection": np.array([12.0])}
clipped_gradients, raw_norm = clip_by_global_norm(gradients, maximum_norm=5.0)
print("raw norm:", raw_norm, "clipped norm:", global_norm(clipped_gradients))
assert np.isclose(raw_norm, 13.0)
assert np.isclose(global_norm(clipped_gradients), 5.0)
ratio = clipped_gradients["projection"] / gradients["projection"]
assert np.allclose(clipped_gradients["embedding"] / gradients["embedding"], ratio)

raw norm: 13.0 clipped norm: 4.999999999999615

The example has a raw norm of 13 and is scaled to 5; both arrays receive the same factor. That common factor is the behavior needed when a Transformer has layers with very different gradient magnitudes. Clipping does not repair an incorrect objective or a consistently unsuitable learning rate, so the unclipped norm and the clipping frequency remain experiment metrics.

Training counters

An epoch count hides the amount of computation when batch size changes. For a sequence task with L tokens per example, a run with batch size B and U optimizer updates processes approximately

\text{tokens}=B\times L\times U

before accounting for the final partial batch. Hold the update budget fixed in the sweep below, vary batch size and learning rate, and report the resulting token budget explicitly. The objective is a small linear regression only because its loss is quick to evaluate; the accounting rule is the same for language-model batches.

sweep_rng = np.random.default_rng(SEED + 1)
features = sweep_rng.normal(size=(64, 3))
true_weights = np.array([1.5, -0.8, 0.4])
responses = features @ true_weights + 0.1 * sweep_rng.normal(size=64)
SEQUENCE_LENGTH = 5


def regression_loss_and_gradient(weights, batch_features, batch_responses):
    errors = batch_features @ weights - batch_responses
    return float(0.5 * np.mean(errors ** 2)), batch_features.T @ errors / len(errors)


def regression_run(batch_size, learning_rate, updates=40, seed=0, maximum_norm=None):
    weights = np.zeros(features.shape[1])
    random_generator = np.random.default_rng(seed)
    for _ in range(updates):
        indices = random_generator.integers(0, len(features), size=batch_size)
        _, gradient = regression_loss_and_gradient(
            weights, features[indices], responses[indices]
        )
        if maximum_norm is not None:
            gradient, _ = clip_by_global_norm([gradient], maximum_norm)
            gradient = gradient[0]
        weights = gd_update(weights, gradient, learning_rate)
    final_loss = regression_loss_and_gradient(weights, features, responses)[0]
    tokens = batch_size * SEQUENCE_LENGTH * updates
    return final_loss, updates, tokens


results = []
for batch_size in (2, 8, 32):
    for learning_rate in (0.03, 0.1):
        result = regression_run(batch_size, learning_rate, seed=SEED + batch_size)
        results.append((batch_size, learning_rate, *result))
        print(
            f"batch={batch_size:>2}, lr={learning_rate:.2f}, "
            f"updates={result[1]:>2}, tokens={result[2]:>4}, loss={result[0]:.5f}"
        )

assert all(row[3] == 40 for row in results)
assert all(row[4] == row[0] * SEQUENCE_LENGTH * row[3] for row in results)
assert len(results) == 6

batch= 2, lr=0.03, updates=40, tokens= 400, loss=0.20330
batch= 2, lr=0.10, updates=40, tokens= 400, loss=0.01465
batch= 8, lr=0.03, updates=40, tokens=1600, loss=0.18862
batch= 8, lr=0.10, updates=40, tokens=1600, loss=0.00678
batch=32, lr=0.03, updates=40, tokens=6400, loss=0.19016
batch=32, lr=0.10, updates=40, tokens=6400, loss=0.00584

Reusable optimizer path

The fixed-surface derivations identify what each state variable means. The training path below applies the same ideas to a real ProofLM module: matrix weights receive decoupled decay, normalization and bias parameters do not, the schedule is indexed by optimizer update, and the raw gradient norm is recorded before clipping.

import torch

from proof_lm.model import DecoderConfig, ProofLM
from proof_lm.optimization import (
    adamw_parameter_groups,
    clip_gradients,
    warmup_cosine_factor,
)

torch.manual_seed(SEED)
optimizer_model = ProofLM(
    DecoderConfig(
        vocab_size=24,
        context_length=8,
        n_layers=1,
        d_model=32,
        n_heads=4,
        d_ff=64,
    )
)
optimizer = torch.optim.AdamW(
    adamw_parameter_groups(optimizer_model, weight_decay=0.1),
    lr=0.02,
)
optimizer_inputs = torch.randint(0, 24, (4, 8))
optimizer_targets = torch.roll(optimizer_inputs, shifts=-1, dims=1)
learning_rates = []
raw_norms = []
for update in range(6):
    factor = warmup_cosine_factor(
        update, warmup_updates=2, total_updates=6, minimum_ratio=0.2
    )
    for group in optimizer.param_groups:
        group["lr"] = 0.02 * factor
    optimizer.zero_grad(set_to_none=True)
    _, optimizer_loss = optimizer_model(
        optimizer_inputs,
        labels=optimizer_targets,
    )
    optimizer_loss.backward()
    raw_norms.append(clip_gradients(optimizer_model.parameters(), 0.5))
    optimizer.step()
    learning_rates.append(optimizer.param_groups[0]["lr"])
processed_tokens = optimizer_inputs.numel() * len(learning_rates)
print("learning rates:", [round(value, 5) for value in learning_rates])
print("raw gradient norms:", [round(value, 4) for value in raw_norms])
print("processed tokens:", processed_tokens)
assert learning_rates[0] < learning_rates[1]
assert learning_rates[-1] < learning_rates[1]
assert all(value >= 0 for value in raw_norms)
assert processed_tokens == 4 * 8 * 6

learning rates: [0.01, 0.02, 0.02, 0.01766, 0.012, 0.00634]
raw gradient norms: [1.9901, 1.2494, 2.7686, 2.1888, 1.9198, 1.0681]
processed tokens: 192

The sweep keeps updates fixed, so the batch-32 runs process sixteen times as many examples and tokens as batch-2 runs. A lower final loss at a larger batch may therefore reflect a larger token budget rather than a more efficient optimizer. Report both axes when comparing learning rates, and add wall-clock time or hardware throughput when the experiment is intended to make a scaling claim.

Summary

• Plain gradient descent, momentum, and AdamW differ in the state they carry and in where weight decay enters the update.
• Warmup and cosine decay are functions of optimizer updates; changing batch size changes the token count per scheduled step.
• Global gradient clipping applies one scale to all parameter arrays and should be logged with the raw norm.
• The fixed-surface and regression experiments make optimizer behavior and learning-rate sensitivity measurable without a deep-learning framework.
• Batch size, updates, examples, and tokens are separate accounting fields; epochs alone cannot compare these runs.

Chapter 07 packages loss, calibration, memorization, overlap, and diversity measurements into composable evaluation functions.

Exercises

Use the exercises to test the chapter’s invariants and connect the derivations to the reusable implementation. Solutions are hidden in the notebook source and are available through the course tooling when needed.

[P6.1] AdamW update

AdamW derivation. Write the first-moment and second-moment recurrences, explain the bias-correction factors at step one, and state where decoupled weight decay appears in the parameter update.

[P6.2] Token-budget accounting

Budget accounting. A run uses batch size 16, sequence length 128, and 250 optimizer updates. Compute processed tokens and explain why comparing it with a batch-size-4 run by epoch count alone is misleading.

[P6.3]

Update-indexed schedules. A cosine schedule is a function of optimizer updates, not epochs. Given warmup_updates=2 and total_updates=6, explain why the first two learning-rate factors differ from the later factors, and state which token-accounting fields must be reported when batch size changes.

def schedule_report(batch_size, context_length, updates):
    # Return processed tokens and the schedule's final update index.
    pass