3.10 Learning Rate Schedules, Warmup, Batch Size & Hyperparameters¶

Hyperparameter tuning is one of the most impactful — and most under-documented — parts of pretraining a large language model. A 10x learning rate error will destroy a run that would otherwise converge cleanly; the right schedule can shave percentage points off final perplexity. Yet the wisdom lives mostly in appendices of papers and in the institutional memory of ML engineering teams.

This chapter makes that implicit knowledge explicit. We cover the full pipeline: why warmup is mandatory at scale, the major schedule families and when to use each, how batch size interacts with learning rate and what the critical batch size tells you about compute efficiency, and finally how muP (maximal-update parameterization) lets you tune hyperparameters on a small model and transfer them to a large one. Every section includes runnable code and concrete numbers.

Related chapters you should read in tandem: Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo covers the adaptive optimizer mechanics that schedules ride on top of; Training Stability, Loss Spikes & Debugging Large Runs covers what happens when you get the schedule wrong; Scaling Laws: Kaplan, Chinchilla & Beyond gives the big picture that informs budget allocation decisions. For every number in this chapter pinned down for one concrete run, see the capstone’s Optimizer & Schedule: Muon + MuonClip and Warmup-Stable-Decay, which sets the peak LR, warmup, and WSD phase lengths for the ~100M-parameter Stack-100M model by measurement rather than by folklore.

Why Schedules and Warmup Exist¶

Deep learning optimizers like Adam carry two moving-average estimates: first moment \(m_t\) (gradient direction) and second moment \(v_t\) (gradient magnitude squared). Both are initialized to zero. During the first few hundred steps, \(v_t\) is still a noisy underestimate of the true gradient variance — the bias-correction terms in Adam partially compensate, but the effective learning rate is still erratic early in training.

For large models, this cold-start instability is catastrophic. At step 1 of a run with 1 billion parameters and a 4096-token context, the weight matrices are random; the gradient norms are large and wildly variable across layers. Applying the full target learning rate immediately produces parameter updates large enough to push weights into regions where softmax logits saturate, norms explode, or residual magnitudes collapse — none of which recover easily. Empirically, runs without warmup frequently spike or diverge within the first thousand steps.

Warmup solves this by linearly ramping the effective learning rate from near-zero to the target value over a set number of steps, giving the optimizer time to calibrate its momentum estimates and giving the network time to find a reasonable initialization basin before the full update magnitude kicks in.

There is a second, complementary mechanism worth holding in your head, because it explains why warmup helps even for optimizers with no bias-correction problem. What actually matters for stability is the relative update size \(\lVert \Delta\theta \rVert / \lVert \theta \rVert\) per layer. Adam’s update has a roughly fixed per-element magnitude of about \(\eta\) (the \(\sqrt{v_t}\) denominator normalizes the gradient scale away), while the weight norm \(\lVert \theta \rVert\) at initialization is small and only grows as training proceeds. So at step 0 a full-magnitude LR is an enormous fractional change to every layer; a few thousand steps later the same LR is a modest one. Warmup is the schedule that keeps that ratio bounded while the weights are still small — an analysis developed carefully by Kosson et al. (Why Warmup the Learning Rate?, 2024). This framing also predicts the practical corollary you will meet below: anything that resets the weight/optimizer scale relationship (a fresh optimizer state, a sudden batch-size jump) calls for re-warming.

The interplay with weight initialization is important. Standard initialization schemes (Transformers: The Transformer Block) ensure variance-preserving forward passes at step 0, but they do not ensure that the gradient landscape is well-behaved. Warmup effectively treats the first \(T_w\) steps as a coarser form of initialization.

The Major Schedule Families¶

Linear Schedule¶

The simplest useful schedule: ramp linearly from \(\eta_{\min}\) to \(\eta_{\max}\) during warmup, then decay linearly to \(\eta_{\min}\) over the remaining steps.

Linear decay is fast to implement and interpretable. It was widely used in BERT-era fine-tuning but is now rarely the first choice for pretraining because it decays too aggressively in the middle of the run.

Cosine Annealing¶

For years the default pretraining schedule, and still one of the two standard choices in 2026 (alongside WSD, below). After warmup, the learning rate follows the right half of a cosine curve, smoothly decaying to a floor \(\eta_{\min}\) (usually \(\eta_{\max}/10\) or a small constant like \(1\text{e-}5\)):

Key properties: - Spends most of the budget near the peak learning rate (the cosine curve is flat near its maximum), which means the network sees aggressive gradient steps for most of training — good for exploration. - The tail naturally slows down near the end, allowing fine-grained convergence. - The exact shape is not sensitive to the precise \(T\) you use, as long as \(T\) is approximately correct.

The main weakness: cosine requires knowing the total token budget \(T\) upfront. If you extend the run or add a second phase, you need to restart the schedule or accept a discontinuity.

Cosine with Restarts (SGDR)¶

Introduced by Loshchilov & Hutter (SGDR: Stochastic Gradient Descent with Warm Restarts, 2016), this runs multiple cosine cycles with exponentially increasing cycle length. Each restart re-warms the LR to \(\eta_{\max}\) and decays again. It was influential in CV but is rarely used in modern LLM pretraining because the LR spike at each restart causes loss spikes and the benefits for language modeling are unclear.

Warmup-Stable-Decay (WSD)¶

WSD, popularized by MiniCPM (Hu et al., 2024) and used in several other recent models, divides training into three explicit phases:

1. Warmup (\(T_w\) steps): linear ramp from near-zero to \(\eta_{\max}\).
2. Stable (\(T_s\) steps): constant \(\eta_{\max}\).
3. Decay (\(T_d\) steps): cosine or linear decay to \(\eta_{\min}\).

where \(f\) is a decay function. Cosine and linear both work; MiniCPM reported that a \(1-\sqrt{\cdot}\) shape (fast at first, flattening at the end) beats linear, and unlike cosine pretraining practice the WSD decay usually goes all the way to (near) zero rather than stopping at \(\eta_{\max}/10\) — the last bit of decay is where the characteristic extra loss drop lives. A common budget is \(T_d \approx 10\%\) of total steps.

The insight driving WSD is that most of the loss reduction happens in the stable phase, and the decay phase mainly “polishes” the model. This decoupling buys three things:

• Late binding of the run length. You can train in the stable phase for as long as resources allow, then trigger the decay when you’re ready to finalize — no commitment to a total step count at the start.
• Branching. One stable checkpoint can seed several independent decay runs (different data mixes, different decay lengths) that you compare head-to-head for a fraction of a full run’s cost. A decayed checkpoint, by contrast, is spent: resuming high-LR training from it throws away the polish.
• A natural seam for a data-mixture switch. Because the highest-value tokens are the ones seen at the lowest LR, the decay window is exactly where you swap the broad web mix for premium data. This is the whole idea behind mid-training/annealing (Data Mixing, Domain Weighting & Curriculum), and the capstone does precisely this: the WSD decay phase is the mid-training phase (Mid-Training).

