14.6 Optimizer & Schedule: Muon + MuonClip and Warmup-Stable-Decay¶

Every choice we have made so far — the deep-and-thin Stack-100M architecture of Chapter 14.4, the ~20B-token over-training budget, the FineWeb-Edu/Cosmopedia data mix — only pays off if the optimizer can actually walk the loss down cleanly across 38,147 steps without spiking, and if the learning-rate schedule leaves a hook for the mid-training annealing phase to come. This chapter fixes both, and it fixes them numerically: every constant here is the one that appears in Chapter 14.7’s TrainConfig and drives the shipped capstone/stacklm package.

We adopt the optimizer stack the capstone spec pins in §5: Muon (Jordan et al., 2024) for the 2D hidden weight matrices, AdamW for embeddings, norms, and every 1D parameter, a QK-clip in the spirit of MuonClip (Kimi K2, Moonshot 2025) for attention-logit stability, and the Warmup-Stable-Decay (WSD) schedule (MiniCPM, Hu et al., 2024; used by DeepSeek).

This is the applied, capstone-specific companion to two deeper reference chapters you should have open in another tab: Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo derives Adam and Muon from first principles, and Learning Rate Schedules, Warmup, Batch Size & Hyperparameters covers warmup, cosine, batch/LR coupling, and muP. We will not re-derive those; we will use them, and explain the three or four things that are specific to training a 100M model on a single A100 to 200 tokens per parameter. When something threatens to blow the run up, we lean on Training Stability, Loss Spikes & Debugging Large Runs.

build_optimizers(model, muon_lr, adamw_lr, weight_decay, betas) -> (muon, adamw)
wsd_lr(step, *, peak_lr, warmup_steps, total_steps,
       decay_steps=None, decay_frac=0.2, final_frac=0.0) -> float
qk_clip_(model, max_logits, tau) -> int   # number of layers that fired

Why Muon, and Why a Hybrid¶

Start with the memory arithmetic, because at 100M parameters it is not the binding constraint — and that itself is the lesson. AdamW keeps two full-precision state tensors per parameter, the first moment \(m\) and second moment \(v\). For Stack-100M’s ~101M parameters that is 101M × 2 × 4 bytes ≈ 0.8 GB of optimizer state, trivial on an 80GB A100. So unlike a 70B run, we are not choosing Muon to save memory. We are choosing it because, at fixed compute, it walks the loss down measurably faster on the 2D weight matrices that dominate a transformer — the attention projections and the SwiGLU MLP. Muon has driven a series of speedrun records on nanoGPT-scale models (KellerJordan/modded-nanogpt), is the optimizer in Karpathy’s nanochat (2025), and was used at frontier scale by Kimi K2 (Moonshot, 2025). For a project whose whole thesis is “spend training compute once, serve forever,” a faster optimizer is free perplexity.

The core idea, developed carefully in Chapter 3.9, is orthogonalization of the momentum update. Adam rescales each coordinate independently by its running gradient magnitude; it never looks at the matrix structure of a weight tensor. Muon does. It takes the momentum buffer \(B_t\) for a weight matrix \(W \in \mathbb{R}^{m\times n}\) and replaces it with the nearest semi-orthogonal matrix — the \(UV^\top\) from its singular value decomposition \(B_t = U\Sigma V^\top\), which is exactly the orthogonal factor in the polar decomposition. Intuitively, the raw momentum update is often dominated by a few large singular directions; the network learns fast along those and starves the rest. Orthogonalizing sets every singular value to 1, so the update pushes equally hard along all directions of the matrix. This is a spectral condition-number fix, and it is why Muon behaves like a cheap approximation to a full matrix preconditioner (Shampoo, Gupta et al., 2018) without the cost of forming and inverting the preconditioner.

Computing an SVD every step of every layer would be far too slow. Muon’s key engineering move is to approximate the orthogonal factor with a fixed number (typically 5) of Newton–Schulz iterations — a matrix polynomial recurrence that converges to \((B B^\top)^{-1/2} B \approx U V^\top\) using only matrix multiplies, which are exactly what GPUs are fastest at. We implement it below.

Why not Muon on everything?¶

Muon’s orthogonalization only makes sense for a matrix whose two dimensions are both “feature” axes that mix. Three parameter groups in Stack-100M fail that test and stay on AdamW:

• The token embedding (32768 × 512, tied to the output head per Press & Wolf 2017). Its rows are per-token and extremely sparse in the gradient — most tokens in a batch touch a tiny fraction of rows. Orthogonalizing across the 32,768-token axis is meaningless; you want per-row adaptive rates, which is exactly Adam.
• RMSNorm gains and any 1D vector. A 1D tensor has no matrix structure to orthogonalize. (This group includes the QK-norm gains — which, as we will see in the stability section, are also the parameters the QK-clip has to reach for.)
• Biases — Stack-100M has essentially none (SwiGLU and RMSNorm are bias-free), but if present they go to AdamW.

This split — Muon for 2D hidden matrices, AdamW for embeddings/norms/1D — is not a compromise; it is the standard hybrid used at scale, including Kimi K2 and Moonlight. The capstone commits to it.

Stack-100M parameter routing
┌─────────────────────────────────────────────┬───────────┬─────────┐
│ parameter group                              │ optimizer │  count  │
├─────────────────────────────────────────────┼───────────┼─────────┤
│ tok_emb.weight  (32768×512, tied)            │  AdamW    │    1    │
│ blocks.*.attn.wq / wo   (512×512)            │  Muon     │   60    │
│ blocks.*.attn.wk / wv   (128×512)            │  Muon     │   60    │
│ blocks.*.mlp.w_gate / w_up (1408×512)        │  Muon     │   60    │
│ blocks.*.mlp.w_down     (512×1408)           │  Muon     │   30    │
│ blocks.*.norm*.weight   (1D RMSNorm gains)   │  AdamW    │   60    │
│ blocks.*.attn.q_norm / k_norm.weight (1D)    │  AdamW    │   60    │
│ final_norm.weight       (1D)                 │  AdamW    │    1    │
└─────────────────────────────────────────────┴───────────┴─────────┘
                                       Muon total: 210 matrices

That is 7 two-dimensional matrices per block (four attention projections + three SwiGLU matrices) × 30 blocks = 210 matrices routed to Muon; everything else goes to AdamW. Keep that 210 in mind — it reappears twice, once when we estimate Muon’s FLOP overhead and once when we fix its wall-clock overhead, which is a different and more interesting problem.

Muon: The Newton–Schulz Orthogonalizer¶

Let us build the numerical core first, because everything else hangs off it. Given the momentum buffer \(B \in \mathbb{R}^{m\times n}\), we want an orthogonal factor \(O = U V^\top\) where \(B = U\Sigma V^\top\). The quintic Newton–Schulz iteration Keller Jordan uses starts from a spectrally-normalized \(X_0 = B / \lVert B \rVert_F\) (so all singular values land in \([0,1]\)) and applies

with the tuned coefficients \((a,b,c) = (3.4445,\,-4.7750,\,2.0315)\). This is a degree-5 odd polynomial \(p(\sigma)\) applied to each singular value \(\sigma\); the coefficients are chosen so that \(p\) pushes every \(\sigma \in (0,1]\) rapidly toward 1 while staying stable near 0. Five iterations get every singular value close enough to 1 for optimization purposes — we do not need machine-precision orthogonality, just a well-conditioned direction.

# stacklm/optim/muon.py
import torch

_NS_COEFFS = (3.4445, -4.7750, 2.0315)          # tuned quintic coefficients