Hägele et al. (Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations, 2024) ran the careful comparison and found that constant-LR-plus-cooldown matches cosine’s final loss at equal compute while producing a whole family of usable checkpoints along the way — which is why WSD is now a first-class choice rather than a curiosity.

RSqrt (Inverse Square Root)¶

Used in original Transformer training (Vaswani et al., 2017) with the combined formula \(\eta(t) = d_{\text{model}}^{-0.5} \cdot \min(t^{-0.5},\; t \cdot T_w^{-1.5})\). The rsqrt schedule never fully plateaus — the LR continues to slowly decrease throughout training. It works well for smaller models and shorter runs but tends to decay too quickly for billion-scale pretraining.

Implementing Schedules From Scratch¶

import math
import torch
from torch.optim.lr_scheduler import LambdaLR

def get_cosine_schedule_with_warmup(
    optimizer: torch.optim.Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    min_lr_fraction: float = 0.1,   # eta_min = eta_max * min_lr_fraction
) -> LambdaLR:
    """
    Cosine annealing with linear warmup.
    The LambdaLR multiplier is relative to the base LR in the optimizer.
    """
    def lr_lambda(current_step: int) -> float:
        # --- Warmup phase ---
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))

        # --- Cosine decay phase ---
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        # progress in [0, 1]; cosine from 1 → min_lr_fraction
        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
        # Rescale so the floor is min_lr_fraction
        return min_lr_fraction + (1.0 - min_lr_fraction) * cosine_decay

    return LambdaLR(optimizer, lr_lambda)


def get_wsd_schedule(
    optimizer: torch.optim.Optimizer,
    num_warmup_steps: int,
    num_stable_steps: int,
    num_decay_steps: int,
    min_lr_fraction: float = 0.1,
) -> LambdaLR:
    """
    Warmup-Stable-Decay (WSD) schedule.
    Advantage: total training length can be decided late — just extend stable phase.
    """
    T_w = num_warmup_steps
    T_s = num_stable_steps
    T_d = num_decay_steps

    def lr_lambda(step: int) -> float:
        if step < T_w:
            # Linear warmup
            return float(step) / float(max(1, T_w))
        elif step < T_w + T_s:
            # Stable plateau at peak LR
            return 1.0
        else:
            # Cosine decay to floor
            decay_progress = float(step - T_w - T_s) / float(max(1, T_d))
            decay_progress = min(decay_progress, 1.0)  # clamp at end
            cosine = 0.5 * (1.0 + math.cos(math.pi * decay_progress))
            return min_lr_fraction + (1.0 - min_lr_fraction) * cosine

    return LambdaLR(optimizer, lr_lambda)


# ---- Quick smoke test ----
if __name__ == "__main__":
    model = torch.nn.Linear(10, 10)
    # Base LR that the scheduler multiplies against
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

    scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=100,
        num_training_steps=1000,
        min_lr_fraction=0.1,
    )

    lrs = []
    for step in range(1000):
        optimizer.step()
        lrs.append(optimizer.param_groups[0]["lr"])
        scheduler.step()

    # Verify: step 50 should be ~50% of peak; step 999 should be near min
    assert abs(lrs[50] / lrs[99] - 50/100) < 0.01, "warmup slope wrong"
    assert lrs[-1] < lrs[99] * 0.15, "floor not reached"
    print(f"Peak LR: {max(lrs):.2e}, Final LR: {lrs[-1]:.2e}")
    # Output: Peak LR: 3.00e-04, Final LR: 3.00e-05

The Library Equivalents¶

You should write the above once to own the mechanics; in a real run you will usually call a library. Every schedule in this chapter already exists in the standard stack, and knowing the exact knob name in each framework is most of the practical battle.

import torch

model = torch.nn.Linear(16, 16)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95))

# --- HuggingFace transformers: the most-used scheduler factory in the ecosystem ---
# (requires `pip install transformers`; get_wsd_schedule needs a recent version)
from transformers import (
    get_linear_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_constant_schedule_with_warmup,
    get_inverse_sqrt_schedule,   # the rsqrt schedule from Vaswani et al.
    get_wsd_schedule,            # Warmup-Stable-Decay
)

sched = get_cosine_schedule_with_warmup(
    optimizer, num_warmup_steps=2000, num_training_steps=100_000
)
# NOTE: get_cosine_schedule_with_warmup decays to ZERO, not to a 10% floor.
# For a floor, pass a min-LR argument if your version exposes one, or use the
# from-scratch LambdaLR above. Check the signature for the version you pinned —
# get_wsd_schedule's phase arguments in particular have changed across releases.

# Via the Trainer, the same thing is one string:
#   TrainingArguments(lr_scheduler_type="cosine", warmup_ratio=0.02,
#                     learning_rate=3e-4, max_grad_norm=1.0, weight_decay=0.1)

# --- Pure PyTorch: compose warmup and decay with SequentialLR ---
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR

warmup = LinearLR(optimizer, start_factor=1e-8, end_factor=1.0, total_iters=2000)
decay = CosineAnnealingLR(optimizer, T_max=100_000 - 2000, eta_min=3e-5)
sched = SequentialLR(optimizer, schedulers=[warmup, decay], milestones=[2000])

Framework-level trainers expose the same schedules as config rather than code:

# DeepSpeed config (ds_config.json): warmup + linear decay, no Python needed.
scheduler:
  type: WarmupDecayLR
  params:
    warmup_min_lr: 0.0
    warmup_max_lr: 3.0e-4
    warmup_num_steps: 2000
    total_num_steps: 100000

# Megatron-LM: the schedule is entirely CLI flags on pretrain_gpt.py.
# --lr-decay-style accepts constant | linear | cosine | inverse-square-root;
# recent versions also ship a WSD style with its own decay-length flags.
# Use --lr-warmup-fraction 0.02 instead of --lr-warmup-iters to express warmup
# as a fraction of --lr-decay-iters.
SCHEDULE_ARGS=(
  --lr 3.0e-4
  --min-lr 3.0e-5
  --lr-decay-style cosine
  --lr-warmup-iters 2000
  --lr-decay-iters 100000
  --clip-grad 1.0
  --weight-decay 0.1
  --adam-beta1 0.9
  --adam-beta2 0.95
)
echo "${SCHEDULE_ARGS[@]}"   # torchrun ... pretrain_gpt.py "${SCHEDULE_ARGS[@]}"

PyTorch’s torchtitan reference stack similarly exposes warmup/stable/decay phase lengths and a decay type in its TOML job config, and TRL/Axolotl fine-tuning configs forward lr_scheduler_type straight through to the transformers factories above. See Megatron-LM, DeepSpeed & Parallelism in Practice for the surrounding launch machinery.

Making the Schedule Restart-Safe¶

Long pretraining runs get preempted, and the schedule must survive the restart exactly. There is a footgun here: LambdaLR.state_dict() cannot serialize a plain-function lr_lambda (it is skipped, with a warning), so a naive torch.save(scheduler.state_dict()) round-trip can silently restore only last_epoch — and if you rebuild the scheduler with different num_training_steps on resume, you get a different curve with no error.

The robust pattern used by most production trainers, including the capstone’s wsd_lr, is to make the schedule a pure function of the global step and set the LR by hand each iteration. There is no hidden state to checkpoint beyond the step counter you were already saving:

import math

def lr_at_step(step: int, peak_lr: float, warmup: int, total: int,
               min_lr_fraction: float = 0.1) -> float:
    """Stateless cosine-with-warmup: the ONLY input is the global step."""
    if step < warmup:
        return peak_lr * (step + 1) / warmup
    progress = min(1.0, (step - warmup) / max(1, total - warmup))
    cos = 0.5 * (1.0 + math.cos(math.pi * progress))
    return peak_lr * (min_lr_fraction + (1.0 - min_lr_fraction) * cos)


# In the training loop, before optimizer.step():
#   lr = lr_at_step(global_step, 3e-4, 2000, 100_000)
#   for g in optimizer.param_groups:
#       g["lr"] = lr * g.get("lr_scale", 1.0)   # lr_scale carries muP / per-group ratios

The lr_scale multiplier in that last line matters whenever different parameter groups run at different peak LRs — muP’s per-layer width factors, or a Muon/AdamW hybrid where the two groups sit at a fixed ratio to each other. Writing one scalar LR into every group collapses the ratio and is a classic silent bug; see Optimizer & Schedule for the concrete Muon-plus-AdamW case. Checkpointing more broadly is covered in Checkpointing, Fault Tolerance & Long-Running Jobs.

Learning Rate vs. Batch Size Scaling¶

The Linear Scaling Rule¶

When you increase the batch size \(B\) by a factor \(k\), each gradient step is an average over \(k\) times more samples, reducing variance by \(\sqrt{k}\) and the signal-to-noise ratio effectively improves. To maintain the same training dynamics — the same total parameter update magnitude per unit of data — you should also scale the learning rate:

This rule was derived and validated for SGD with momentum, for small-to-moderate batch changes (say, 256 to 4096). The intuition: with \(k\times\) larger batches, each step is \(k\times\) more expensive in wall time (same flops per sample), but covers \(k\times\) more data, so to move at the same “rate through the data manifold,” the step size should scale linearly.

The Square-Root Scaling Rule¶

Two independent arguments land on the square root, and for LLM pretraining they are the ones that matter.

It is the correct rule for adaptive optimizers. This is the point most treatments of the linear rule get wrong. Adam’s update divides the gradient by \(\sqrt{v_t}\), an estimate of the gradient’s second moment — so the noise that averaging removes appears inside a square root, and the noise-preserving scaling becomes \(\eta \propto \sqrt{k}\) rather than \(\eta \propto k\). Malladi et al. (On the SDEs and Scaling Rules for Adaptive Gradient Algorithms, 2022) derive this from the stochastic-differential-equation limit of Adam and RMSProp and verify it empirically. Since essentially all LLM pretraining uses AdamW (or a normalized-update optimizer like Muon, whose update magnitude is fixed by construction and so behaves similarly), square-root scaling should be your default when changing batch size, and linear scaling the special case you reach for only with SGD.

It is also the conservative choice near saturation. Independently of the optimizer, at very large batch the gradient variance stops falling as \(1/B\) — you hit the intrinsic noise floor of the data distribution rather than sampling noise — so any rule that keeps growing the LR with \(B\) eventually over-scales. The square root degrades gracefully where linear does not.

The capstone uses exactly this rule to move a peak LR measured on a small probe batch (65,536 tokens/step) to the real run’s 524,288 tokens/step: \(\times\sqrt{8} \approx 2.83\) (Optimizer & Schedule).

Critical Batch Size¶

The critical batch size \(B^*\) is the regime boundary between these two laws. It was formalized by McCandlish et al. (An Empirical Model of Large-Batch Training, 2018). The key quantity is the gradient noise scale:

where \(\Sigma\) is the covariance of the per-sample gradient and \(G\) is the mean gradient. When \(B \ll B_{\text{noise}}\), batches are too small to average out noise — increasing batch size linearly reduces steps needed. When \(B \gg B_{\text{noise}}\), you are in the saturated regime where more data per step doesn’t help; gradient noise is already small and increasing \(B\) wastes compute.

The practical takeaway: there is an optimal batch size for a given compute budget. Doubling the batch size beyond \(B^*\) halves your throughput efficiency. For typical LLM pretraining, \(B^*\) for cross-entropy loss on language is on the order of a few million tokens — consistent with the token batches used in Chinchilla-optimal runs — while for a ~100M model it sits closer to a few hundred thousand tokens, which is why the capstone lands on 524,288 tokens/step.

An important 2024–2025 refinement: the critical batch size scales primarily with the amount of data trained on, not with parameter count. Zhang et al. (How Does Critical Batch Size Scale in Pre-training?, 2024) hold one of the two fixed at a time and find \(B^*\) grows with tokens seen and is close to flat in model size once data is controlled. The practical consequence is that a long over-trained run of a small model can tolerate a surprisingly large batch, and that \(B^*\) increases as your run proceeds — a reason ramping batch size over training (see the practitioner tip at the end of this chapter) is principled rather than folklore.

You do not have to take \(B^*\) on faith: the noise scale is directly measurable from two batch sizes. McCandlish et al. give unbiased estimators for \(\lVert G\rVert^2\) and \(\operatorname{tr}(\Sigma)\) from the squared gradient norms at a small and a large batch, and the ratio is \(B_{\text{noise}}\). Run this for a few hundred steps on your actual model and data before you pick a batch size.

import torch


@torch.no_grad()
def _sq_grad_norm(model) -> float:
    """Squared global L2 norm of the current .grad tensors."""
    return float(sum((p.grad.detach() ** 2).sum() for p in model.parameters()
                     if p.grad is not None))


def gradient_noise_scale(model, loss_fn, batch, b_small: int, b_big: int) -> float:
    """
    Estimate B_noise = tr(Sigma) / ||G||^2 (McCandlish et al., 2018) from one
    large batch and its first `b_small` examples.

    The estimators exploit E[||g_B||^2] = ||G||^2 + tr(Sigma)/B, evaluated at two
    batch sizes and solved as a 2x2 linear system:
        ||G||^2   ~ (b_big*||g_big||^2 - b_small*||g_small||^2) / (b_big - b_small)
        tr(Sigma) ~ (||g_small||^2 - ||g_big||^2) / (1/b_small - 1/b_big)
    Both are noisy per-step; average them over ~100 steps before dividing.
    """
    model.zero_grad(set_to_none=True)
    loss_fn(model, batch[:b_small]).backward()
    g_small = _sq_grad_norm(model)

    model.zero_grad(set_to_none=True)
    loss_fn(model, batch[:b_big]).backward()
    g_big = _sq_grad_norm(model)

    g_norm_sq = (b_big * g_big - b_small * g_small) / (b_big - b_small)
    trace_sigma = (g_small - g_big) / (1.0 / b_small - 1.0 / b_big)
    return trace_sigma / max(g_norm_sq, 1e-12)