@torch.no_grad()
def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Approximate the orthogonal (polar) factor UV^T of a 2D matrix G via a
    quintic Newton-Schulz iteration. Returns a matrix with singular values ~1.

    This is the numerical heart of Muon (Jordan et al., 2024). It uses only
    matmuls, so it runs fast on tensor cores. We deliberately do NOT aim for
    exact orthogonality -- 5 steps is enough to give a well-conditioned update
    direction. The reference implementation runs this in bf16; the capstone repo
    runs it in fp32, which keeps the CPU smoke test agreeing with the GPU path to
    fp32 tolerance (not bit-for-bit -- CPU and CUDA GEMMs reduce in different
    orders -- but far closer than bf16 would be).
    """
    assert G.ndim == 2, "Newton-Schulz orthogonalization is only for 2D matrices"
    a, b, c = _NS_COEFFS
    X = G.float()
    X = X / (X.norm() + 1e-7)                    # normalize so all sigma <= 1
    transposed = G.size(0) > G.size(1)          # keep the short side as rows
    if transposed:
        X = X.T                                  # iterate on the smaller Gram matrix
    for _ in range(steps):
        A = X @ X.T                              # A = X X^T  (small: min(m,n)^2)
        P = b * A + c * (A @ A)                  # polynomial in the Gram matrix
        X = a * X + P @ X                        # quintic Newton-Schulz update
    if transposed:
        X = X.T
    return X

Note the transposed trick: the Gram matrix \(XX^\top\) has size \(\min(m,n)^2\), so we always iterate on the smaller of the two dimensions. For a SwiGLU weight of shape 1408 × 512 that means the inner products live in \(512^2\), not \(1408^2\) — a real speedup. (We name the polynomial term P, not b, to avoid shadowing the coefficient; a subtle way to break this iteration is to reuse a coefficient name inside the loop.)

One property of that first line matters far more than it looks, and we will cash it in twice: because \(X_0 = B/\lVert B\rVert_F\), the output is invariant to any positive rescaling of the input. \(O(cB) = O(B)\) for every \(c > 0\). Muon does not care how big your gradient is, only which way it points.

The full Muon step and RMS matching¶

Two scale factors turn the orthogonal direction into a usable update.

First, the orthogonalized update has a root-mean-square element magnitude of roughly \(1/\sqrt{\max(m,n)}\): a semi-orthogonal \(m\times n\) matrix has \(\min(m,n)\) singular values equal to 1, hence Frobenius norm \(\sqrt{\min(m,n)}\), spread over \(mn\) entries, so RMS \(=\sqrt{\min(m,n)/(mn)} = 1/\sqrt{\max(m,n)}\). That is shape-dependent: the same learning rate would move a 1408×512 matrix and a 128×512 matrix by different relative amounts.

Second, to put Muon’s step magnitude in a fixed, shape-free regime, the Moonshot “Muon is Scalable” report (Liu et al., 2025) multiplies the Muon update by \(0.2\cdot\sqrt{\max(m,n)}\), which cancels the shape dependence and brings the RMS to exactly \(0.2\). Where does \(0.2\) come from? Not from the idealized AdamW bound. AdamW’s update is \(m/(\sqrt{v}+\epsilon)\), whose per-element magnitude is \(\approx 1\) only if \(m\) and \(v\) are estimated over the same window on a stationary gradient. In real LLM training the two EMAs disagree (\(\beta_1 = 0.9\) vs \(\beta_2 = 0.95\)) and the gradient is noisy, so the measured update RMS across a transformer’s weight matrices sits well below 1 — on the order of \(0.2\)–\(0.4\). Moonshot picked \(0.2\) to land inside that measured band. That is the honest statement; the “AdamW step has RMS 1” folklore is the theoretical ceiling, not the observed value.

What RMS matching does and does not buy. It makes Muon’s learning rate shape-invariant: one number governs all 210 matrices, from 128×512 to 512×1408. That is the real win, and it is why you can sweep Muon’s LR as a single scalar. It does not imply that the two optimizer groups want the same number, because the AdamW group in this hybrid is not a set of hidden matrices at all — it is the tied embedding plus every 1D gain, objects with entirely different gradient statistics. The capstone freezes

a ratio of about 6.7:1. Treat that ratio as empirical, not derived — it is the band the reference recipes live in (Moonlight and nanochat both run Muon in the low \(10^{-2}\)s with a much smaller embedding LR), and it is the pair Ch. 14.7’s TrainConfig ships. The reason the embedding wants the smaller number is concrete: its gradient is row-sparse, so a rare token’s row receives a full-magnitude Adam update from the handful of batches that contain it, and a smaller step is cheap insurance against those rows thrashing. The practical point is that you tune a line — one LR and a fixed ratio — not a plane.

# stacklm/optim/muon.py  (continued)
from torch.optim.optimizer import Optimizer

class Muon(Optimizer):
    """Muon: momentum + Newton-Schulz orthogonalization, for 2D matrices ONLY.
    Route embeddings/norms/1D params to AdamW instead (see build_optimizers()).

    Args mirror the Jordan et al. reference plus the Moonshot RMS-matching scale
    and decoupled weight decay (Liu et al., 2025). `batched=True` switches to the
    shape-bucketed fast path described below."""

    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True,
                 weight_decay=0.1, ns_steps=5, batched=False):
        defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov,
                        weight_decay=weight_decay, ns_steps=ns_steps)
        super().__init__(params, defaults)
        self.batched = batched

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure is not None else None
        for group in self.param_groups:
            lr, mu = group["lr"], group["momentum"]
            nesterov, wd, ns = group["nesterov"], group["weight_decay"], group["ns_steps"]

            live, dirs = [], []
            for p in group["params"]:
                if p.grad is None:
                    continue
                g = p.grad
                assert g.ndim == 2, "Muon received a non-2D param; check routing"
                state = self.state[p]
                if "momentum_buffer" not in state:
                    state["momentum_buffer"] = torch.zeros_like(g)
                buf = state["momentum_buffer"]
                buf.mul_(mu).add_(g)                    # B <- mu*B + g
                live.append(p)
                # Nesterov look-ahead: orthogonalize g + mu*B, not the bare buffer.
                dirs.append(g.add(buf, alpha=mu) if nesterov else buf)

            if not live:
                continue
            if self.batched:
                orths = _orthogonalize_bucketed(dirs, ns)     # one bmm per shape class
            else:
                orths = [zeropower_via_newtonschulz5(d, steps=ns) for d in dirs]

            for p, o in zip(live, orths):
                # RMS-match: cancel the shape dependence, land the update RMS at 0.2.
                scale = 0.2 * max(p.size(0), p.size(1)) ** 0.5
                # Decoupled weight decay (AdamW-style): shrink the weight itself.
                if wd != 0.0:
                    p.mul_(1.0 - lr * wd)
                p.add_(o.to(p.dtype), alpha=-lr * scale)   # W <- W - lr*scale*O
        return loss

Two design points worth pausing on. We orthogonalize the Nesterov look-ahead direction \(g + \mu B\) rather than the bare buffer — this is the reference Muon default and gives a slightly more responsive update. And weight decay is decoupled exactly as in AdamW (Chapter 3.9): we shrink the weight by (1 - lr*wd) before applying the orthogonalized step, so the decay is a clean L2 pull toward zero and not entangled with the (unit-scale) orthogonal update. The capstone uses weight_decay = 0.1 for both optimizers.

FLOPs are cheap; kernel launches are not¶

The FLOP overhead of Muon is genuinely negligible, and the worked example below computes it exactly: ~1.04 TFLOP per optimizer step against the model’s ~319 TFLOP, i.e. 0.33%. But a FLOP count is not a wall-clock claim, and this is the single most common way Muon disappoints in practice.

Count the launches instead. The naive loop visits 210 matrices; each runs 5 iterations; each iteration issues roughly four GPU kernels (X @ X.T, A @ A, the fused b*A + c*(A@A) elementwise, P @ X), plus the momentum and weight-decay ops. That is on the order of 4,200 kernel launches per optimizer step, on matrices whose inner dimension is 128–1408. At those sizes an A100 finishes each matmul in single-digit microseconds while a kernel launch costs several microseconds — this regime is launch-latency bound, not FLOP bound. Measured naively, Muon’s overhead at 100M scale is typically several percent of step time, not 0.3%.

Two fixes, both standard, both worth doing:

1. Compile the iteration. Wrapping zeropower_via_newtonschulz5 in torch.compile fuses the elementwise work and can capture the whole 5-step loop into a CUDA graph, collapsing most of the launch overhead (Kernel Fusion, torch.compile, CUDA Graphs & Compilers).
2. Batch same-shaped matrices. After the orientation-normalizing transpose, Stack-100M’s 210 matrices collapse into just three shape classes: 60 of 512×512 (wq, wo), 60 of 128×512 (wk, wv), and 90 of 512×1408 (w_gate, w_up, and w_down transposed). Stack each class and run the iteration as a single bmm per class: 3 classes × 5 iterations × 4 kernels ≈ 60 launches, a ~70× reduction.

The second fix is the one that matters, and it is shipped — Muon(..., batched=True), or build_optimizers(..., batched_muon=True):

# stacklm/optim/muon.py  (the batched fast path)
@torch.no_grad()
def zeropower_batched(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
    """Same quintic on a STACK of identically-shaped matrices: G is (K, m, n).

    Collapses K*steps*4 kernel launches into steps*4 batched ones. Algebraically
    identical to looping zeropower_via_newtonschulz5; expect last-ulp differences,
    because bmm and mm use different reduction orders and split-k heuristics.
    """
    assert G.ndim == 3, "zeropower_batched expects a (K, m, n) stack"
    a, b, c = _NS_COEFFS
    X = G.float()
    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)   # per-matrix Frobenius
    transposed = X.size(-2) > X.size(-1)
    if transposed:
        X = X.mT                                          # batched transpose
    for _ in range(steps):
        A = X @ X.mT                                      # (K, r, r) via bmm
        P = b * A + c * (A @ A)
        X = a * X + P @ X
    if transposed:
        X = X.mT
    return X


@torch.no_grad()
def _orthogonalize_bucketed(dirs, steps: int):
    """Orthogonalize a list of 2D directions, bucketing identical shapes.

    Key on the SORTED shape, so (1408, 512) and (512, 1408) land in the SAME
    bucket after the orientation-normalizing transpose. That is what collapses
    Stack-100M's 210 matrices into 3 classes rather than 4.
    """
    out = [None] * len(dirs)
    buckets: dict[tuple, list[int]] = {}
    for i, d in enumerate(dirs):
        m, n = d.shape
        buckets.setdefault((min(m, n), max(m, n)), []).append(i)
    for idxs in buckets.values():
        flip = [dirs[i].size(0) > dirs[i].size(1) for i in idxs]
        stack = torch.stack([dirs[i].T if f else dirs[i] for i, f in zip(idxs, flip)])
        O = zeropower_batched(stack, steps=steps)
        for j, (i, f) in enumerate(zip(idxs, flip)):
            out[i] = O[j].T if f else O[j]
    return out

The bucket key is the interesting line. Keying on the raw shape gives you four classes (1408×512 and 512×1408 are different tensors); keying on the sorted shape and transposing the tall ones gives three, because the iteration is orientation-agnostic. On Stack-100M’s shapes the batched path agrees with the looped one to \(\sim4\times10^{-7}\) absolute — pure reduction-order noise, invisible next to the \(\pm20\%\) orthogonality slack we deliberately accept.

torch.cuda.synchronize(); t0 = time.perf_counter()
muon.step(); adamw.step()
torch.cuda.synchronize(); print("opt step ms:", 1e3 * (time.perf_counter() - t0))

Setting the Peak Learning Rate (A Protocol, Not a Guess)¶

The peak LR is the most consequential number in this chapter, and “0.02-ish, like the reference recipes” is not an answer a reproducible book is allowed to give. Chapter 14.5 already built a ladder of tiny models ({4M, 9M, 19M, 43M} params) under a shared recipe precisely so we can measure hyperparameters instead of guessing them. Here is the protocol, and it is cheap: the whole sweep costs well under an hour on the same A100.

Step 1 — sweep on the top ladder rung. Take the 43M rung. Train it for 500–1000 optimizer steps at a reduced effective batch of one micro-batch, 32 × 2048 = 65,536 tokens/step, with the same WSD warmup shape (scale warmup proportionally, e.g. 250 steps). Sweep the Muon peak over a ×2 grid, {2e-3, 4e-3, 8e-3, 1.6e-2, 3.2e-2}, holding adamw_lr at the fixed ratio throughout so you are searching a line, not a plane.

Step 2 — read the curve, then back off one grid point. Plot final held-out loss vs peak LR. You get a U: too low and the model is under-trained at this step count; too high (3.2e-2 here) and the loss either spikes or plateaus above the minimum. Take the argmin — 1.6e-2 — then step one grid point down, to 8e-3. The full run is ~40× longer than the probe, and the highest LR that survives 1,000 steps is not always the highest LR that survives 38,147.

Step 3 — transfer across batch size. The probe ran at 65,536 tokens/step; the real run is 524,288, an 8× increase. Below the critical batch size, the standard rule of thumb for adaptive/normalized updates is \(\eta \propto \sqrt{B}\) (Chapter 3.10), so multiply by \(\sqrt{8} \approx 2.83\): \(8\text{e-}3 \times 2.83 = 2.26\text{e-}2\). Round down to 0.02 — which is exactly TrainConfig.muon_peak_lr, and squarely inside the band Moonlight and nanochat report.

Step 4 — transfer across width. The 43M rung is narrower than d_model = 512. Two defensible options: (a) muP width scaling, which says matrix-like parameters want \(\eta \propto 1/\text{fan\_in}\) (Chapter 3.10); or (b) rely on the fact that Muon’s RMS-matched update has a fixed per-element magnitude by construction, which makes its LR far closer to width-invariant than Adam’s — this is one of Muon’s practical selling points. We take (b) and then verify with a 200-step confirmation run at full width and full batch: if the loss curve at step 200 is not monotone-decreasing and the grad-norm log is not flat, drop the LR by 2× and repeat.

Attention-Logit Stability: QK-Norm, QK-Clip, and Soft-Caps¶

Muon’s aggressiveness has a failure mode that AdamW largely hides: because it pushes equally along every singular direction, the query and key projections can grow until attention logits explode. The pre-softmax score is \(s_{ij} = q_i^\top k_j / \sqrt{d_h}\); if \(\lVert W_Q\rVert\) and \(\lVert W_K\rVert\) drift upward together, \(\max_{ij} s_{ij}\) climbs into the hundreds, the softmax saturates to a one-hot, gradients through it vanish, and — in bf16 — you get a loss spike or an outright NaN (Chapter 3.11). Kimi K2 hit exactly this at trillion-parameter scale and introduced MuonClip, whose active ingredient is QK-clip: a per-head, post-step rescaling that caps the maximum attention logit.

Now the part most treatments get wrong. Kimi K2 needed QK-clip on \(W_Q\)/\(W_K\) precisely because it does not QK-norm. Stack-100M does (Chapter 14.4 sets qk_norm = True), and under QK-norm, rescaling \(W_Q\) or \(W_K\) changes nothing at all. The two mechanisms are alternatives operating on the same knob, not layers that stack.

The invariance, and the bound QK-norm actually buys¶

Stack-100M’s attention computes, per head,

with \(\gamma_q,\gamma_k \in \mathbb{R}^{d_h}\) the learned RMSNorm gains. RMSNorm is positively homogeneous of degree 0: for any \(\eta > 0\), \(\operatorname{RMSNorm}(\eta z) = \operatorname{RMSNorm}(z)\) up to the \(\epsilon\) in the denominator. Therefore

Scaling the projections is a no-op. RoPE does not rescue it either: RoPE is a rotation applied after the norm, and rotations preserve norms.

What does set the logit scale is the pair of gains. Since \(u = z/\operatorname{rms}(z)\) has unit RMS, \(\lVert u \rVert_2 = \sqrt{d_h}\), so \(\lVert q \rVert_2 \le \lVert \gamma_q \rVert_\infty \sqrt{d_h}\), and by Cauchy–Schwarz

For Stack-100M, \(d_h = 64\) and both gains initialize to 1, so the max attention logit is bounded by 8 at initialization, with no dependence on \(W_Q, W_K\) whatsoever. Blow-up is only possible if the gains grow. That is a much narrower failure surface — and it is where the clip must act.

You can verify all of this in about ten seconds against the capstone package:

# verify_qk_invariance.py -- run from capstone/
import torch
from stacklm.config import toy_config
from stacklm.model.transformer import Stack100M

torch.manual_seed(0)
cfg = toy_config()                       # head_dim=16 -> a-priori bound sqrt(16)=4
model = Stack100M(cfg)
x = torch.randint(0, cfg.vocab_size, (2, 32))

def max_logit(m):
    rec = {}
    for b in m.blocks:                   # exact per-head max, not the cheap bound
        b.attn.record_exact = True
    with torch.no_grad():
        m(x, record=rec)
    return max(float(v.max()) for v in rec.values())

a = model.blocks[0].attn
gq, gk = a.q_norm.weight.detach(), a.k_norm.weight.detach()
bound = cfg.head_dim ** 0.5 * gq.abs().max() * gk.abs().max()
print(f"a-priori bound        {float(bound):.4f}")
print(f"observed max logit    {max_logit(model):.4f}")

with torch.no_grad():                    # (1) rescale the PROJECTIONS
    for b in model.blocks:
        b.attn.wq.weight.mul_(0.3); b.attn.wk.weight.mul_(0.3)
print(f"after W_Q,W_K x0.3    {max_logit(model):.4f}   <- unchanged (no-op)")

with torch.no_grad():                    # (2) rescale the GAINS
    for b in model.blocks:
        b.attn.q_norm.weight.mul_(0.5); b.attn.k_norm.weight.mul_(0.5)
print(f"after gains x0.5 each {max_logit(model):.4f}   <- exactly 0.25x")

a-priori bound        4.0000
observed max logit    3.3529
after W_Q,W_K x0.3    3.3388   <- unchanged (no-op)
after gains x0.5 each 0.8347   <- exactly 0.25x

The projection rescale moved the max logit by 0.4% — that residue is entirely the RMSNorm \(\epsilon\), not a real effect. The gain rescale multiplied it by \(0.5\times0.5 = 0.25\) exactly, as the algebra promises. And the a-priori bound of 4 (here \(d_h = 16\)) is tight to within 20% of the observed maximum.

Harvesting \(S_{\max}\) without giving up FlashAttention¶

Before you can clip, you need the per-head maximum logit — and this is not free, contrary to what a casual reading of the Kimi recipe suggests. Stack-100M’s fast path is F.scaled_dot_product_attention, which dispatches to a FlashAttention-style kernel (FlashAttention I: IO-Awareness & The Online Softmax). That kernel’s entire point is that it never materializes the \((B, H, T, T)\) score matrix. There is no .amax() to take, because there is nothing to take it of.

If you want the exact maximum you must fall back to eager attention and build the score tensor. At the capstone’s micro_batch_size = 32, n_heads = 8, T = 2048, in fp32:

2.15 GB in bf16, plus \(O(T^2)\) FLOPs and the loss of FlashAttention’s memory win. Doing that every step, on every micro-batch, of all 30 layers is not a rounding error; it is a different training run.

The capstone therefore offers three tiers, and defaults to the cheap one:

Tiers 0 and 1 are Cauchy–Schwarz upper bounds, so they fire the clip slightly early — a conservative error, which is what you want in a safety net. Tier 1 is tight in practice because attention keys and queries are not adversarially aligned. Here is the real plumbing in stacklm/model/attention.py:

# stacklm/model/attention.py  (inside Attention.forward, after QK-norm + RoPE,
#  after k/v have been repeat_interleave'd up to n_heads)
scale = 1.0 / (self.head_dim ** 0.5)

if record is not None:
    with torch.no_grad():
        if self.record_exact:
            # Tier 2: exact, but materializes (B, n_heads, T, T) and gives up SDPA.
            att = (q.float() @ k.float().transpose(-2, -1)) * scale
            causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))
            att = att.masked_fill(~causal, float("-inf"))
            record[self.layer_idx] = att.amax(dim=(0, 2, 3))     # (n_heads,)
        else:
            # Tier 1 (default): Cauchy-Schwarz bound. O(B*H*T*d_h) work,
            # no T^2 tensor, FlashAttention/SDPA fast path preserved below.
            qn = q.float().norm(dim=-1).amax(dim=(0, 2))         # (n_heads,)
            kn = k.float().norm(dim=-1).amax(dim=(0, 2))         # (n_heads,)
            record[self.layer_idx] = qn * kn * scale

The second half of the answer is frequency, and the tier you pick decides it:

• Tier 1 is cheap enough to run every step. This is what Ch. 14.7’s train.py does: it keeps a preallocated (n_layers, n_heads) running-max buffer, updates it during the normal accumulation loop, calls qk_clip_ after the optimizer steps, and zeroes it. No probe forward, no extra pass.
• Tier 2 must be gated. Attention-logit drift is slow — the gains move by a decayed learning rate per step — so there is no reason to measure exactly 38,147 times. stacklm/train/loop.py::pretrain() exposes qk_clip_every for this; set it to 200 (~190 measurements over the run) and take the reading on a single small probe batch with post-step weights. See the canonical optimizer_step below.

Either way, log the return value of qk_clip_: it is the single most informative stability signal this loop produces.

The clip that actually works here¶

Since the gains are the knob, the clip scales the gains. Scaling \(\gamma_q\) and \(\gamma_k\) each by \(\sqrt{\tau/S_{\max}}\) multiplies every logit in the layer by exactly \(\tau/S_{\max}\) — this is exactly the bilinearity argument Kimi uses, just applied to the parameters that are actually free.

One structural consequence you must not paper over: Stack-100M allocates one gain vector of length head_dim per layer, shared by all heads (RMSNorm(cfg.head_dim) in Ch. 14.4’s Attention.__init__). So the clip’s granularity is the layer, driven by the layer’s worst head — not the head. That is the price of the cheaper parameterization, and it is fine: over-shrinking the well-behaved heads in a layer by the same factor is conservative and reversible (they will regrow their gains). If you want Kimi’s per-head granularity, widen the gains to (n_heads, 1, head_dim) and (n_kv_heads, 1, head_dim), which costs 8*64 + 2*64 = 640 parameters per layer instead of 128 — and then the GQA group logic below transfers verbatim (Exercise 7 works this through).

# stacklm/optim/qk_clip.py
import torch

@torch.no_grad()
def qk_clip_(model, max_logits, tau: float = 100.0) -> int:
    """QK-clip (the core of MuonClip, Kimi K2 / Moonshot 2025), architecture-aware.

    `max_logits[layer_idx]` is an (n_heads,) tensor recorded by the attention
    module this forward pass. Returns the number of layers that fired -- log it.

    A dense (n_layers, n_heads) tensor is also accepted, because Ch. 14.7's
    train.py harvests into a preallocated running-max buffer. It is CONVERTED,
    not indexed directly: `layer_idx in tensor` is elementwise VALUE membership
    in PyTorch, not key lookup, and would silently do the wrong thing.

    The clip acts on whatever parameter actually controls the logit scale:
      * QK-norm ON  -> the RMSNorm gains (rescaling W_Q/W_K is provably a no-op).
      * QK-norm OFF -> W_Q / W_K, per Kimi K2, GQA-aware.
    """
    if torch.is_tensor(max_logits):
        max_logits = {i: max_logits[i] for i in range(max_logits.size(0))}
    fired = 0
    for layer_idx, block in enumerate(model.blocks):
        if layer_idx not in max_logits:
            continue
        s_max = max_logits[layer_idx].float()          # (n_heads,)
        attn = block.attn
        if isinstance(attn.q_norm, torch.nn.Identity):
            fired += _clip_projections_(attn, s_max, tau)    # Kimi K2 config
        else:
            fired += _clip_qk_norm_gains_(attn, s_max, tau)  # Stack-100M
    return fired


@torch.no_grad()
def _clip_qk_norm_gains_(attn, s_max, tau: float) -> int:
    """QK-norm path: the gains ARE the temperature, so clip them.

    Scaling gamma_q and gamma_k each by sqrt(tau/S) multiplies every logit in the
    layer by exactly tau/S (RoPE is a rotation and commutes with the scale).
    Stack-100M shares one gain vector across heads, so this is a per-LAYER clip
    driven by the layer's worst head.
    """
    s = float(s_max.max())
    if s <= tau:
        return 0                                  # this layer is fine
    eta = (tau / s) ** 0.5                        # sqrt so q AND k share it
    attn.q_norm.weight.mul_(eta)
    attn.k_norm.weight.mul_(eta)
    return 1

Choosing \(\tau\) — and why the shipped default is 100. Ch. 14.7’s TrainConfig ships qk_clip_tau = 100.0, Kimi’s number, and that is a deliberate, defensible choice with a caveat you must understand. Under QK-norm the logits start bounded by 8, so reaching 100 requires \(\lVert\gamma_q\rVert_\infty\lVert\gamma_k\rVert_\infty > 12.5\) — a state so pathological the run is already lost. In other words, at \(\tau = 100\) with qk_norm = True the clip is an inert backstop: it costs nothing, it will essentially never fire, and it is exactly the right value the moment you flip qk_norm = False (the Ch. 14.4 ablation), where there is no a-priori bound at all. That is why it is the config default: it is correct across both configurations.

If you want the clip to do the other job — be an early-warning sensor rather than a last rite — pin \(\tau\) to your architecture’s natural scale and set qk_clip_tau = 30, roughly \(4\times\) the initialization bound of 8. That is comfortably above normal gain drift, far below the softmax-saturation regime, and low enough that a rising trigger count tells you the LR is too high days before the loss does. The capstone’s recommendation is: ship 100, and switch to 30 the moment you are debugging. Exercises 4 and 7 work both settings.

If you drop QK-norm: Kimi K2’s per-head, GQA-aware clip¶

Run Stack-100M with qk_norm = False and \(W_Q, W_K\) become the knob again. Now the original MuonClip applies — and here GQA introduces a wrinkle worth its own code path.

In Stack-100M, n_heads = 8 query heads share n_kv_heads = 2 key/value heads (a 4:1 ratio; see Multi-Head Attention, MQA, GQA & MLA), so four query heads read the same \(W_K\) slice. We must scale each \(W_Q\) slice per query head, but each shared \(W_K\) slice only once — otherwise a KV slice read by four tripping heads gets multiplied four times over, badly over-shrinking it. Two separate loops:

# stacklm/optim/qk_clip.py  (continued)
@torch.no_grad()
def _clip_projections_(attn, s_max, tau: float) -> int:
    """No-QK-norm path (Kimi K2): per-query-head W_Q, per-KV-head shared W_K.

    Attribute names (n_heads, n_kv_heads, head_dim, groups) are exactly those of
    stacklm.model.attention.Attention -- keep them stable across chapters.
    """
    hd = attn.head_dim
    group = attn.groups                          # q-heads per kv-head (= 4)
    fired = 0

    # (1) Per-query-head scale on W_Q.
    for h in range(attn.n_heads):
        if float(s_max[h]) <= tau:
            continue                             # this head is fine
        eta = (tau / float(s_max[h])) ** 0.5     # sqrt so q AND k share it
        attn.wq.weight[h * hd:(h + 1) * hd].mul_(eta)
        fired += 1

    # (2) Per-KV-head scale on the SHARED W_K, using the group's worst logit.
    for kv in range(attn.n_kv_heads):
        s_grp = float(s_max[kv * group:(kv + 1) * group].max())
        if s_grp <= tau:
            continue
        eta = (tau / s_grp) ** 0.5
        attn.wk.weight[kv * hd:(kv + 1) * hd].mul_(eta)
    return fired

The head that set the group maximum lands exactly at \(\tau\); the other three query heads in its group get their shared key scaled by the same amount and so are pulled slightly further below \(\tau\) — a conservative, safe outcome. Exercise 4 walks the arithmetic. In this configuration the shipped \(\tau = 100\) is a genuine, load-bearing safety net.

The third lever: attention soft-capping¶

StackConfig already carries attn_soft_cap (default 0.0 = off). Set it to 50.0 and the attention scores pass through Gemma-2’s \(c\tanh(s/c)\) before the softmax: a smooth, differentiable ceiling that makes blow-up impossible rather than merely corrected-after-the-fact. Its cost is that the naive implementation materializes the score matrix and abandons SDPA — the same \(T^2\) tax as Tier-2 recording. The modern answer is PyTorch’s FlexAttention (torch.nn.attention.flex_attention), which lets you express the cap as a score_mod function that the compiler fuses into a FlashAttention-style kernel, keeping the memory win (Kernel Fusion, torch.compile, CUDA Graphs & Compilers).

The three levers form a clean hierarchy, and you should be able to say which is which in an interview:

QK-norm removes the failure mode structurally. Soft-cap bounds it analytically. QK-clip is the monitor plus correction — its real value at 100M scale is not that it fires, but that its trigger log tells you your LR is too high days before the loss does.

The WSD Schedule: Warmup–Stable–Decay¶

Now the schedule. The dominant pretraining schedule of 2023–2024 was cosine annealing (Chapter 3.10): warm up, then decay the LR along a cosine from peak to a small floor over the entire run. Cosine has one fatal inconvenience for a project like ours: you must know the total step count in advance, because the cosine is stretched to fit it. Decide to train longer and the whole curve is wrong; the model trained to step \(T\) was on a schedule optimized for \(T\), not \(2T\). And cosine gives you no natural place to swap the data distribution.

Warmup–Stable–Decay (WSD; MiniCPM, Hu et al., 2024; adopted by DeepSeek) fixes both. It has three phases:

where \(f\) decays from 1 to (near) 0 over the final phase. MiniCPM found a \(1-\sqrt{\cdot}\) or exponential-style decay works better than linear; we use the \(1-\sqrt{\cdot}\) form. Two properties make WSD ideal for the capstone:

1. The stable phase is schedule-agnostic in length. You hold \(\eta_{\max}\) constant for as long as your token budget allows. Because there is no baked-in endpoint, you can decide to keep going — extend the stable phase, or branch off multiple decay runs from the same stable checkpoint to compare data mixes. This is the “continuable pretraining” property MiniCPM highlighted, and it is why Ch. 14.7 deliberately stops at step 34,332 and hands over a pre-decay checkpoint.
2. The decay phase is where the loss drops fastest — and that is precisely where we anneal on premium data. Empirically WSD shows a characteristic sharp loss drop once decay begins. The capstone exploits this: the WSD decay phase is the mid-training annealing phase of Chapter 14.8. We spend the constant-LR stable phase on the bulk FineWeb-Edu mix, then, as the LR decays, we shift the data toward higher-quality Cosmopedia, math, and code. Low LR + high-quality data = the model “polishes” on the best tokens without the high-LR thrash that would otherwise wash them out. Cosine cannot cleanly do this because its LR is already low through most of the middle of training.

# stacklm/optim/schedule.py
import math

def wsd_lr(step: int, *, peak_lr: float, warmup_steps: int, total_steps: int,
           decay_steps: int | None = None, decay_frac: float = 0.2,
           final_frac: float = 0.0) -> float:
    """Warmup-Stable-Decay learning rate (MiniCPM, Hu et al. 2024).

    - Linear warmup for `warmup_steps`.
    - Constant `peak_lr` through the stable phase.
    - 1 - sqrt() decay over the final phase, down to `final_frac * peak_lr`
      (we use 0.0, i.e. anneal fully to ~0).

    Give EITHER `decay_steps` (absolute) or `decay_frac` (a fraction of
    total_steps -- what Ch. 14.7's TrainConfig passes, 0.10). `decay_steps` wins.

    Call it with peak_lr=1.0 to get a pure MULTIPLIER, which is how the training
    loop drives Muon's 0.02 and AdamW's 3e-3 off one shared curve.
    """
    if decay_steps is None:
        decay_steps = int(decay_frac * total_steps)
    stable_end = total_steps - decay_steps               # first decay step
    if step < warmup_steps:                              # --- warmup ---
        return peak_lr * (step + 1) / warmup_steps
    if step < stable_end:                                # --- stable ---
        return peak_lr
    # --- decay: 1 - sqrt(progress) ---
    progress = (step - stable_end) / max(1, decay_steps)
    # CLAMP: without min(...,1.0) any step past total_steps gives sqrt(p) > 1 and
    # a NEGATIVE learning rate -- live risk, since Ch. 14.8 continues past the
    # pretrain horizon. Past the end we sit at the floor.
    decay_mult = 1.0 - math.sqrt(min(progress, 1.0))     # 1 -> 0
    floor = final_frac * peak_lr
    return floor + (peak_lr - floor) * decay_mult

Putting It Together: Building the Optimizers¶

Here is the routing function the training loop calls once at startup. It walks the model’s named parameters, sends every 2D hidden matrix to Muon and everything else to AdamW, and returns both. The dimensionality test (p.ndim == 2 and “not the embedding”) is the whole trick.

# stacklm/optim/build.py
import torch
from stacklm.optim.muon import Muon

def build_optimizers(model, muon_lr=0.02, adamw_lr=3e-3,
                     weight_decay=0.1, betas=(0.9, 0.95), batched_muon=False):
    """Route Stack-100M params: 2D hidden matrices -> Muon, everything else
    (tied embedding, RMSNorm/QK-norm gains, any 1D/bias) -> AdamW.
    Returns (muon, adamw). Capstone peaks: 0.02 and 3e-3 (~6.7:1, empirical).

    The two peaks are NOT a derivation of each other. Callers must scale both by
    the SAME WSD multiplier, never overwrite both with one shared LR.
    """
    muon_params, adamw_params = [], []
    # The tied embedding is exposed as model.tok_emb.weight (output head shares it).
    embed_ids = {id(model.tok_emb.weight)}
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if id(p) in embed_ids:
            adamw_params.append(p)                 # embedding -> AdamW (sparse rows)
        elif p.ndim == 2:
            muon_params.append(p)                  # hidden matrix -> Muon
        else:
            adamw_params.append(p)                 # 1D norms/biases -> AdamW
    muon = Muon(muon_params, lr=muon_lr, momentum=0.95, nesterov=True,
                weight_decay=weight_decay, ns_steps=5, batched=batched_muon)
    # Weight-decay split for AdamW: decay the 2D embedding lightly, leave 1D norm
    # gains un-decayed (shrinking a RMSNorm gain toward 0 fights the network --
    # and would fight the QK-clip, which uses those same gains as its knob).
    decay, no_decay = [], []
    for p in adamw_params:
        (decay if p.ndim >= 2 else no_decay).append(p)
    adamw = torch.optim.AdamW(
        [{"params": decay,    "weight_decay": weight_decay},
         {"params": no_decay, "weight_decay": 0.0}],
        lr=adamw_lr, betas=betas, eps=1e-8, fused=torch.cuda.is_available())
    return muon, adamw

The AdamW betas are (0.9, 0.95), the LLM-standard choice (Chapter 3.9): the lower \(\beta_2=0.95\) (vs the vision default 0.999) makes the second-moment estimate more responsive, which matters because language gradients are noisier and more non-stationary than image gradients. We enable fused=True on CUDA for a single-kernel AdamW step. Note the param-group split so that 1D norm gains get no weight decay — and notice this is now doubly important: decaying \(\gamma_q,\gamma_k\) toward zero would slowly starve attention of temperature while the QK-clip is trying to manage exactly that quantity.

The full step: grad accumulation, clipping, bf16, periodic QK-clip¶

The capstone’s effective batch is 524,288 tokens — micro_batch_size 32 × seq_len 2048 × grad_accum_steps 8, the number PLAN §5 rounds to “~0.5M”. Everything runs under bf16 autocast (Mixed Precision, bf16 & FP8 Training); bf16’s wide exponent means we need no loss scaler, unlike fp16.

This is the canonical inner step. Everything earlier in the chapter that says “the training loop does X” means this function.

# stacklm/train/loop.py  (the inner optimizer step, abbreviated)
import torch
from stacklm.optim import build_optimizers, qk_clip_, wsd_lr

MICRO_BSZ      = 32       # sequences per micro-step
GRAD_ACCUM     = 8        # micro-steps per optimizer step  -> 524,288 tokens
GRAD_CLIP      = 1.0
QK_CLIP_EVERY  = 200      # Tier-2 gating; with Tier-1 harvesting, clip every step

def optimizer_step(model, batches, muon, adamw, peaks, step, cfg):
    # ONE curve, TWO peaks: peak_lr=1.0 turns wsd_lr into a pure multiplier.
    mult = wsd_lr(step, peak_lr=1.0, warmup_steps=cfg.warmup_steps,
                  total_steps=cfg.total_steps, decay_frac=cfg.decay_frac)
    for opt in (muon, adamw):
        for g in opt.param_groups:
            g["lr"] = peaks[id(opt)] * mult      # 0.02 * mult / 3e-3 * mult

    muon.zero_grad(set_to_none=True)
    adamw.zero_grad(set_to_none=True)
    for _ in range(GRAD_ACCUM):
        x, y, seq_ids = next(batches)                # (MICRO_BSZ, 2048)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            _, loss = model(x, targets=y, seq_ids=seq_ids)
            loss = loss / GRAD_ACCUM                 # average, not sum, over accum
        loss.backward()                              # accumulate grads

    # Clip the GLOBAL grad norm across ALL params (both optimizers) to 1.0.
    # The RETURN VALUE is the real prize: it is the pre-clip norm. Log it.
    gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)

    muon.step()
    adamw.step()

    # MuonClip / QK-clip: PERIODIC, on a fresh probe batch with post-step weights.
    qk_fired = 0
    if step % QK_CLIP_EVERY == 0:
        rec = {}                                     # dict[layer_idx] -> (n_heads,)
        with torch.no_grad():
            xp, _, sp = next(batches)
            model(xp, seq_ids=sp, record=rec)        # Tier-1 cheap bound by default
        qk_fired = qk_clip_(model, rec, tau=cfg.qk_clip_tau)
    return float(loss) * GRAD_ACCUM, float(gnorm), peaks[id(muon)] * mult, qk_fired

Four details that bite people:

1. Average, don’t sum. We divide the loss by GRAD_ACCUM inside the loop so the gradient magnitude — and thus the effective LR — is independent of how many micro-steps we chose; change the accumulation count and the run behaves identically.

2. Clip globally — but know what the clip actually does to Muon. The grad-norm clip runs over all parameters at once, max_norm = 1.0; clipping each optimizer’s group separately would let one group’s blow-up hide behind the other’s normal-sized gradients. Here is the non-obvious part. Global clipping multiplies every gradient by one scalar \(c < 1\), and zeropower_via_newtonschulz5 begins by dividing its input by \(\lVert X\rVert_F\) — so the orthogonalized direction is exactly invariant to that rescale. Scale all gradients by \(c\) and the Muon update is unchanged (up to how \(c\) perturbs the momentum EMA on subsequent steps). Grad clipping is close to a no-op for the 210 Muon-routed matrices. It bites the AdamW group, where the spike survives: with \(\beta_2 = 0.95\), a one-step \(10\times\) gradient spike raises \(\sqrt{v}\) by only \(\sqrt{0.95 + 0.05\cdot 100} \approx 2.4\times\), so the update still grows ~4× and clipping genuinely helps.

   The consequence for what you monitor: for the Muon group, the clip is not the spike defense — the pre-clip norm it returns is. That log, the non-finite-gradient skip guard (Ch. 14.7), and the QK-clip trigger count are your three real sensors. This is why item 1 of the divergence checklist above is item 1.

3. record is a dict, not a tensor. It maps layer_idx -> (n_heads,). It is created fresh each time and thrown away; there is no .zero_() and no persistent buffer to forget to reset. (Ch. 14.7 uses the preallocated-tensor variant, which does need the zero_() — qk_clip_ accepts both.)

4. QK-clip is post-step. It reads weights after .step(). With Tier-1 harvesting it is cheap enough to run every step; with Tier-2 it must be gated by qk_clip_every. Doing Tier 2 every step on every micro-batch is the mistake that turns a 0.05% safety net into a 10%+ tax (Exercise 8).

Why 524,288 tokens, and the critical batch size¶

The batch is not arbitrary. Chapter 3.10 develops the critical batch size (McCandlish et al., 2018): below it, doubling the batch nearly halves the step count for the same loss (compute is “wasted” on gradient noise you could average away); above it, extra batch buys diminishing returns and mostly costs memory. For a 100M model the critical batch sits on the order of a few hundred thousand to ~1M tokens, so ~0.5M is a sweet spot — large enough that the gradient estimate is clean (which Muon likes, since orthogonalizing a noisy buffer wastes the preconditioner) and small enough that we still take 38,147 update steps to spend the budget. Push the batch to 4M tokens and you would take only ~4,800 steps, each barely better than at 0.5M — a worse use of the same tokens.

Micro-batch sizing is then a pure memory game on the A100. Activation memory scales with MICRO_BSZ × seq_len × d_model × n_layers; at MICRO_BSZ = 32 and seq_len = 2048, Stack-100M fits comfortably in 80GB without activation checkpointing (Ch. 14.7 does the arithmetic). On a 24GB 4090 you drop to MICRO_BSZ = 8 and raise GRAD_ACCUM to 32 — same effective batch, same recipe, more wall-clock, exactly as the spec’s secondary tier promises. The invariant to preserve is the product:

One table for the shared knobs¶

Worked Example: Step Budget, Overhead, and a Newton–Schulz Trace¶

The takeaway: Muon buys faster convergence for essentially free compute at this scale — the Newton–Schulz iterations are dwarfed by the transformer’s own matmuls — provided you do not squander it on kernel launches.

Contrasting the Choices, and What Would Break¶

It is worth stating plainly why each alternative was rejected, because an interviewer will ask.

• Pure AdamW (no Muon). Perfectly safe, and what nanoGPT-2024 used. It simply converges slower per token on the 2D matrices; at a fixed 20B-token budget you leave perplexity on the table. If you distrust Muon’s stability, this is the fallback — route all 2D matrices to AdamW too and drop the peak to 6e-4.
• Pure Muon (no AdamW). Breaks on the embedding: orthogonalizing a 32768 × 512 matrix with sparse per-row gradients is both wrong and slow, and 1D norms cannot be orthogonalized at all. The hybrid is not optional.
• Shampoo / full-matrix preconditioning. Strictly more expressive than Muon and strictly more expensive: you must form and invert (or root) the preconditioner, and keep its state. Muon is the “just the orthogonal factor, via matmuls only” corner of that design space, and at 100M the extra expressiveness does not pay for its wall-clock.
• Cosine instead of WSD. Works, but forces you to commit to a step count and gives no clean data-annealing hook. Since the capstone’s entire mid-training story (Chapter 14.8) rides on the decay phase, WSD is the natural fit. A useful mental model: WSD is cosine with the middle stretched into a flat plateau you can extend at will.
• No QK-clip. Defensible here, and worth saying out loud: with QK-norm on, the logits are bounded a priori by \(\sqrt{d_h}\lVert\gamma_q\rVert_\infty\lVert\gamma_k\rVert_\infty\), so the run is already structurally safe — which is precisely why the shipped \(\tau = 100\) never fires. We keep the clip because at \(\tau = 30\) its trigger log is a free instability sensor. Without QK-norm, it is not optional; it is the thing standing between you and a NaN at hour 14.
• Scaling out. If you move past one GPU, note that Muon’s orthogonalization needs the whole matrix, which fights naive optimizer-state sharding. Moonshot’s Moonlight release shows the fix: a ZeRO-1-style distributed Muon that gathers each matrix, orthogonalizes, and re-shards (Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP). At 100M you will never need it — but it is the reason “just wrap it in FSDP” is not a complete answer.

Further Reading¶

• Keller Jordan, Muon: An optimizer for hidden layers in neural networks (blog, 2024) — the original writeup deriving Newton–Schulz orthogonalization, the tuned quintic, and the nanoGPT speedrun context.
• Liu et al. (Moonshot AI), Muon is Scalable for LLM Training (2025) — the Moonlight report: decoupled weight decay for Muon, the \(0.2\sqrt{\max(m,n)}\) update-RMS matching constant, and ZeRO-1 distributed Muon.
• Kimi Team, Kimi K2: Open Agentic Intelligence (2025) — MuonClip and the per-head QK-clip that tamed attention-logit blow-up under Muon at trillion-parameter scale.
• Henry et al., Query-Key Normalization for Transformers (2020) — the QK-norm whose scale-invariance is the reason this chapter’s clip targets the gains, not the projections.
• Hu et al., MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies (2024) — the Warmup–Stable–Decay schedule and its continuable-pretraining / annealing properties.
• DeepSeek-AI, DeepSeek-V3 Technical Report (2024) — a frontier-scale run on a WSD-style multi-stage schedule.
• McCandlish, Kaplan & Amodei, An Empirical Model of Large-Batch Training (2018) — the gradient-noise-scale / critical-batch-size argument behind the ~0.5M-token effective batch.
• Gupta, Koren & Singer, Shampoo: Preconditioned Stochastic Tensor Optimization (2018) — the full-matrix preconditioner Muon cheaply approximates.
• Kingma & Ba, Adam (2015) and Loshchilov & Hutter, Decoupled Weight Decay Regularization (2019) — the baseline this whole stack extends, and the decoupling recipe Muon reuses.
• Yang & Hu et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer (muP, 2022) — the principled alternative to this chapter’s empirical LR-transfer protocol.

Exercises¶

1. In build_optimizers() the tied token embedding (32768 x 512) is routed to AdamW even though it is a 2D tensor and would pass a naive p.ndim == 2 test. Explain why orthogonalizing this matrix with Muon is both wrong and wasteful, and name the property of embedding gradients that makes AdamW the right choice.

2. The RMS-matching scale is \(0.2\sqrt{\max(m,n)}\). Take the SwiGLU gate weight w_gate of shape 1408 x 512. (a) What is the root-mean-square element magnitude of the raw orthogonalized update \(O = UV^\top\) before scaling? (b) Compute the scale factor. © Confirm the post-scale RMS. (d) State precisely what the scale buys you and what it does not — in particular, why the capstone still runs Muon at 0.02 and AdamW at 3e-3 rather than one shared number.

3. Suppose you re-budget the run to 30B tokens at the same 524,288-token effective batch, keeping warmup_steps = 2_000 and decay_frac = 0.10. (a) How many total optimizer steps? (b) Give the warmup / stable / decay step counts. © Using the \(1-\sqrt{\cdot}\) decay, what Muon LR and what AdamW LR are in effect at one quarter of the way through the decay phase? (d) At which step would Ch. 14.7’s pretraining script hand over its pre-decay checkpoint?

4. (The no-QK-norm configuration.) You set qk_norm = False to reproduce Kimi K2’s setup, so qk_clip_ takes the _clip_projections_ path, and you keep the shipped \(\tau = 100\) (which is now load-bearing, since there is no a-priori bound). A block has n_heads = 8 query heads and n_kv_heads = 2 (GQA group size 4). The per-head max pre-softmax logits recorded this step are

s_max = [120, 90, 100, 80, 150, 60, 70, 200]   # heads 0..7

(a) Which query heads get their \(W_Q\) slice scaled, and by what factor? (b) What single factor scales each shared \(W_K\) group? © Verify that the head that set group 1’s maximum lands exactly at \(\tau\), and show that head 4 (also in group 1) ends up below \(\tau\). (d) What would go wrong if you instead scaled \(W_K\) once per query head?

5. The chapter stresses WSD’s continuable property: you branch multiple decay runs off the same stable checkpoint, and the decay length is not tied to any pre-committed total. wsd_lr() accepts decay_steps directly — but it still computes stable_end = total_steps - decay_steps, so a horizon is baked in. Implement wsd_lr_branch() that takes stable_steps and decay_steps directly, with no total_steps at all. Keep the linear warmup, the constant stable phase, and the \(1-\sqrt{\cdot}\) decay to final_frac * peak_lr.

# stacklm/optim/schedule.py  (continued)
import math

def wsd_lr_branch(step: int, *, peak_lr: float, warmup_steps: int,
                  stable_steps: int, decay_steps: int,
                  final_frac: float = 0.0) -> float:
    """WSD LR with phase lengths given DIRECTLY (no baked-in total horizon).

    Decouples the decay from any pre-committed step count, so multiple decay
    runs of different lengths can branch off one stable-phase checkpoint
    (MiniCPM's continuable-pretraining property). Call with peak_lr=1.0 to
    get a multiplier and preserve the Muon/AdamW ratio.
    """
    stable_end = warmup_steps + stable_steps          # first decay step
    if step < warmup_steps:                           # --- warmup ---
        return peak_lr * (step + 1) / warmup_steps
    if step < stable_end:                             # --- stable ---
        return peak_lr
    # --- decay: 1 - sqrt(progress), clamped to [0, 1] ---
    progress = (step - stable_end) / max(1, decay_steps)
    progress = min(progress, 1.0)                     # stay at the floor past the end
    decay_mult = 1.0 - math.sqrt(progress)            # 1 -> 0
    floor = final_frac * peak_lr
    return floor + (peak_lr - floor) * decay_mult

6. You want a sanity check that zeropower_via_newtonschulz5 really conditions the update. Write a short script that builds an ill-conditioned random 256 x 128 matrix (singular values spanning a couple orders of magnitude, condition number \(\sim100\)), runs the orthogonalizer, and prints the singular-value spread of the result. Explain what you expect to see and why the chapter says we do not need machine-precision orthogonality. Then extend it to check the batched path against the looped one.

import torch
from stacklm.optim.muon import (zeropower_via_newtonschulz5,
                                _orthogonalize_bucketed)

torch.manual_seed(0)
m, n = 256, 128
# Random orthonormal U, V and a bad spectrum from 1e-1 to 1e1 (cond ~1e2).
U, _ = torch.linalg.qr(torch.randn(m, m))
V, _ = torch.linalg.qr(torch.randn(n, n))
s = torch.logspace(-1, 1, n)                      # 1e-1 ... 1e1, span 1e2
G = (U[:, :n] * s) @ V.T                          # ill-conditioned (m x n)

print("input  sigma: min %.3e  max %.3e  cond %.3e"
      % (s.min(), s.max(), s.max() / s.min()))

O = zeropower_via_newtonschulz5(G, steps=5).float()
sv = torch.linalg.svdvals(O)
print("output sigma: min %.4f  max %.4f  cond %.4f"
      % (sv.min(), sv.max(), sv.max() / sv.min()))

# Batched path must agree to reduction-order noise, on Stack-100M's shapes.
dirs = [torch.randn(*s_) for s_ in [(512, 512), (128, 512),
                                    (1408, 512), (512, 1408)]]
looped = [zeropower_via_newtonschulz5(d) for d in dirs]
batched = _orthogonalize_bucketed(dirs, 5)
print("max |looped - batched| =",
      max(float((a - b).abs().max()) for a, b in zip(looped, batched)))

input  sigma: min 1.000e-01  max 1.000e+01  cond 1.000e+02
output sigma: min 0.6818     max 1.2023     cond 1.7634
max |looped - batched| = 4.02e-07

7. (The invariance, from scratch.) Prove that with QK-norm enabled, replacing \(W_Q \leftarrow \eta W_Q\) leaves every attention logit unchanged (ignore \(\epsilon\)), and that replacing \(\gamma_q \leftarrow \eta\gamma_q\) multiplies every logit by exactly \(\eta\). Then: Stack-100M has \(d_h = 64\) and shares one gain vector per layer across all 8 query heads. (a) State the a-priori bound on \(|s_{ij}|\) at initialization. (b) You are debugging, so you set qk_clip_tau = 30. A layer’s worst head reports \(S_{\max} = 45\). What factor does _clip_qk_norm_gains_ apply to \(\gamma_q\) and to \(\gamma_k\), and what happens to a head in the same layer that was sitting at \(S_{\max} = 12\)? © Why is that acceptable, and what one-line architecture change would make the clip per-head instead?

# Ch. 14.4 Attention.__init__, per-head QK-norm gains.
# q, k are (B, n_heads, T, head_dim) at this point, so the gain must carry a
# singleton time axis to broadcast correctly: (n_heads, 1, head_dim).
self.q_norm = RMSNorm((self.n_heads,    1, self.head_dim), cfg.norm_eps)
self.k_norm = RMSNorm((self.n_kv_heads, 1, self.head_dim), cfg.norm_eps)

8. (Cost accounting.) Your training loop records the exact per-head max logit (Tier 2) on every micro-batch of every step, with micro_batch_size = 32, seq_len = 2048, n_heads = 8, n_layers = 30, grad_accum_steps = 8, in fp32. (a) You computed the per-layer score-tensor size in §”Harvesting \(S_{\max}\)”: 4.29 GB. Explain why it does not accumulate across the 30 layers, and what the real damage is instead. (b) Roughly how many extra FLOPs per optimizer step does the score matmul cost, versus the model’s \(6ND \approx 3.19\times10^{14}\)? © Give two ways to cut this by more than 100× and state what each gives up.

9. (The clip that isn’t.) A colleague reports that their Muon run spikes badly at step 4,000, and proposes tightening the global gradient clip from max_norm = 1.0 to 0.25. (a) Using the first line of zeropower_via_newtonschulz5, argue precisely what effect that change has on the update applied to the 210 Muon-routed matrices. (b) What effect does it have on the AdamW group? Quantify roughly, using \(\beta_2 = 0.95\) and a one-step \(10\times\) gradient spike. © What should they change instead?