if __name__ == "__main__":
    torch.manual_seed(0)
    model = torch.nn.Linear(32, 4)
    data = torch.randn(256, 32)
    targets = torch.randint(0, 4, (256,))

    def loss_fn(m, idx_slice):
        n = idx_slice.shape[0]
        return torch.nn.functional.cross_entropy(m(data[:n]), targets[:n])

    # Average the two moments over several draws, then divide (never average ratios).
    est = sum(gradient_noise_scale(model, loss_fn, data, 8, 128) for _ in range(20)) / 20
    print(f"B_noise estimate: {est:.1f} examples")
    # In a real run: repeat every few hundred steps and watch B_noise GROW as
    # the gradient shrinks -- that growth is the signal to ramp the batch size.

In a distributed run the same estimate falls out for free: compare the per-rank gradient norm (small batch) with the post-all-reduce global gradient norm (large batch), which is what production monitoring typically logs. See Distributed Training I.

Gradient Accumulation¶

When you can’t fit the full desired batch size in GPU memory in a single forward-backward pass, gradient accumulation (GA) simulates a larger effective batch by running \(k\) micro-batches before calling optimizer.step():

The effective batch size is micro_batch_size * accumulation_steps * world_size. Gradient accumulation is mathematically equivalent to a full batch if loss is averaged (not summed) — a subtle but critical distinction.

import torch
import torch.nn as nn
from contextlib import nullcontext

def train_step_with_grad_accumulation(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    data_loader,
    accumulation_steps: int = 4,
    scaler=None,  # Optional GradScaler for mixed precision
    device: str = "cuda",
) -> float:
    """
    One effective step = accumulation_steps micro-forward-backward passes.
    Returns mean loss over the effective batch.
    """
    model.train()
    optimizer.zero_grad()
    total_loss = 0.0

    for micro_step, batch in enumerate(data_loader):
        input_ids = batch["input_ids"].to(device)
        labels = batch["labels"].to(device)

        # Only sync gradients across DDP replicas on the LAST micro-step.
        # This avoids expensive all-reduce on every micro-step.
        sync_ctx = (
            model.no_sync()
            if hasattr(model, "no_sync") and micro_step < accumulation_steps - 1
            else nullcontext()
        )

        with sync_ctx:
            # Use autocast for bf16/fp16 if scaler is provided
            amp_ctx = torch.autocast("cuda", dtype=torch.bfloat16) if scaler else nullcontext()
            with amp_ctx:
                logits = model(input_ids)
                # CRITICAL: divide by accumulation_steps so effective batch
                # average == sum-of-microbatch-averages / k.
                # If your loss already averages over tokens in the microbatch,
                # this gives the right weight for each sample.
                loss = nn.functional.cross_entropy(
                    logits.view(-1, logits.size(-1)),
                    labels.view(-1),
                ) / accumulation_steps

            if scaler:
                scaler.scale(loss).backward()
            else:
                loss.backward()

        total_loss += loss.item()

        if (micro_step + 1) == accumulation_steps:
            break  # done accumulating

    # Unscale, clip, step
    if scaler:
        scaler.unscale_(optimizer)
    # Gradient clipping (see next section)
    nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    if scaler:
        scaler.step(optimizer)
        scaler.update()
    else:
        optimizer.step()

    return total_loss * accumulation_steps  # report unscaled mean loss

One important caveat with distributed training: gradient synchronization (the all-reduce across data-parallel ranks) should happen only at the final accumulation step. Using model.no_sync() in PyTorch DDP avoids the all-reduce on intermediate steps. See Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP for the full picture.

Weight Decay and Gradient Clipping¶

Weight Decay¶

Weight decay (L2 regularization) penalizes large weights and prevents any single parameter from dominating. In modern deep learning it is implemented as decoupled weight decay (as in AdamW, Loshchilov & Hutter, 2019):

The decay is applied to the raw parameter values, not the gradient estimate, which prevents the adaptive scaling of Adam from interfering with regularization. Typical values: \(\lambda = 0.1\) for pretraining (used in GPT-3, Llama, and most modern runs). Embeddings and bias terms are usually excluded from decay since they have a different scale and semantics.

Notice what the update equation implies: the shrinkage per step is \(\lambda\eta\), not \(\lambda\). Weight decay and the learning rate are therefore not independent knobs — they multiply. Two consequences follow, and both bite in practice:

• Your schedule decays your regularization too. As \(\eta\) falls along a cosine or a WSD decay phase, the effective pull toward zero falls with it. The characteristic timescale over which a weight is forgotten is roughly \(\tau \approx 1/(\lambda\eta)\) steps; at \(\lambda=0.1,\ \eta=3\text{e-}4\) that is about 33,000 steps — comparable to a whole run, which is why weight decay in LLM pretraining is a slow, global effect rather than a per-batch one.
• If you change the peak LR, you have changed the regularization strength. When comparing two LR settings on equal footing, either hold \(\lambda\eta\) fixed or report that you did not.

This coupling is also why weight decay works at all in a regime with essentially no overfitting: for one-epoch LLM pretraining, decay is not primarily preventing memorization of a re-seen training set. Andriushchenko et al. (Why Do We Need Weight Decay in Modern Deep Learning?, 2023) argue its main role there is to keep weight norms — and hence the effective per-step learning rate seen by each layer — in a favorable range, i.e. it is an optimization-shaping term more than a classical regularizer. Practically: keep \(\lambda = 0.1\), exclude 1D parameters, and treat \(\lambda\) and \(\eta\) as one joint decision.

def get_optimizer_with_decay(
    model: nn.Module,
    lr: float,
    weight_decay: float = 0.1,
    betas: tuple = (0.9, 0.95),  # common for LLM pretraining
    eps: float = 1e-8,
) -> torch.optim.AdamW:
    """
    AdamW with weight decay applied only to weight matrices (not biases/norms).
    betas=(0.9, 0.95) is standard for LLM pretraining — beta2=0.999 (default)
    can slow adaptation to gradient changes late in training.
    """
    # Partition params: decay weights but NOT biases, LayerNorm params, embeddings
    decay_params = []
    no_decay_params = []

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        # Common no-decay criteria: 1D params (bias, LN scale/bias),
        # and sometimes embedding matrices
        if param.ndim == 1 or "bias" in name or "norm" in name.lower():
            no_decay_params.append(param)
        else:
            decay_params.append(param)

    param_groups = [
        {"params": decay_params, "weight_decay": weight_decay},
        {"params": no_decay_params, "weight_decay": 0.0},
    ]

    optimizer = torch.optim.AdamW(
        param_groups, lr=lr, betas=betas, eps=eps, fused=True  # fused=True uses fast CUDA kernel
    )
    return optimizer

Setting betas=(0.9, 0.95) rather than the default (0.9, 0.999) is a commonly used pretraining choice — the lower \(\beta_2\) makes the optimizer more responsive to recent gradient magnitude changes, which helps during the warmup phase and when learning rate jumps occur.

Gradient Clipping¶

Gradient clipping prevents parameter updates from being catastrophically large due to occasional gradient spikes (common when training on noisy web data or when the schedule is aggressive). The global-norm clip is by far the most common form:

where \(c\) is the clip threshold (typically 1.0) and \(\|g\|_2 = \sqrt{\sum_i g_i^2}\) is the global L2 norm across all parameters. This preserves gradient direction while bounding the step magnitude.

A persistent misconception: clipping is not a substitute for a well-designed schedule. If your global gradient norm is routinely hitting the clip threshold on more than ~20% of steps, something is wrong with your initialization, learning rate, or data distribution.

# Gradient clipping with norm tracking for monitoring
def clip_and_log_grad_norm(
    model: nn.Module,
    max_norm: float = 1.0,
) -> float:
    """Returns the pre-clip gradient norm for monitoring dashboards."""
    # Computes global L2 norm across all parameters
    total_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_norm)
    return float(total_norm)

clip_grad_norm_ is a one-liner in PyTorch, but it’s worth implementing from scratch once so the mechanics are unambiguous. The whole thing is two steps: (1) compute a single global L2 norm as the square root of the summed sum-of-squares across all parameters — not a separate norm per parameter — and (2) scale every gradient in place by \(\min(1,\, c / \|g\|)\), i.e. you only ever scale gradients down, never up.

@torch.no_grad()
def clip_grad_norm_from_scratch(params, max_norm: float = 1.0, eps: float = 1e-6) -> float:
    """From-scratch global-norm clip; matches torch.nn.utils.clip_grad_norm_."""
    grads = [p.grad for p in params if p.grad is not None]
    # ONE global L2 norm across ALL params, not per-parameter norms.
    total_norm = torch.sqrt(sum((g.detach() ** 2).sum() for g in grads))
    clip_coef = max_norm / (total_norm + eps)   # torch adds eps for stability
    if clip_coef < 1.0:                          # only ever scale DOWN
        for g in grads:
            g.mul_(clip_coef)
    return float(total_norm)


if __name__ == "__main__":
    torch.manual_seed(0)
    layer_a = nn.Linear(64, 64)
    layer_b = nn.Linear(64, 64)
    params = list(layer_a.parameters()) + list(layer_b.parameters())

    x = torch.randn(8, 64)
    loss = (layer_b(layer_a(x)) ** 2).sum()
    loss.backward()

    # Save the un-clipped grads so both implementations start from the same state.
    original_grads = [p.grad.clone() for p in params]

    total_norm_scratch = clip_grad_norm_from_scratch(params, max_norm=1.0)
    scratch_clipped_grads = [p.grad.clone() for p in params]

    for p, g in zip(params, original_grads):
        p.grad.copy_(g)  # reset to un-clipped values before the torch reference call
    total_norm_torch = float(nn.utils.clip_grad_norm_(params, max_norm=1.0))
    torch_clipped_grads = [p.grad.clone() for p in params]

    assert abs(total_norm_scratch - total_norm_torch) < 1e-5
    for g_scratch, g_torch in zip(scratch_clipped_grads, torch_clipped_grads):
        assert torch.allclose(g_scratch, g_torch, atol=1e-6)
    print(f"scratch total_norm={total_norm_scratch:.6f}, torch total_norm={total_norm_torch:.6f}")
    # Expected: both report the same total_norm and identical clipped grads

Always log the pre-clip gradient norm at every step. A sudden spike — say, from 0.5 to 50 — is an early warning of a loss spike before it becomes visible in the loss itself. See Training Stability, Loss Spikes & Debugging Large Runs for the full debugging playbook.

muP: Maximal-Update Parameterization¶

The Problem with Standard Parameterization at Scale¶

When you tune hyperparameters on a small model (say, 125M parameters) and transfer them to a large model (7B+), something breaks. The optimal learning rate changes because the width of the network changes: wider networks have larger forward activations and can receive larger gradient signal, so they need smaller learning rates to maintain stable training dynamics.

In standard parameterization (SP), the optimal LR scales roughly as \(1/\sqrt{d}\) or \(1/d\) depending on the layer type, where \(d\) is the hidden dimension. This means every time you scale the model, you have to re-tune LR.

muP: What Changes¶

Maximal-update parameterization (Yang et al., Tensor Programs V, 2022) is a reparameterization of the network that makes the optimal hyperparameters — especially learning rate and initialization scale — independent of model width. You can then:

1. Run a full HP sweep on a tiny “proxy model” (e.g., width 256).
2. Transfer the optimal HPs directly to the full-scale run at width 4096 or beyond.

The key changes relative to standard PyTorch initialization:

The precise prescription comes from requiring that all feature updates \(\Delta h^{(l)}\) (the pre-activation change per step) remain \(O(1)\) as width \(d \to \infty\) — the “maximal” in muP means every neuron participates maximally in learning without causing instability.

import torch
import torch.nn as nn
import math


class MuPLinear(nn.Linear):
    """
    Linear layer with muP-compatible initialization and LR scaling.
    In muP:
      - hidden layers: init std = base_std / sqrt(fan_in), LR *= 1/fan_in
      - readout layer: init std = base_std / fan_in, LR *= 1/fan_in
    We implement LR scaling via a per-parameter LR multiplier convention
    compatible with mup (microsoft/mup on GitHub).
    """
    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        is_readout: bool = False,
        base_std: float = 1.0,
        inf_width: int = None,  # width at "infinite" (reference) model scale
    ):
        super().__init__(in_features, out_features, bias)
        self.is_readout = is_readout
        self.inf_width = inf_width or in_features

        # muP initialization
        if is_readout:
            # Readout: std ∝ 1/d so activations stay O(1)
            std = base_std / in_features
        else:
            # Hidden: same as standard He/fan-in but with explicit inf_width scaling
            std = base_std / math.sqrt(in_features)

        nn.init.normal_(self.weight, mean=0.0, std=std)
        if bias:
            nn.init.zeros_(self.bias)

    def get_lr_multiplier(self) -> float:
        """
        Returns the per-layer LR multiplier.
        Base LR should be tuned at proxy (small) model; this scales it correctly.
        At proxy model width d_proxy, multiplier = 1.
        At width d >> d_proxy, multiplier = d_proxy / d.
        For simplicity, return 1/in_features (absorbed into optimizer param groups).
        """
        return 1.0 / self.in_features


def build_mup_optimizer(
    model: nn.Module,
    base_lr: float,
    proxy_width: int,
    weight_decay: float = 0.1,
) -> torch.optim.AdamW:
    """
    Build AdamW where each layer's effective LR = base_lr * (proxy_width / layer_width).
    base_lr is tuned at proxy_width; this transfers the HP to any larger model.
    """
    param_groups = []

    for name, module in model.named_modules():
        if isinstance(module, MuPLinear):
            # Scale LR inversely with layer width to maintain muP invariance
            actual_width = module.in_features
            lr_scale = proxy_width / actual_width  # == 1 at proxy, <1 at larger models
            # muP scales ONLY the 2D matrix weight by 1/width. Vector params
            # (biases, ndim==1) stay width-invariant under muP+Adam, so exclude
            # them here; they fall through to the base-LR group below.
            matrix_params = [p for p in module.parameters() if p.ndim >= 2]
            param_groups.append({
                "params": matrix_params,
                "lr": base_lr * lr_scale,
                "weight_decay": weight_decay,
                "name": name,
            })

    # All other parameters (norms, embeddings) get base LR
    # Only the width-scaled matrix weights were grouped above; MuPLinear biases
    # (ndim==1) deliberately fall through to the width-invariant group below.
    named_param_set = {
        id(p)
        for m in model.modules()
        if isinstance(m, MuPLinear)
        for p in m.parameters()
        if p.ndim >= 2
    }
    other_params = [p for p in model.parameters() if id(p) not in named_param_set]
    if other_params:
        param_groups.append({
            "params": other_params,
            "lr": base_lr,
            "weight_decay": 0.0,
            "name": "other",
        })

    return torch.optim.AdamW(param_groups, lr=base_lr, betas=(0.9, 0.95))

Verifying muP: The Coordinate Check¶

The coordinate check is the single test that catches the large majority of real-world muP bugs. It plots the per-layer activation (or update) scale against width. Under a correct muP implementation, those curves are flat across widths — the whole point of the parameterization. Under standard parameterization (SP), the same curves blow up or shrink monotonically as width grows, because activation scale is exactly what SP fails to control.

import torch
import torch.nn as nn


def make_mlp(width: int, mup: bool) -> nn.Module:
    """4-layer MLP: input -> hidden -> hidden -> readout, ReLU between."""
    if mup:
        layers = [
            MuPLinear(64, width),
            nn.ReLU(),
            MuPLinear(width, width),
            nn.ReLU(),
            MuPLinear(width, width),
            nn.ReLU(),
            MuPLinear(width, 64, is_readout=True),
        ]
    else:
        # Standard parameterization contrast: plain nn.Linear, default init.
        layers = [
            nn.Linear(64, width),
            nn.ReLU(),
            nn.Linear(width, width),
            nn.ReLU(),
            nn.Linear(width, width),
            nn.ReLU(),
            nn.Linear(width, 64),
        ]
    return nn.Sequential(*layers)


def last_hidden_mean_abs_act(model: nn.Module, x: torch.Tensor) -> float:
    """Captures mean(abs(activation)) at the output of the last hidden ReLU."""
    activations = {}

    def hook(module, inp, out):
        activations["last_hidden"] = out.detach().abs().mean().item()

    # Index 5 is the ReLU right after the third (last hidden) linear layer.
    handle = model[5].register_forward_hook(hook)
    model(x)
    handle.remove()
    return activations["last_hidden"]


if __name__ == "__main__":
    print(f"{'width':>6} | {'muP mean|act|':>14} | {'SP mean|act|':>14}")
    for width in [256, 512, 1024, 2048]:
        for mup_flag, label in [(True, "mup"), (False, "sp")]:
            torch.manual_seed(0)
            model = make_mlp(width, mup=mup_flag)
            optimizer = (
                build_mup_optimizer(model, base_lr=1e-2, proxy_width=256)
                if mup_flag
                else torch.optim.AdamW(model.parameters(), lr=1e-2)
            )

            torch.manual_seed(0)
            x = torch.randn(32, 64)
            target = torch.randn(32, 64)

            for _ in range(5):
                optimizer.zero_grad()
                out = model(x)
                loss = nn.functional.mse_loss(out, target)
                loss.backward()
                optimizer.step()

            act = last_hidden_mean_abs_act(model, x)
            if mup_flag:
                mup_act = act
            else:
                sp_act = act
        print(f"{width:>6} | {mup_act:>14.4f} | {sp_act:>14.4f}")
    # Expected: the muP column stays roughly constant (within ~2x) across the
    # 8x width sweep 256 -> 2048; the SP column drifts several-fold over the
    # same sweep (in a typical run it shrinks ~15-20x, e.g. ~0.24 -> ~0.014),
    # i.e. it is NOT flat.

If your muP column is not flat, the bug is almost always in the init std, the per-layer LR multiplier, or the attention/readout scaling.

Practical muP Workflow¶

The Microsoft mup library (github.com/microsoft/mup) provides a plug-and-play implementation. Its central concept is that you never describe a model in absolute terms — you give it a base model and a delta model differing only in the widths you intend to scale, and it infers which dimensions are “infinite” and rewrites init and per-parameter LR accordingly:

# Sketch of the mup API (pip install mup). Requires MuReadout on the output head.
import torch.nn as nn
import mup

def MyTransformer(d_model):                      # stand-in for your real model
    return nn.Sequential(nn.Linear(64, d_model), nn.ReLU(),
                         mup.MuReadout(d_model, 64))

base_lr = 1e-2                                   # tuned at the base width
base = MyTransformer(d_model=256)                # base shapes
delta = MyTransformer(d_model=512)               # differs ONLY in scaled dims
model = MyTransformer(d_model=4096)              # the model you will train

mup.set_base_shapes(model, base, delta=delta)    # MUST come before init + optimizer
for p in model.parameters():
    if p.ndim >= 2:
        mup.init.normal_(p, mean=0.0, std=0.02)  # mup-aware init, reads base shapes

optimizer = mup.MuAdam(model.parameters(), lr=base_lr)   # or MuAdamW / MuSGD
# mup.coord_check.get_coord_data(...) produces the coordinate-check plot above.

Two ordering rules cause most mup bugs: set_base_shapes must run before you initialize weights and before you construct the optimizer, and the readout layer must be mup.MuReadout (or hand-scaled by \(1/d\)) or the width invariance silently fails.

The typical workflow is:

1. Define a base config with a small proxy width (e.g., 256 hidden dim).
2. Run a dense HP sweep over LR, initialization std, and possibly weight decay — this is cheap at small width.
3. Identify the optimal HP at proxy width.
4. Scale width to the full model. muP guarantees the same HP is optimal (up to reasonable approximation in practice).
5. Validate on a medium model (e.g., 1B) before launching the full run.

The evidence that this works is now substantial: Microsoft’s Phi models, various internal runs at other labs, and controlled ablations in the Tensor Programs V paper all show that muP-transferred LRs closely match the empirically optimal LRs found by grid search at the large scale — saving orders-of-magnitude in tuning compute.

Two notes on how this looks in 2026 practice. First, many teams no longer depend on the mup package: because the prescription reduces to a handful of rules (init std \(\propto 1/\sqrt{\text{fan\_in}}\), readout scaled by \(1/d\), per-matrix LR \(\propto 1/\text{fan\_in}\), attention logits divided by \(d_k\) rather than \(\sqrt{d_k}\)), it is often written directly into the model definition — which is also what makes it survive torch.compile and FSDP wrapping without surprises. Second, normalized-update optimizers get part of this for free. Muon’s orthogonalized update has a fixed per-element RMS by construction, independent of gradient scale and largely of width, so its peak LR transfers across width far better than Adam’s — which is why the capstone can carry a peak LR measured on a 43M proxy up to full width with only a short confirmation run instead of a full muP apparatus (Optimizer & Schedule). muP and normalized optimizers are attacking the same problem — making the update magnitude scale-invariant — from opposite ends.

Practical Hyperparameter Recipes¶

Consolidating the above into a reference table for common pretraining scales:

These are starting points synthesized from published literature (GPT-3, Llama ½/3, Mistral, Falcon, OLMo) — treat them as reasonable defaults, not ground truth. The right value for any specific run depends on architecture choices (RMSNorm vs LayerNorm, activation function, depth/width ratio), the optimizer, and the data mixture. The peak-LR column assumes AdamW on every parameter; a Muon/AdamW hybrid runs the 2D matrices an order of magnitude higher (the capstone’s ~100M model uses Muon at \(0.02\) with AdamW at \(3\text{e-}3\) for embeddings and norms), because Muon’s orthogonalized update has a fixed per-element magnitude rather than a gradient-scaled one. See Optimizers and, for the full worked recipe at this scale, Optimizer & Schedule.

from dataclasses import dataclass, field
from typing import Literal


@dataclass
class PretrainingHParams:
    """
    Reference hyperparameter config for LLM pretraining.
    Start here and tune with muP proxy sweeps.
    """
    # Optimizer
    optimizer: Literal["adamw", "lion", "adafactor"] = "adamw"
    peak_lr: float = 3e-4
    min_lr_fraction: float = 0.1       # eta_min = peak_lr * min_lr_fraction
    beta1: float = 0.9
    beta2: float = 0.95
    eps: float = 1e-8
    weight_decay: float = 0.1

    # Schedule
    schedule: Literal["cosine", "wsd", "linear", "rsqrt"] = "cosine"
    warmup_steps: int = 2000
    # For cosine: total_steps must be set before training
    # For WSD: set stable_steps and decay_steps instead
    total_steps: int = 100_000
    wsd_stable_fraction: float = 0.85  # fraction of (total-warmup) in stable phase

    # Gradient
    grad_clip_norm: float = 1.0
    grad_accumulation_steps: int = 1

    # Batch
    micro_batch_size: int = 4          # per-GPU, per-step
    tokens_per_sample: int = 2048

    def effective_batch_tokens(self, world_size: int) -> int:
        return (
            self.micro_batch_size
            * self.tokens_per_sample
            * self.grad_accumulation_steps
            * world_size
        )

    def total_tokens(self, world_size: int) -> int:
        return self.effective_batch_tokens(world_size) * self.total_steps


def scale_lr_for_batch_size(
    base_lr: float,
    base_batch: int,
    target_batch: int,
    rule: Literal["linear", "sqrt"] = "linear",
) -> float:
    """
    Scale learning rate when changing effective batch size.
    Use 'sqrt' for AdamW/Muon (the SDE-derived rule for adaptive and
    normalized-update optimizers) -- this is the LLM-pretraining default.
    Use 'linear' only for SGD-with-momentum, and only well below B*.
    """
    ratio = target_batch / base_batch
    if rule == "linear":
        return base_lr * ratio
    elif rule == "sqrt":
        return base_lr * math.sqrt(ratio)
    else:
        raise ValueError(f"Unknown rule: {rule}")


# Example: scale 1M-token/step config to 4M tokens/step under AdamW
base_cfg = PretrainingHParams(peak_lr=3e-4)
new_lr = scale_lr_for_batch_size(
    base_lr=base_cfg.peak_lr,
    base_batch=1_000_000,
    target_batch=4_000_000,
    rule="sqrt",
)
print(f"Scaled LR (sqrt): {new_lr:.2e}")  # 6.00e-04
# The SGD-style linear rule would give 1.20e-03 -- 2x higher, and near B* that
# is exactly the over-scaling that shows up as a loss spike a few thousand
# steps in rather than as an immediate divergence.

Warmup Duration Heuristics¶

A useful rule of thumb: warm up for at least \(T_w = \max(1000, 0.02 \times T_{\text{total}})\) steps — that is, at least 1000 steps or 2% of the total budget, whichever is larger. Shorter warmups are fine for fine-tuning (where the model is already initialized near a good basin) but dangerous for pretraining from scratch.

State it in tokens when you compare across runs, since that is the quantity that is actually invariant: 2% of a 38,147-step run at 524,288 tokens/step is ~1B warmup tokens, which is squarely the ~0.5–2B band used by open pretraining recipes at every scale from 100M to 70B (The Pretraining Run uses exactly 2,000 steps ≈ 1.05B tokens). Warmup is cheap insurance: over-warming costs you a fraction of a percent of final loss, while under-warming can cost you the run. When a run diverges early, lengthen warmup before you lower the peak LR — it preserves the peak you tuned.

For continued pretraining (e.g., domain adaptation starting from a released checkpoint), a short warmup of 100–500 steps is usually sufficient — the parameters are already in a well-behaved regime. See Continual & Domain-Adaptive Pretraining.

Putting It All Together: A Training Launch Checklist¶

Before launching a large run, verify each item:

Hyperparameter Launch Checklist
================================
[ ] Peak LR set (use muP proxy sweep if scaling architecture)
[ ] Warmup steps >= max(1000, 2% of total) for full pretraining
[ ] Batch size chosen: effective tokens/step in range [0.5M, 16M]
    depending on model scale; check against critical batch size estimate
[ ] LR scaled for batch size if changed from reference run (sqrt for AdamW/Muon)
[ ] Cosine or WSD schedule: total_steps (cosine) or stable_steps set
[ ] Min LR floor = 10% of peak LR
[ ] Weight decay = 0.1 (decoupled, excluded from embeddings/norms/biases)
[ ] Grad clip norm = 1.0; grad norm logged every step
[ ] beta2 = 0.95 (not 0.999 default)
[ ] Gradient accumulation: effective_batch = micro_batch * accum_steps * world_size
[ ] DDP: model.no_sync() used on non-final accumulation steps
[ ] Optimizer state saved in checkpoint; schedule reproducible from global_step
    alone (prefer a stateless lr_at_step over a pickled LambdaLR)
[ ] Per-group LR ratios (muP multipliers / Muon-vs-AdamW) preserved on resume
[ ] Monitoring: log {lr, grad_norm, loss, step} every step

Further Reading¶

• Yang, G. et al. “Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer.” NeurIPS 2022. (The muP paper; defines maximal-update parameterization and the theoretical framework behind HP transfer.)
• McCandlish, S. et al. “An Empirical Model of Large-Batch Training.” arXiv, 2018. (Introduces the gradient noise scale and critical batch size.)
• Goyal, P. et al. “Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour.” arXiv, 2017. (The linear scaling rule for batch size — derived for SGD with momentum.)
• Malladi, S. et al. “On the SDEs and Scaling Rules for Adaptive Gradient Algorithms.” NeurIPS 2022. (The square-root scaling rule for Adam/RMSProp; the correct default for LLM pretraining.)
• Hägele, A. et al. “Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations.” NeurIPS 2024. (Constant LR plus cooldown matches cosine; the empirical case for WSD.)
• Loshchilov, I. & Hutter, F. “Decoupled Weight Decay Regularization.” ICLR 2019. (Defines AdamW and decoupled weight decay.)
• Loshchilov, I. & Hutter, F. “SGDR: Stochastic Gradient Descent with Warm Restarts.” ICLR 2017. (Cosine annealing with warm restarts.)
• Hu, S. et al. “MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies.” arXiv, 2024. (Popularizes the WSD schedule and demonstrates its practical advantages.)
• Brown, T. et al. “Language Models are Few-Shot Learners.” NeurIPS 2020. (GPT-3; documents the hyperparameter choices used for large-scale pretraining at the time.)
• microsoft/mup GitHub repository. Reference implementation of maximal-update parameterization in PyTorch.

Exercises¶

1. Adam carries first- and second-moment estimates \(m_t\) and \(v_t\) that are both initialized to zero. Using this fact, explain (a) why pretraining a billion-parameter model without warmup frequently diverges in the first thousand steps, and (b) why the “Common Pitfall” admonition insists you re-warm the learning rate after a checkpoint restart that resets the optimizer state — even though the same admonition says a clean continuation (which keeps optimizer state) needs no re-warm.

2. Use the chapter’s cosine-with-warmup schedule with \(\eta_{\max} = 3\text{e-}4\), num_warmup_steps \(=100\), num_training_steps \(=1000\), and min_lr_fraction \(=0.1\). Compute the learning rate at (a) step 50, (b) step 550, and © step 1000. Give each to three significant figures.

3. Your reference run uses \(\eta_{\max} = 3\text{e-}4\) at an effective batch of \(0.5\text{M}\) tokens/step. You want to scale to \(4\text{M}\) tokens/step. (a) Give the peak LR the linear scaling rule prescribes and the peak LR the square-root rule prescribes. (b) The chapter says the critical batch size \(B^*\) for language cross-entropy is “on the order of a few million tokens.” Given that, which of your two candidate LRs is the safer choice for the \(4\text{M}\)-token batch, and why? © If you hold the total token budget fixed while going from \(0.5\text{M}\) to \(4\text{M}\) tokens/step, by what factor does the number of optimizer steps change?

4. The gradient-accumulation code divides the per-microbatch loss by accumulation_steps before calling .backward(). (a) Assume each of \(k\) microbatches contains exactly \(m\) tokens and its loss is the mean cross-entropy over those \(m\) tokens. Show that summing the \(k\) divided-and-backpropagated microbatch gradients equals the gradient of the mean loss over all \(km\) tokens. (b) Now suppose the microbatches have different token counts \(m_1, \dots, m_k\). Explain why dividing every microbatch loss by the same constant \(k\) no longer reproduces the true full-batch mean gradient, and state the correct weighting.

5. Implement a get_linear_schedule_with_warmup function in the same LambdaLR style as the chapter’s get_cosine_schedule_with_warmup: linear ramp over num_warmup_steps, then a linear decay to a floor of min_lr_fraction at num_training_steps. Add a smoke test that checks the midpoint of the decay phase and the final value.

import math
import torch
from torch.optim.lr_scheduler import LambdaLR


def get_linear_schedule_with_warmup(
    optimizer: torch.optim.Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    min_lr_fraction: float = 0.1,
) -> LambdaLR:
    """Linear warmup, then linear decay from peak to min_lr_fraction * peak."""
    def lr_lambda(current_step: int) -> float:
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        progress = min(progress, 1.0)  # clamp so we never go below the floor
        # Linear interpolation from 1.0 down to min_lr_fraction.
        return 1.0 - (1.0 - min_lr_fraction) * progress

    return LambdaLR(optimizer, lr_lambda)


if __name__ == "__main__":
    model = torch.nn.Linear(10, 10)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=100, num_training_steps=1000,
        min_lr_fraction=0.1,
    )

    lrs = []
    for _ in range(1000):
        optimizer.step()
        lrs.append(optimizer.param_groups[0]["lr"])
        scheduler.step()

    # Peak is 3e-4. Decay midpoint is step 550 (progress = 0.5):
    # multiplier = 1 - 0.9 * 0.5 = 0.55  ->  1.65e-4.
    assert abs(lrs[550] - 3e-4 * 0.55) < 1e-9, "decay midpoint wrong"
    # Step 999 (progress = 899/900 ~ 0.999) is just above the 3e-5 floor.
    assert lrs[-1] < 3e-4 * 0.101, "floor not reached"
    print(f"Peak LR: {max(lrs):.2e}, Step 550 LR: {lrs[550]:.2e}, "
          f"Final LR: {lrs[-1]:.2e}")
    # Output: Peak LR: 3.00e-04, Step 550 LR: 1.65e-04, Final LR: 3.03e-05

6. You tune hyperparameters on a muP proxy model of width \(d_{\text{proxy}} = 256\) and find an optimal base LR of \(1\text{e-}2\). (a) Using the chapter’s build_mup_optimizer convention (lr_scale = proxy_width / actual_width), what effective LR does a hidden MuPLinear layer of width \(2048\) receive when you scale up? (b) In standard parameterization the optimal LR for hidden matrices scales roughly as \(1/d\). If you instead grid-searched at width 256 in SP and naively reused that LR at width \(4096\), by what factor would you likely be off? © When you run the chapter’s coordinate check, what qualitative signature in the “muP mean|act|” column versus the “SP mean|act|” column tells you the muP implementation is correct?