14.5 Mini Scaling Laws: Fit Your Own Law Before Spending the Budget¶

You are about to spend 20 billion tokens training Stack-100M — roughly 22 to 29 A100-hours at USD 25 to 50, with the exact config frozen in The Architecture and the budget derived in Overview & Landscape. Before you commit that, you should be able to answer one question with a number, not a shrug: what final loss should this run reach, and is the 100M/20B allocation actually the right one? If your prediction and your run disagree by more than a couple of tenths of a nat, something is broken — a data-pipeline bug, a mistuned learning rate, a tokenizer mismatch — and you want to know that on hour two, not hour thirty.

The general theory of how loss scales with parameters \(N\) and tokens \(D\) is developed in Scaling Laws: Kaplan, Chinchilla & Beyond. That chapter gives you the form of the law and the folk constants. This chapter does the thing the theory chapter tells you to do but cannot do for your specific setup: it fits the law to your own recipe, on your own data, with your own tokenizer, and extrapolates to the exact model you are about to train. The constants \(E\), \(A\), \(B\), \(\alpha\), \(\beta\) are not universal — they depend on the corpus (FineWeb-Edu + Cosmopedia + a little code and math), the vocabulary (32,768), and every architecture choice frozen into the recipe. Importing another lab’s numbers is how you mis-budget a run.

The plan of attack follows Hoffmann et al.’s Chinchilla methodology at miniature scale: train a ladder of four tiny models (about 4M, 9M, 19M, 44M non-embedding parameters) under the identical recipe, each at several token budgets; fit \(L(N,D)=E+A/N^{\alpha}+B/D^{\beta}\) to the 17 measured losses; cross-check with IsoFLOP profiles; then extrapolate to 84.5M and predict Stack-100M’s loss. The ladder costs about 4.7 GPU-hours (~USD 8), plus another 0.9 for a data-mixture screen — together under 14% of the flagship run, and the cheapest insurance you will ever buy. Then we use the fitted law to justify the plan’s most counterintuitive decision: deliberately over-training to ~200 tokens/parameter, an order of magnitude past compute-optimal.

Along the way we fix three things that quietly wreck small-scale ladders and that most treatments skip: the \(6ND\) FLOP rule is wrong by 1.6×–3.6× at this scale and wrong by a different factor at every rung; the learning rate does not transfer across width unless you make it; and a fixed batch size gives your cheapest runs too few optimizer steps for their own schedule to mean anything.

A Ladder Under One Frozen Recipe¶

A scaling-law fit is only as trustworthy as the invariance of everything you are not varying. The single rule that makes a ladder valid is this: hold the entire recipe fixed and change only \(N\) and \(D\). Same tokenizer, same data mixture and ordering, same optimizer (Muon + AdamW), same WSD schedule shape, same RMSNorm/RoPE/NoPE/GQA/SwiGLU/QK-norm stack from the plan. If the architecture drifts between rungs you are no longer fitting a law — you are fitting noise, which is precisely the confound that broke Kaplan et al.’s original allocation (a learning-rate schedule that was not matched to each run’s token count).

“Hold the recipe fixed” is subtler than it sounds, because two things you must change with width are the learning rate and the batch size. We return to both after the parameter accounting.

Scaling the ladder: deep-and-thin at every rung¶

Stack-100M is deep-and-thin (30 layers × 512 wide), following the small-model result that at fixed parameters, more layers × narrower width beats the reverse (MobileLLM, Liu et al. 2024). We scale the ladder the same way — growing d_model and n_layers together so every rung keeps roughly the target’s depth-to-width aspect ratio. Head dimension stays pinned at 64; the small rungs collapse GQA to a single KV head (MQA), the natural small-model limit of the target’s 8:2 grouping.

The ladder spans about 1.05 decades in \(N\) (3.93M → 43.95M) and the target sits just 0.28 decades above the top rung — a short, honest extrapolation. That closeness is a feature of the USD-100 project: you are not predicting a 400B model from 100M runs (a 3.5-decade leap of faith), you are predicting 84.5M from 44M. The law barely has to reach.

The embedding subtlety: which \(N\) do you fit?¶

Here is a trap unique to small models that the frontier-scale literature glosses over. The recipe uses a 32,768-token vocabulary, and the tied input/output embedding is a \(32768 \times 512 = 16.78\text{M}\)-parameter table (Press & Wolf 2017; the vocab-size tradeoff is analyzed in Tokenization: BPE, WordPiece, Unigram & Byte-Level). At the 4M rung that table is larger than the entire rest of the model, and its share of total parameters falls monotonically up the ladder — the last column above. That is a systematically size-correlated quantity sitting inside your capacity variable.

Two conventions are defensible, and the literature is split, which is the single biggest source of confusion in this area:

• Kaplan et al. (2020) fit on non-embedding parameters, \(N_{\text{nonembed}}\), paired with \(C = 6 N_{\text{nonembed}} D\).
• Hoffmann et al. (Chinchilla, 2022) fit and report total parameters. Chinchilla’s famous “20 tokens per parameter” is 20 tokens per total parameter.

This is not stylistic. Pearce & Song, Reconciling Kaplan and Chinchilla Scaling Laws (arXiv 2406.12907, TMLR 2024) show that counting non-embedding rather than total parameters — combined with running the study at small scale — is a primary driver of the discrepancy between Kaplan’s \(N^\star \propto C^{0.73}\) and Chinchilla’s \(N^\star \propto C^{0.50}\). When they simulate the Chinchilla study under Kaplan’s conventions, they reproduce Kaplan-like biased coefficients. The effect is largest precisely below 1B — where we live.

# stacklm/scaling/ladder.py  -- ladder configs + parameter accounting.
# The nonembed_params() arithmetic mirrors stacklm.config.count_params (Ch. 14.4),
# minus the RMSNorm gains (35,072 params at the target: 0.03% of the total, and
# not what the A/N^alpha capacity term models).
from dataclasses import dataclass

@dataclass(frozen=True)
class LadderConfig:
    """A single rung. Every field except (d_model, n_layers) is *derived* so the
    recipe stays frozen: head_dim pinned at 64, SwiGLU width ~= 2.75*d_model
    rounded to a multiple of 64, KV heads = 1 on the small rungs (MQA limit)."""
    name: str
    d_model: int
    n_layers: int
    head_dim: int = 64
    n_kv_heads: int = 1
    vocab_size: int = 32768
    seq_len: int = 2048            # pretraining context; enters the FLOP count

    @property
    def n_heads(self) -> int:
        assert self.d_model % self.head_dim == 0
        return self.d_model // self.head_dim

    @property
    def intermediate(self) -> int:                 # SwiGLU hidden width
        return int(round(2.75 * self.d_model / 64) * 64)

    def nonembed_params(self) -> int:
        """Parameters in the transformer blocks (the capacity variable we fit).
        Per block: Q,O are d*d; K,V are d*(n_kv*head_dim) under GQA; SwiGLU is
        3*d*intermediate (gate, up, down)."""
        d, kv, hd, inter = self.d_model, self.n_kv_heads, self.head_dim, self.intermediate
        attn = 2 * d * d + 2 * d * (kv * hd)       # Q,O  +  K,V
        mlp  = 3 * d * inter                       # SwiGLU: gate, up, down
        return self.n_layers * (attn + mlp)

    def embed_params(self) -> int:
        return self.vocab_size * self.d_model      # tied: counted once

    def total_params(self) -> int:
        return self.nonembed_params() + self.embed_params()

LADDER = [
    LadderConfig("S1", d_model=192, n_layers=10),
    LadderConfig("S2", d_model=256, n_layers=13),
    LadderConfig("S3", d_model=320, n_layers=17),
    LadderConfig("S4", d_model=448, n_layers=21),
]
TARGET = LadderConfig("Stack-100M", d_model=512, n_layers=30, n_kv_heads=2)
BY_NAME = {c.name: c for c in LADDER}

def family(d_model: int, n_kv_heads: int = 2) -> LadderConfig:
    """Smooth continuation of the ladder's deep-and-thin aspect ratio,
    n_layers ~= d_model/19.2. This is the SEARCH SPACE for allocation questions;
    the target itself is pinned by the plan at d=512 / L=30."""
    return LadderConfig(f"d{d_model}", d_model=d_model,
                        n_layers=max(4, round(d_model / 19.2)), n_kv_heads=n_kv_heads)

if __name__ == "__main__":
    for c in LADDER + [TARGET]:
        print(f"{c.name:10s} d={c.d_model:3d} L={c.n_layers:2d} "
              f"heads={c.n_heads} kv={c.n_kv_heads} inter={c.intermediate:4d}  "
              f"N_nonemb={c.nonembed_params():>10,}  total={c.total_params():>11,}")
    # Stack-100M prints N_nonemb=84,541,440  total=101,318,656

The target line reads N_nonemb=84,541,440 total=101,318,656, reproducing the plan’s parameter budget (16.78M tied embedding + 84.54M in blocks; add the 35,072 norm gains and you land on count_params’s 101,353,728). That arithmetic reproducibility is the point: if your own counter disagrees, fix it now, because every downstream FLOP and dollar estimate rides on it.

The other confound: hyperparameters do not transfer across width¶

Checklist item 3 below insists every rung gets its own decayed schedule. There is a second, equally silent confound: the optimal learning rate depends on width, and our rungs span \(d_{\text{model}}\) from 192 to 512. Train them all at one naive LR and you systematically under-train the wide rungs or destabilize the narrow ones — a width-correlated error, exactly the kind that rotates \(\alpha\) rather than merely adding noise.

The principled fix is muP / µTransfer (Yang et al., Tensor Programs V, arXiv 2203.03466): reparameterize initialization scales, per-group learning rates, and the output multiplier so the optimal hyperparameters become width-invariant, then tune once and transfer. Microsoft’s mup package implements this for PyTorch by recording base shapes from a reference model; Cerebras-GPT (Dey et al., 2023) is the public example of a Chinchilla-style ladder trained under muP. The mechanics are derived in Learning Rate Schedules, Warmup, Batch Size & Hyperparameters; here we need only the operational consequence, and it is worth stating muP’s actual prescription because it is routinely misquoted:

Under muP with Adam, hidden weights and the readout take \(\eta \propto 1/\text{fan\_in}\); input embeddings, biases, and all 1D gains keep a width-independent \(\eta\).

Norms are not rescaled. Applying a \(1/d\) scaling to RMSNorm gains is not muP — it is a width-correlated error dressed up as one.

Our optimizer is the Muon + AdamW hybrid from Chapter 14.6, which splits the parameters into three groups that map cleanly onto that prescription:

• Muon — the 2D hidden matrices. Muon orthogonalizes the momentum and rescales by \(0.2\sqrt{\max(m,n)}\) (the RMS-matching trick, Liu et al., Muon is Scalable, 2025). The orthogonalized update’s spectral norm is set by the shape and the rescaling divides that shape dependence out, so Muon’s peak LR is far closer to width-invariant than Adam’s. Chapter 14.6 measures it — a \(\times 2\) sweep on this chapter’s S4 rung, transferred across batch size — and lands on muon_lr = 6e-3. We use that number, hold it constant across every rung, and re-verify width-invariance with a short probe at S1 and S3 rather than assuming it.
• AdamW, tied embedding/head. The tied table is simultaneously the input embedding (muP: \(\Theta(1)\)) and the readout (muP: \(\Theta(1/d)\)), which muP does not cleanly cover. We follow the readout rule, \(\eta \propto 1/d_{\text{model}}\), because the output side is where the width-dependent logit scale actually bites, anchored at the target’s \(3\times10^{-3}\) (= muon_lr / 2, the 2:1 ratio Chapter 14.6 fixes) with base width \(d_{\text{base}}=512\).
• AdamW, norms and 1D. Width-independent, held at \(3\times10^{-3}\) on every rung.

Everything else is fixed by the run’s own step count: linear warmup for min(2000, 5% of steps), a constant stable phase, then the WSD \(1-\sqrt{\cdot}\) decay over the final 20% to ~0, exactly as in Chapter 14.6. Weight decay 0.1 (0.0 on 1D), betas \((0.9,0.95)\), grad-clip 1.0, bf16.

from mup import MuAdam, set_base_shapes, MuReadout   # microsoft/mup
base  = build_model(d_model=512)         # base width == the target's width
delta = build_model(d_model=1024)        # a second width, to infer what scales
model = build_model(d_model=192)         # the rung you are about to train
set_base_shapes(model, base, delta=delta)   # tags every param with its scaling
opt = MuAdam(model.parameters(), lr=3e-3)   # per-group 1/fan_in applied for you

Compute Accounting: What \(6ND\) Leaves Out¶

To lay out the sweep we must convert each (N, D) run into FLOPs and dollars. The workhorse everyone quotes is the dense-transformer rule from the scaling-laws chapter: 2 FLOPs per parameter per token forward, 4 backward, so \(C \approx 6ND\). It is an excellent rule at 7B. At 100M it is not, and at 4M it is off by a factor of 3.6.

Three terms make up the real per-token training cost of a decoder-only transformer (all “×3” factors are forward + backward \(\approx 3\times\) forward):

1. Block matmuls. \(2 N_{\text{nonembed}}\) FLOPs/token forward \(\Rightarrow 6 N_{\text{nonembed}}\). This is the \(6ND\) term.
2. Attention scores and the value-weighted sum. Per layer, \(QK^\top\) over a length-\(s\) sequence costs \(2 s^2 d_{\text{model}}\) and \(AV\) another \(2 s^2 d_{\text{model}}\) — but causal masking halves both (a FlashAttention-style kernel never computes the masked blocks), so per token it is \(2 s\, d_{\text{model}}\) per layer forward. Over \(L\) layers, ×3: \(6 L s\, d_{\text{model}}\). GQA does not change this: the K/V heads are broadcast back to all query heads, so the score matmul still sums over the full \(d_{\text{model}}\).
3. The tied output head. A \(d_{\text{model}} \times V\) matmul, \(2 d_{\text{model}} V\) forward, ×3: \(6 d_{\text{model}} V\). The input embedding is a gather and costs no meaningful FLOPs; the head does all of the vocabulary work.

Working the rule out for our ladder at \(s=2048\), \(V=32768\) is sobering:

Read the last column again. \(6ND\) undercounts by a factor that varies by more than 2× across the ladder. At S1 the vocabulary head alone costs more than the entire block stack. Two runs you naively call “exactly \(C\)” can differ by 2.3× in real compute — which means IsoFLOP slices built on \(6ND\) are not iso-FLOP at all, and their vertices are pulled toward whichever rung’s true cost you underestimated least. It also means a ladder costed with \(6ND\) is 2–3× cheaper on paper than in the cloud.

Keep \(6ND\) as the mental rule of thumb — it is right about the shape of the cost, and it is what everyone quotes. Just do not let it be your sweep’s cost model.

# stacklm/scaling/flops.py
def flops_per_token(cfg) -> dict:
    """Full training FLOPs per token: blocks + causal attention + tied head.
    Each term is (forward MACs x 2) x 3 for forward+backward.
      blocks:    2 * N_nonembed fwd -> 6 * N_nonembed
      attention: QK^T and AV are 2*s^2*d each; causal masking halves both, so
                 per TOKEN it is 2*s*d per layer fwd -> 6 * L * s * d
      head:      the d x V logits matmul, 2*d*V fwd -> 6 * d * V
                 (the input embedding is a gather: no meaningful FLOPs)"""
    blocks = 6.0 * cfg.nonembed_params()
    attn   = 6.0 * cfg.n_layers * cfg.seq_len * cfg.d_model
    head   = 6.0 * cfg.d_model * cfg.vocab_size
    return dict(blocks=blocks, attn=attn, head=head, total=blocks + attn + head)

def training_flops(cfg, n_tokens: float) -> float:
    """Total training FLOPs -- the number you budget and schedule against."""
    return flops_per_token(cfg)["total"] * n_tokens

def training_flops_6nd(cfg, n_tokens: float) -> float:
    """The 6ND approximation, kept ONLY for contrast and for reproducing
    published numbers that were quoted that way."""
    return 6.0 * cfg.nonembed_params() * n_tokens

def gpu_hours(flops: float, peak_flops_per_s: float = 312e12,
              mfu: float = 0.35) -> float:
    """Wall-clock on ONE accelerator. 312 TFLOP/s ~ A100 bf16 dense peak. MFU is
    measured against the FULL model FLOPs above -- the only honest denominator.
    Ch. 14.1 derives the 0.30-0.45 band for this shape; Ch. 14.7 measures it."""
    return flops / (peak_flops_per_s * mfu) / 3600.0

Applied to the flagship, \(D = 2.0\times10^{10}\) tokens at \(7.97\times10^{8}\) FLOPs/token gives \(C = \mathbf{1.59\times10^{19}}\) FLOPs — against \(1.01\times10^{19}\) from \(6N_{\text{nonembed}}D\). At Chapter 14.1’s attention-inclusive MFU band this is the canonical ≈22–29 GPU-hour, USD 25–50 stable-phase envelope. We do not re-derive it here; that chapter owns the budget, Chapter 14.7 owns the measurement, and Chapter 14.12 owns the invoice — this chapter’s job is only to make sure the sweep is costed on the same footing as the run it is predicting.

Designing the sweep: an IsoFLOP backbone plus off-diagonal points¶

Two forces pull against each other. To identify the law you want spread in both directions — some runs param-limited (few tokens per param), some data-limited. To keep the sweep cheap you cannot run the 44M rung to 400 tokens/param (that single run is \(8.2\times10^{18}\) FLOPs, half the flagship). The resolution is a hybrid grid:

• An IsoFLOP backbone. Four compute budgets a factor of 3 apart, \(C \in \{10^{16},\,3\times10^{16},\,9\times10^{16},\,2.7\times10^{17}\}\). Inside each budget, run every rung whose implied \(D = C/\text{FLOPs-per-token}\) lands at a sane tokens/param (roughly 6–500). Because every run in a slice costs the same true \(C\), these slices give the IsoFLOP method its raw material.
• A few off-diagonal fixed-model points. IsoFLOP slices all lie on lines of constant compute — a degenerate direction that cannot separate \(\alpha\) from \(\beta\) in the additive form. Adding six runs (every rung pushed down to 12 tokens/param, plus the two cheap rungs pushed to 400 and 150) breaks that degeneracy and lets the parametric fit pin the exponents.

Batch size: every rung needs enough steps for its schedule to mean anything¶

A fixed batch_tokens = 2**19 for all runs is the other silent ladder-killer. The cheapest run in our sweep sees \(4.7\times10^{7}\) tokens; at 0.5M tokens per step that is 90 optimizer steps including warmup and the WSD decay tail — not a measurement, a rounding error. Worse, a 0.5M-token batch is far above the critical batch size for a 4M-parameter model at that token count (McCandlish et al., 2018; Zhang et al., How Does Critical Batch Size Scale in Pre-training?, arXiv 2410.21676, ICLR 2025, report a fit in which critical batch size grows roughly as \(D^{0.47}\) and depends only weakly on \(N\)). Past the critical batch size the extra data parallelism buys almost no per-step progress, so those runs land systematically high in loss — and they are exactly the low-tokens-per-param points you added to pin \(\beta\). The bias lands on the exponent you were trying to measure.

So we scale the batch with the run, under two constraints — at least ~2000 optimizer steps and at or below the critical batch size — rounded down to a power of two:

# stacklm/scaling/sweep.py -- design the sweep, size its batches, cost it out.
import math
from stacklm.scaling.ladder import LADDER, TARGET, BY_NAME
from stacklm.scaling.flops import flops_per_token, training_flops, gpu_hours

# (1) IsoFLOP backbone: (true_compute_budget_C, rung). Every run in a slice costs
#     exactly C under the FULL FLOP model -- so the slices really are iso-FLOP.
ISO = [(1e16, "S1"), (1e16, "S2"),
       (3e16, "S1"), (3e16, "S2"), (3e16, "S3"),
       (9e16, "S1"), (9e16, "S2"), (9e16, "S3"),
       (2.7e17, "S2"), (2.7e17, "S3"), (2.7e17, "S4")]
# (2) Off-diagonal fixed-model points: (rung, tokens_per_param on N_nonembed).
EXTRA = [("S1", 12), ("S2", 12), ("S3", 12), ("S4", 12), ("S1", 400), ("S2", 150)]

MIN_STEPS = 2000

def critical_batch_tokens(D: float) -> float:
    """Empirical critical-batch-size fit (Zhang et al. 2025): B* grows ~ D^0.47
    and depends only weakly on N. Illustrative constants -- re-fit on your own
    runs; the point is that B* is a function of D, not a constant."""
    return 22.91 * D ** 0.47

def batch_tokens(D: float) -> int:
    """Tokens per optimizer step: small enough for >= MIN_STEPS steps AND at or
    below the critical batch size, rounded DOWN to a power of two, clipped to
    [2**14, 2**19]. Sanity check: at the flagship's D=2.0e10 this returns
    2**19 = 524,288 -- exactly the batch Ch. 14.6 freezes."""
    b = min(D / MIN_STEPS, critical_batch_tokens(D))
    b = 2 ** int(math.floor(math.log2(b)))
    return int(min(max(b, 2 ** 14), 2 ** 19))

def build_runs():
    runs = []                                  # each run is a dict the harness consumes
    for C, name in ISO:
        c = BY_NAME[name]
        D = C / flops_per_token(c)["total"]    # tokens s.t. TRUE cost == C exactly
        runs.append(dict(cfg=c, N=c.nonembed_params(), D=D, C=C,
                         tpp=D / c.nonembed_params(), kind="iso"))
    for name, tpp in EXTRA:
        c = BY_NAME[name]
        D = float(tpp * c.nonembed_params())
        runs.append(dict(cfg=c, N=c.nonembed_params(), D=D,
                         C=training_flops(c, D), tpp=float(tpp), kind="fixed"))
    for r in runs:                             # size the batch, then the schedule
        r["batch_tokens"] = batch_tokens(r["D"])
        r["steps"] = int(r["D"] / r["batch_tokens"])
        r["warmup_steps"] = min(2000, max(50, round(0.05 * r["steps"])))
    return runs

if __name__ == "__main__":
    runs = build_runs()
    for r in runs:
        print(f'{r["cfg"].name:3s} {r["kind"]:5s} C={r["C"]:.2e} '
              f'D={r["D"]/1e9:7.4f}B tpp={r["tpp"]:6.1f} '
              f'batch={r["batch_tokens"]:7d} steps={r["steps"]:6d}')
    ladder = sum(r["C"] for r in runs)
    D_big  = 20e9                                          # PLAN sec.2: ~20B tokens
    big    = training_flops(TARGET, D_big)
    print(f"{len(runs)} runs; ladder = {ladder:.3e} FLOPs "
          f"= {100*ladder/big:4.1f}% of the flagship ({big:.3e})")
    print(f"ladder wall-clock (1xA100, 35% MFU) = {gpu_hours(ladder):.2f} GPU-hr "
          f"~ USD {gpu_hours(ladder)*1.75:.2f}")
    print(f"flagship = {gpu_hours(big, mfu=0.45):.0f}/{gpu_hours(big):.0f}/"
          f"{gpu_hours(big, mfu=0.30):.0f} GPU-hr at MFU 0.45/0.35/0.30")

The printout is the sweep you are about to run — 17 jobs, none of them degenerate:

S1  iso   C=1.00e+16 D= 0.1177B tpp=  29.9  batch=  32768 steps=  3593
S2  iso   C=1.00e+16 D= 0.0684B tpp=   7.5  batch=  32768 steps=  2087
S1  iso   C=3.00e+16 D= 0.3532B tpp=  89.8  batch= 131072 steps=  2694
S2  iso   C=3.00e+16 D= 0.2052B tpp=  22.4  batch=  65536 steps=  3131
S3  iso   C=3.00e+16 D= 0.1237B tpp=   6.6  batch=  32768 steps=  3774
S1  iso   C=9.00e+16 D= 1.0596B tpp= 269.5  batch= 262144 steps=  4042
S2  iso   C=9.00e+16 D= 0.6157B tpp=  67.2  batch= 262144 steps=  2348
S3  iso   C=9.00e+16 D= 0.3710B tpp=  19.7  batch= 131072 steps=  2830
S2  iso   C=2.70e+17 D= 1.8471B tpp= 201.7  batch= 262144 steps=  7045
S3  iso   C=2.70e+17 D= 1.1131B tpp=  59.2  batch= 262144 steps=  4246
S4  iso   C=2.70e+17 D= 0.5777B tpp=  13.1  batch= 262144 steps=  2203
S1  fixed C=4.01e+15 D= 0.0472B tpp=  12.0  batch=  16384 steps=  2880
S2  fixed C=1.61e+16 D= 0.1099B tpp=  12.0  batch=  32768 steps=  3354
S3  fixed C=5.47e+16 D= 0.2256B tpp=  12.0  batch=  65536 steps=  3442
S4  fixed C=2.47e+17 D= 0.5275B tpp=  12.0  batch= 262144 steps=  2012
S1  fixed C=1.34e+17 D= 1.5729B tpp= 400.0  batch= 262144 steps=  6000
S2  fixed C=2.01e+17 D= 1.3738B tpp= 150.0  batch= 262144 steps=  5240
17 runs; ladder = 1.846e+18 FLOPs = 11.6% of the flagship (1.593e+19)
ladder wall-clock (1xA100, 35% MFU) = 4.70 GPU-hr ~ USD 8.22
flagship = 32/41/47 GPU-hr at MFU 0.45/0.35/0.30

Every rung gets 2,000–7,000 steps, so its WSD warmup/stable/decay phases all mean something; the batch never exceeds the critical size; and the same batch_tokens() rule, evaluated at the flagship’s 20B tokens, returns exactly the plan’s 0.5M-token batch — hence Chapter 14.6’s 38,147 steps. The sweep is ~12% of the flagship compute, about 4.7 GPU-hours, roughly USD 8 — which is precisely the “scaling ladder ~4–5 GPU-hours, ~USD 8” line in Chapter 14.1’s cost table. Note the honest scale effect: at the frontier a scaling ladder is well under 1% of the target run. Here the top rung (44M) is half the target (84.5M), so the ladder is an unavoidably larger fraction — but 12% of a day-and-a-half run to know the answer before you spend the other 88% is still the best trade in the project.

The sweep harness¶

The harness is thin: it drives the real training loop from Chapter 14.7 and records one JSON line per run plus the full loss curve (we need those curves later, for the live monitor). Every name below is the one the shipped stacklm package actually exports — Stack100M/StackConfig from Chapter 14.4, PackedMemmapDataset from Chapter 14.2, pretrain from Chapter 14.7. In CI the loop runs a handful of steps on the tiny synthetic corpus (hermetic, CPU-only); on the A100 it runs for real.

# stacklm/scaling/run_sweep.py
import json, os
from stacklm.config import StackConfig
from stacklm.model import Stack100M
from stacklm.data import PackedMemmapDataset          # Ch. 14.2 memmap shards
from stacklm.train import pretrain                    # Ch. 14.7 loop
from stacklm.scaling.sweep import build_runs, rung_lrs

SEQ_LEN, MICRO_BS = 2048, 16                          # 32,768 tokens per micro-batch

def run_one(run, data_dir, out_path, device="cuda"):
    c = run["cfg"]
    model = Stack100M(StackConfig(                    # frozen recipe: only d/L vary
        vocab_size=c.vocab_size, d_model=c.d_model, n_layers=c.n_layers,
        n_heads=c.n_heads, n_kv_heads=c.n_kv_heads, head_dim=c.head_dim,
        intermediate=c.intermediate, max_seq_len=SEQ_LEN, rope_theta=10000.0,
        tie_embeddings=True, nope_every=4, qk_norm=True,   # identical across rungs
    ))
    train = PackedMemmapDataset(os.path.join(data_dir, "train"))
    val   = PackedMemmapDataset(os.path.join(data_dir, "val"))

    grad_accum = max(1, run["batch_tokens"] // (MICRO_BS * SEQ_LEN))
    lrs = rung_lrs(c)                                 # Muon 6e-3; AdamW head 3e-3*512/d
    out = pretrain(
        model, train, device=device,
        steps=run["steps"], total_steps=run["steps"],  # schedule ends at THIS run's D
        micro_batch_size=MICRO_BS, grad_accum=grad_accum,
        muon_lr=lrs["muon"], adamw_lr=lrs["adamw_head"],
        warmup_steps=run["warmup_steps"], decay_frac=0.20,
        weight_decay=0.1, grad_clip=1.0, qk_clip_every=200,
        log_every=max(1, run["steps"] // 200),        # ~200 curve points for the monitor
        eval_dataset=val, seed=1234,
    )
    row = dict(name=c.name, N=run["N"], D=run["D"], C=run["C"],
               tpp=run["tpp"], kind=run["kind"], steps=run["steps"],
               batch_tokens=run["batch_tokens"],
               val_loss=out["val_loss"],              # LR fully decayed to THIS run's D
               curve=[(i * run["batch_tokens"], l)    # [(tokens, train_loss), ...]
                      for i, l in enumerate(out["loss_history"])])
    with open(out_path, "a") as f:
        f.write(json.dumps(row) + "\n")
    return row

if __name__ == "__main__":
    out = "ladder_results.jsonl"
    if os.path.exists(out):
        os.remove(out)
    for run in build_runs():
        r = run_one(run, data_dir=os.environ["STACKLM_DATA"], out_path=out)
        print(f"{r['name']:4s} N={r['N']/1e6:5.1f}M D={r['D']/1e9:5.3f}B "
              f"tpp={r['tpp']:6.1f} steps={r['steps']:5d} -> val_loss={r['val_loss']:.4f}")

Designing your own ladder: a checklist¶

If you adapt this to a different target, the design choices above generalize. Every item exists to defeat a specific failure mode.

1. Span at least a decade in \(N\). Ours covers 3.9M → 44M (1.05 decades). Less spread and \(\alpha\) is unconstrained; the power law needs leverage.
2. Keep the top rung within ~0.3 decades of the target. We extrapolate 44M → 84.5M. Extrapolating a factor of 2 is honest; a factor of 100 is a prayer.
3. Give every rung its own decayed schedule. Non-negotiable — this is the Kaplan confound, and it is a one-line bug (a shared total_steps) that silently rotates your whole fit.
4. Cost the sweep with a full FLOP model, not \(6ND\). At sub-1B scale the correction is 1.6×–3.6× and rung-dependent, so \(6ND\) slices are not iso-compute.
5. Scale the LR with width (muP), or prove you do not have to. State the per-rung LR in the config, not in a comment. A width-correlated LR error rotates \(\alpha\) just as surely as a schedule error does.
6. Size the batch per run: \(\ge\) ~2000 steps and \(\le\) the critical batch size. Echo steps in the printout so a 90-step run is visible and gets rejected before it poisons \(\beta\).
7. Freeze everything else byte-for-byte. Same tokenizer, data order, nope_every, qk_norm, optimizer split. If you must change one thing (say, MLA vs GQA), that is a separate ladder, not a mixed one.
8. Budget 10–20% of the flagship compute at this scale (well under 1% at frontier scale). Ours is ~12% for the law plus ~2% for the data-mixture screen.
9. Add a few seeds at the cheapest rung. Two or three reruns of S1 at one budget measure your noise floor directly, which tells you the Huber delta to use and whether two rungs really differ or just wiggled.
10. Hold out the top rung, refit, and check the prediction. The only success criterion that matters is extrapolation.

The through-line: a scaling ladder is a measurement instrument, and like any instrument it is worthless if uncalibrated. Every item above is a calibration step.

Fitting Your Own L(N, D)¶

With ladder_results.jsonl in hand we fit the Chinchilla parametric form:

\(E\) is the irreducible floor (the entropy of the mixture under this tokenizer), \(A/N^{\alpha}\) the finite-capacity penalty, \(B/D^{\beta}\) the finite-data penalty. Two robustness tricks — both from Chinchilla’s appendix — separate a real fit from a toy one: build \(\log L\) via logsumexp of the three log-terms, which is numerically stable across the orders of magnitude the terms span; and minimize the Huber loss of the log-space residual, so the few noisiest runs cannot dominate.

The code below is self-contained: because you cannot run 17 GPU jobs inside a textbook, it synthesizes the loss table from a plausible ground-truth law and fits that — proving the machinery recovers what it should, and giving us a truth to check against. In practice you delete the synthesis block and json.load your real results; the fitting code is byte-for-byte identical.

import numpy as np
from scipy.optimize import minimize
from scipy.special import logsumexp
from stacklm.scaling.sweep import build_runs

# --- Load the measured ladder. In production: read ladder_results.jsonl. ---------
# Here we synthesize from a ground-truth law so the demo is hermetic. These
# constants are what a real Stack-recipe ladder on this mix would look like: a
# higher floor E than raw web text (small vocab, educational corpus) and a strong
# data term (tiny models are data-hungry at these token counts).
GROUND_TRUTH = dict(E=2.45, A=124.0, alpha=0.33, B=234.0, beta=0.30)
def _law(N, D, p=GROUND_TRUTH):
    return p["E"] + p["A"]*N**(-p["alpha"]) + p["B"]*D**(-p["beta"])

rng = np.random.default_rng(0)
runs = build_runs()
# dtype=float matters: an integer N array silently truncates the fit's E term.
N_obs = np.array([r["N"] for r in runs], dtype=float)
D_obs = np.array([r["D"] for r in runs], dtype=float)
# ~1% multiplicative noise mimics seed / data-order variation between runs:
L_obs = _law(N_obs, D_obs) * (1.0 + 0.01*rng.standard_normal(len(runs)))

# --- The fit. Parameterize positives via logs: E=exp(e), A=exp(a), B=exp(b). -----
def predict_log_loss(theta, N, D):
    e, a, b, alpha, beta = theta
    terms = np.stack([
        np.full_like(N, e),            # log E
        a - alpha * np.log(N),         # log(A * N^-alpha)
        b - beta  * np.log(D),         # log(B * D^-beta)
    ])
    return logsumexp(terms, axis=0)    # = log(E + A N^-alpha + B D^-beta), stably

def huber(r, delta=1e-3):
    ar = np.abs(r)
    return np.where(ar <= delta, 0.5*r**2, delta*(ar - 0.5*delta))

def fit(N, D, L, rng, bounds_exp=(0.25, 0.42), n_starts=60):
    """Multi-start L-BFGS-B on a non-convex surface. We CONSTRAIN the exponents
    to the band the scaling literature consistently reports -- with only 4 rungs
    spanning ~1 decade in N, an unconstrained fit will happily rail alpha to an
    absurd 0.1 or 0.6 and take the extrapolation with it. Bounding the exponents
    to physically-sane values is the single most important regularizer here."""
    def objective(theta):
        return np.sum(huber(predict_log_loss(theta, N, D) - np.log(L)))
    best, best_val = None, np.inf
    for _ in range(n_starts):
        x0 = np.array([rng.uniform(0.4, 1.1), rng.uniform(3, 7), rng.uniform(3, 7),
                       rng.uniform(*bounds_exp), rng.uniform(*bounds_exp)])
        res = minimize(objective, x0, method="L-BFGS-B",
                       bounds=[(0.2, 1.3), (0, 12), (0, 12), bounds_exp, bounds_exp])
        if res.fun < best_val:
            best, best_val = res.x, res.fun
    return best

best = fit(N_obs, D_obs, L_obs, rng)
e, a, b, alpha, beta = best
print(f"fitted:  E={np.exp(e):.3f}  A={np.exp(a):.1f}  alpha={alpha:.3f}  "
      f"B={np.exp(b):.1f}  beta={beta:.3f}")
print(f"allocation exponent  beta/(alpha+beta) = {beta/(alpha+beta):.3f}")

Read the fit honestly: the loss extrapolates, the constants do not¶

Run that across ten noise seeds and a very specific pattern appears. The allocation exponent \(\beta/(\alpha+\beta)\) scatters over 0.37–0.60 (median 0.45) around the ground truth \(0.30/0.63 = 0.476\). The individual constants swing far more wildly: \(E\) from 2.13 to 2.69 against a true 2.45, and \(A\) from 58 to 430 against a true 124 — a factor of seven.

That is not a bug, and it is the single most important thing to internalize about scaling-law fits: \(E\), \(A\), and \(B\) are strongly correlated and only weakly identified. The additive surface is nearly flat near its optimum, so many \((E,A,B)\) triples fit the same 17 points almost equally well; they trade off against each other and only their combination is pinned. If your colleague’s fit reports different constants, neither of you is necessarily wrong. Epoch AI’s replication of Chinchilla (Besiroglu et al., 2024) found the original published constants fragile for exactly this reason, with implausibly narrow confidence intervals. Be equally honest about the exponents: at frontier-ladder sizes \(\beta/(\alpha+\beta)\) is fairly robust, but at our size, with 17 points and 1% run-to-run noise, it wanders by ±0.1. Quote it as “roughly one-for-one, Chinchilla-like,” never to three digits.

What is reliably pinned is the quantity you actually care about — the extrapolated loss. Because the fit interpolates the measured surface faithfully, and the target sits only 0.28 decades beyond the top rung, the compensating errors in \(E\), \(A\), \(B\) largely cancel where it matters:

def predicted_loss(theta, N, D):
    # dtype=float AGAIN: np.full_like on an integer array truncates log E to 0
    # and silently returns a loss ~1.2 nats too low. This is the single most
    # common bug in home-grown scaling-law code.
    return float(np.exp(predict_log_loss(theta, np.array([N], dtype=float),
                                         np.array([D], dtype=float)))[0])

N100 = 84_541_440                       # non-embedding params
for label, D in (("Chinchilla ~20 tok/TOTAL-param", 20 * 101_318_656),
                 ("the plan's 20B budget",          20e9)):
    print(f"{label:32s} D={D/1e9:5.2f}B  predicted L={predicted_loss(best, N100, D):.3f}"
          f"  (ground truth {_law(N100, D):.3f})")
# across 10 noise seeds:
#   @ D=2.03B : predicted L in 3.06-3.19   (ground truth 3.128)
#   @ D=20.0B : predicted L in 2.84-3.07   (ground truth 2.940)
# The ~0.1-nat BAND is the stable thing; the individual digits are not.

Across noise seeds the extrapolation stays within about 0.1 nats of the truth, with no consistent bias. That is exactly the resolution you should expect from a four-rung ladder — and exactly enough. Report the predicted loss with an honest ±0.1-nat band, not five decimal places of \(A\). A ±0.1-nat band is useless for splitting hairs between two good runs but perfect for its real job: a broken run misses by 0.3+ nats, and that you will catch on hour two.

The sharpest version of that validation is checklist item 10 — fit without the top rung and predict it:

# Hold out S4 entirely, refit on S1-S3, predict the two held-out S4 points.
mask = np.array([r["cfg"].name != "S4" for r in runs])
theta_ho = fit(N_obs[mask], D_obs[mask], L_obs[mask], np.random.default_rng(7))
pred = np.exp(predict_log_loss(theta_ho, N_obs[~mask], D_obs[~mask]))
print("held-out S4:", np.round(pred, 3), "vs actual", np.round(L_obs[~mask], 3))
# Across seeds the held-out S4 predictions land within 0.02-0.09 nats of the
# measured values. THAT is the number that licenses the leap to 84.5M --
# not the fitted A, and not the R^2 on the training points.

Which parameter count? Refit under both conventions¶

We chose \(N_{\text{nonembed}}\) and promised evidence. Here it is — the same 17 measurements, fit twice, with the exponent bounds widened to \([0.15,0.60]\) so we can see the exponent move rather than watch it hit a wall:

N_total = np.array([r["cfg"].total_params() for r in runs], dtype=float)

for seed in range(6):
    rng_s = np.random.default_rng(seed)
    L_s = _law(N_obs, D_obs) * (1 + 0.01*rng_s.standard_normal(len(runs)))
    t_ne = fit(N_obs,   D_obs, L_s, np.random.default_rng(1000+seed), bounds_exp=(0.15, 0.60))
    t_to = fit(N_total, D_obs, L_s, np.random.default_rng(1000+seed), bounds_exp=(0.15, 0.60))
    print(f"seed{seed}: non-embedding alpha={t_ne[3]:.3f} | total-param alpha={t_to[3]:.3f}")
# non-embedding: alpha in 0.29-0.60, median ~0.43  (ground truth 0.33)
# total params : alpha in 0.43-0.60, median ~0.59, at or within 0.02 of the
#                upper bound in four of six seeds

The non-embedding fits scatter around the truth. The total-parameter fits sit systematically higher and pile up against the bound. The mechanism is the embed-share column of the ladder table: the embedding is \(32768 \cdot d\), growing only linearly in width, while block parameters grow like \(n_{\text{layers}} \cdot d^{2}\), so the embedding’s share of \(N_{\text{total}}\) falls monotonically (62% → 17%). That shrinking near-constant baseline flattens the low-\(N\) end of the curve and forces the fitter to rotate \(\alpha\) upward to compensate. That rotation is Kaplan’s \(C^{0.73}\) in miniature, and it is what Pearce & Song reconstruct at full scale.

Two honest caveats. First, this demo is rigged in the non-embedding convention’s favour, because the synthetic ground truth was generated on \(N_{\text{nonembed}}\); it shows that mixing conventions between generation and fitting corrupts \(\alpha\), not that nature prefers one. Second, the deeper point stands regardless: whichever variable you fit, say so, and make your compute variable consistent with it.

IsoFLOP profiles: Chinchilla’s second method¶

The parametric fit is Chinchilla’s Approach 3 (fit the whole surface, differentiate). Its independent cross-check is Approach 2, the IsoFLOP method, which never commits to the parametric form and is therefore robust to its misspecification. At each fixed budget \(C\), the loss as a function of \(\log N\) (with \(D\) forced by the FLOP constraint) is a U-shaped valley — too-small models are param-limited, too-large models are data-starved. Fit a parabola in \(\log N\), read off the vertex, and you have the compute-optimal \(N^\star(C)\) without ever assuming a power law. Do it for several budgets and the slope of \(\log N^\star\) vs \(\log C\) is the allocation exponent \(a\) in \(N^\star \propto C^{a}\).

There is a catch at our scale: a parabola needs enough points bracketing the minimum on both arms, or its vertex is garbage. Our four-rung ladder puts only two to four points in each slice — too few. So we do the honest thing: demonstrate the method on a properly dense grid (7 models per slice, the way Hoffmann et al. actually ran it), and use the real ladder’s coarse valleys only as a loose sanity check. A production ladder would simply include more rungs.

Because the slices must be genuinely iso-compute, the demo searches over real configurations in the deep-and-thin family() and forces \(D\) from the full FLOP model:

import math, numpy as np
from scipy.optimize import minimize_scalar
from stacklm.scaling.ladder import family
from stacklm.scaling.flops import flops_per_token

def slice_configs(C, n=7, span=0.6):
    """Seven real configs for one IsoFLOP slice: coarse-scan for the valley,
    then span +/- `span` decades of N around it. CENTERING MATTERS -- an
    off-centre grid is one of the documented biases of the parabola method."""
    coarse = [family(d) for d in range(64, 1345, 64)]
    losses = [_law(c.nonembed_params(), C / flops_per_token(c)["total"]) for c in coarse]
    N_c = coarse[int(np.argmin(losses))].nonembed_params()
    return [min((family(d) for d in range(64, 2049, 32)),
                key=lambda c: abs(math.log(c.nonembed_params() / t)))
            for t in N_c * np.logspace(-span, span, n)]

def true_optimum(C):
    """The exact valley of the deep-and-thin family, from a CONTINUOUS relaxation
    (n_layers = d/19.2, intermediate = 2.75d, no rounding). A discrete config grid
    is a staircase whose lumps can move a vertex by 20% -- do not use one as your
    ground truth, or you will 'measure' a bias that is really a rounding artifact."""
    def shape(d):                                  # (N_nonembed, FLOPs per token)
        L, inter = d / 19.2, 2.75 * d
        N = L * (2*d*d + 2*d*128 + 3*d*inter)      # n_kv_heads=2 -> 2*d*(2*64)
        return N, 6*N + 6*L*2048*d + 6*d*32768
    def objective(log_d):
        N, c = shape(math.exp(log_d))
        return _law(N, C / c)                      # D forced by the constraint
    r = minimize_scalar(objective, bounds=(math.log(48), math.log(4096)),
                        method="bounded")
    return shape(math.exp(r.x))[0]

rng2 = np.random.default_rng(1)
C_slices = np.array([1e16, 3e16, 9e16, 2.7e17])
N_star = []
for C in C_slices:
    cfgs = slice_configs(C)
    Ns = np.array([c.nonembed_params() for c in cfgs], dtype=float)
    Ds = np.array([C / flops_per_token(c)["total"] for c in cfgs])   # TRUE iso-FLOP
    Ls = _law(Ns, Ds) * (1 + 0.01*rng2.standard_normal(len(cfgs)))
    c2, c1, c0 = np.polyfit(np.log(Ns), Ls, 2)      # parabola in log N
    N_star.append(np.exp(-c1 / (2.0 * c2)))         # vertex = valley = optimal N
    print(f"C={C:.1e}: parabola N*={N_star[-1]/1e6:6.2f}M   "
          f"true N*={true_optimum(C)/1e6:6.2f}M")

a_iso, _ = np.polyfit(np.log(C_slices), np.log(N_star), 1)
a_true, _ = np.polyfit(np.log(C_slices), np.log([true_optimum(C) for C in C_slices]), 1)
print(f"IsoFLOP a = {a_iso:.3f}   (true a = {a_true:.3f})")

The true valleys of this family under this law sit at 8.59M, 14.31M, 23.77M, 39.47M, giving \(a_{\text{true}} = 0.463\). Two results are worth separating carefully, because they are usually conflated:

• On noise-free data the parabola is accurate here. Run it with the noise term removed and the vertices reproduce the true valleys to better than 1%, giving \(a = 0.464\)–\(0.468\); deliberately widening the grid to \(\pm 1.2\) decades or off-centring it by \(\pm 0.4\) decades still keeps \(a\) inside \(0.458\)–\(0.468\). For this law and this family the systematic bias is small.
• Noise dominates. Add the same 1% run-to-run noise the real ladder has and \(a\) across ten seeds spans 0.44–0.57 (median 0.49), with individual vertices moving by up to 2× at a single budget. Four slices of seven runs simply cannot pin an exponent to better than ±0.05.

That distinction is the practical lesson. Czech et al. (2026) show that the parabolic approximation can carry systematic bias even on noise-free data — driven by grid width, off-centre sampling, and the asymmetry of the loss surface when \(\alpha \ne \beta\) — and that Approach 3 removes it if you exploit the objective’s partially-linear structure (variable projection) rather than throwing a generic optimizer at it. We measured our own bias and found it small; you can only make that statement because you have ground truth, and on real data you do not. So: centre each slice on its valley, keep the span tight, run enough seeds to see the noise, and treat the vertex as an estimate with an error bar.

The practical reading: run both methods and require them to agree on the story, not the digit. Ours do — parametric \(\beta/(\alpha+\beta) \approx 0.45\), IsoFLOP \(a \approx 0.49\), true \(0.463\); all of them say “grow \(N\) and \(D\) together, roughly one-for-one.” That is the Chinchilla result, reproduced from scratch, with an honest error bar.

Extrapolating to Stack-100M and Predicting Its Loss¶

Now the payoff. We have a fitted law; plug in the target’s non-embedding count, \(N=84.54\text{M}\), and read off the loss at the two allocations that matter.

Use the book’s interactive Scaling-Law-Optimal tool to poke at these tradeoffs — drop in the fitted \(E,A,B,\alpha,\beta\) and slide the compute budget to watch \((N^\star, D^\star)\) move:

Using the ladder as a live monitor¶

The chapter opened with a promise: catch a mismatch on hour two. Delivering it requires one more step, because the fitted law predicts the final, LR-decayed loss, while at hour two your run is mid-stable-phase at full LR and therefore sitting well above its eventual value. Comparing a live number to a final prediction makes every healthy run look broken. Two ingredients fix it, and you already paid for both when you saved the ladder’s loss curves.

1. Measure your own WSD decay drop. The anneal at the end of a WSD run buys a characteristic extra drop (MiniCPM’s headline observation; see Chapter 14.6). Its size depends on your schedule shape and data, so measure it across rungs rather than importing a number.

2. Project the live stable-phase curve to its endpoint. Within the stable phase the loss falls like the law’s data term, so it is linear in \(D^{-\beta}\) with the \(\beta\) your ladder just measured. Fit a straight line in that coordinate, extrapolate to the full budget, subtract the decay drop.

# stacklm/scaling/monitor.py
import numpy as np

def measure_decay_drop(curves, decay_frac=0.20) -> float:
    """Median (loss at start of decay) - (final loss) across the ladder rungs.
    `curves` is a list of (tokens, val_loss) sequences from ladder_results.jsonl.
    Expect a tenth of a nat at this scale -- but YOUR ladder tells you, free."""
    drops = []
    for tok, loss in curves:
        tok, loss = np.asarray(tok, float), np.asarray(loss, float)
        i = min(int(np.searchsorted(tok, tok[-1] * (1.0 - decay_frac))), len(loss) - 1)
        drops.append(loss[i] - loss[-1])
    return float(np.median(drops))

def project_final_loss(tokens, losses, total_tokens, beta, decay_drop,
                       warmup_skip=0.25) -> float:
    """Project a live STABLE-PHASE curve to the run's final, decayed loss.
    beta        : the data exponent from YOUR ladder fit.
    decay_drop  : measure_decay_drop() over the ladder curves.
    warmup_skip : drop the first fraction of points; the warmup transient is not
                  on the power law and will drag the line."""
    t, l = np.asarray(tokens, float), np.asarray(losses, float)
    keep = t > t[-1] * warmup_skip
    x = t[keep] ** (-beta)                        # law is linear in D^-beta ...
    slope, intercept = np.polyfit(x, l[keep], 1)  # ... so a straight line fits
    return float(intercept + slope * total_tokens ** (-beta) - decay_drop)

# Hour two of the flagship: ~1.5B tokens in (7% of the run).
projected = project_final_loss(live_tokens, live_losses, total_tokens=2.0e10,
                               beta=beta, decay_drop=measure_decay_drop(ladder_curves))
print(f"projected final loss {projected:.3f}  vs  ladder prediction 2.94 +/- 0.10")

3. Compare shapes, not just endpoints. Overlay the S4 rung’s curve on the flagship’s with tokens-per-parameter on the x-axis. Under a frozen recipe the two should lie almost on top of each other over their shared range; a healthy flagship simply continues where S4 stopped. A shape divergence — a kink, a plateau, a different slope — localizes the bug in time far better than a single scalar, and it shows up in the first hour.

When the check fails, triage before you restart:

The instrument you built to choose the run is the same instrument that supervises it. That is the whole return on the 12%.

Second Use of the Ladder: Choosing the Data Mixture¶

The plan fixes an approximate mixture — 70% FineWeb-Edu, 15% Cosmopedia v2, 10% code (StarCoder), 5% math (FineMath) — and explicitly delegates tuning those weights to this chapter. A ladder is the right instrument, because the same frozen-recipe discipline that lets you vary \((N, D)\) lets you vary the mixture instead: fix \(N\) and \(D\), vary the weights, rank by held-out loss. This is the small-proxy-model methodology behind DoReMi (Xie et al., 2023), RegMix (Liu et al., 2024) and Aioli (Chen et al., 2024), and it is the cheapest lever on final quality you have.

Two design rules make the ranking mean something.

Rule 1: hold the evaluation set fixed and matched to your deployment target. If you score each candidate on a validation set drawn from its own mixture, you are measuring “how predictable is this mix by a model trained on it,” which rewards low-entropy mixes (crank the code fraction and watch aggregate loss fall). Instead, fix one validation set whose composition reflects what the model must actually do, and also report per-domain losses so you can see the trade you are making.

Rule 2: screen cheap, confirm one rung up. Ai2’s DataDecide (Magnusson et al., ICML 2025) is the definitive study: 25 corpora × 14 model sizes from 4M to 1B, over 30k checkpoints. Its headline is both encouraging and cautionary — ranking candidate corpora at a single small size predicts the best corpus at 1B in roughly 80% of pairwise comparisons, using ~0.01% of the compute. Eighty percent is a wonderful screen and a terrible oracle. So: screen at S1, confirm the finalists at S2, and never let a mixture decision rest on one 4M-parameter run.

# stacklm/scaling/mixture.py -- second use of the ladder: rank data mixtures.
# Every API here is the one Ch. 14.2 / 14.7 / 14.11 actually ship: re-weighting is
# a `DataMixEntry.weight` edit, and per-domain scoring is `compute_perplexity`
# looped over single-domain validation shard directories.
import dataclasses, os
from stacklm.config import StackConfig
from stacklm.model import Stack100M
from stacklm.data import STACK100M_MIX, PackedMemmapDataset, build_corpus
from stacklm.eval import compute_perplexity
from stacklm.train import pretrain
from stacklm.scaling.ladder import BY_NAME
from stacklm.scaling.flops import training_flops, gpu_hours
from stacklm.scaling.sweep import batch_tokens, rung_lrs

# Candidate weights over (fineweb_edu, cosmopedia_v2, starcoder, finemath), in
# STACK100M_MIX order. The plan's default is first; the others probe one axis at
# a time so a win is attributable.
CANDIDATES = {
    "plan":       (0.70, 0.15, 0.10, 0.05),   # PLAN sec.2 default
    "more_synth": (0.55, 0.30, 0.10, 0.05),   # is Cosmopedia worth more weight?
    "less_synth": (0.82, 0.03, 0.10, 0.05),   # ...or is it hurting diversity?
    "more_code":  (0.60, 0.15, 0.20, 0.05),   # code as a reasoning/structure prior
    "more_math":  (0.60, 0.15, 0.10, 0.15),   # math for the downstream agent
    "web_only":   (1.00, 0.00, 0.00, 0.00),   # the null hypothesis -- always include it
}
# A FIXED validation weighting, composed for the deployment target, NOT per
# candidate. Ch. 14.10's agent needs grounded QA over clean prose plus arithmetic.
VAL_WEIGHTS = {"fineweb_edu": 0.55, "cosmopedia_v2": 0.20,
               "starcoder": 0.10, "finemath": 0.15}

def reweight(weights):
    """A new mix = the same sources with new weights. Zero-weight sources are
    dropped so `build_corpus`'s weight-sum assertion still holds."""
    return [dataclasses.replace(e, weight=w)
            for e, w in zip(STACK100M_MIX, weights) if w > 0]

def screen(tokenizer, data_root, rung="S1", tokens_per_param=100, device="cuda"):
    """Train one model per candidate at a FIXED (N, D) and score them all on the
    SAME per-domain held-out shards. Returns rows sorted by weighted loss."""
    cfg = BY_NAME[rung]
    D, rows = int(tokens_per_param * cfg.nonembed_params()), []
    bt = batch_tokens(D)
    lrs = rung_lrs(cfg)
    # Built ONCE and shared by every candidate: one single-domain val corpus per
    # source (build_corpus with entries=[that one entry at weight 1.0]), so we can
    # report per-domain losses instead of one opaque aggregate.
    val = {e.name: PackedMemmapDataset(os.path.join(data_root, "val", e.name, "val"))
           for e in STACK100M_MIX}

    for name, w in CANDIDATES.items():
        corpus = os.path.join(data_root, "mix", name)
        build_corpus(corpus, tokenizer, total_tokens=D, entries=reweight(w),
                     seq_len=cfg.seq_len, holdout_per_mille=0, seed=1234)
        model = Stack100M(StackConfig(
            vocab_size=cfg.vocab_size, d_model=cfg.d_model, n_layers=cfg.n_layers,
            n_heads=cfg.n_heads, n_kv_heads=cfg.n_kv_heads, head_dim=cfg.head_dim,
            intermediate=cfg.intermediate, max_seq_len=cfg.seq_len,
            tie_embeddings=True, nope_every=4, qk_norm=True))
        steps = D // bt
        pretrain(model, PackedMemmapDataset(os.path.join(corpus, "train")), device=device,
                 steps=steps, total_steps=steps,
                 micro_batch_size=16, grad_accum=max(1, bt // (16 * cfg.seq_len)),
                 muon_lr=lrs["muon"], adamw_lr=lrs["adamw_head"],
                 warmup_steps=min(2000, max(50, steps // 20)), decay_frac=0.20,
                 seed=1234)                                  # identical for all mixes
        per_dom = {d: compute_perplexity(model, ds, device=device)["loss_nats_per_token"]
                   for d, ds in val.items()}
        rows.append(dict(name=name, weights=w, per_domain=per_dom,
                         score=sum(VAL_WEIGHTS[d] * per_dom[d] for d in VAL_WEIGHTS)))
    return sorted(rows, key=lambda r: r["score"])

if __name__ == "__main__":                                   # cost the protocol
    c1, c2 = BY_NAME["S1"], BY_NAME["S2"]
    screen_flops  = len(CANDIDATES) * training_flops(c1, 100 * c1.nonembed_params())
    confirm_flops = 2 * training_flops(c2, 60 * c2.nonembed_params())   # top-2 at S2
    total = screen_flops + confirm_flops
    print(f"screen {screen_flops:.2e} + confirm {confirm_flops:.2e} = {total:.2e} FLOPs")
    print(f"= {gpu_hours(total):.2f} GPU-hr ~ USD {gpu_hours(total)*1.75:.2f}")
    # -> 2.00e17 + 1.61e17 = 3.61e17 FLOPs = 0.92 GPU-hr ~ USD 1.61
    #    i.e. 2.3% of the flagship to make the mixture decision on evidence.

Six candidates at S1 (393M tokens each, 3,000 steps) plus two confirmations at S2 (550M tokens, ~2,100 steps) cost \(3.6\times10^{17}\) FLOPs — 2.3% of the flagship, about 0.9 GPU-hours, roughly USD 1.60. That is the entire price of replacing “70/15/10/5 felt right” with a measurement.

Three practical notes. Build the candidate mixes with real tooling: HuggingFace datatrove handles the streaming, filtering, dedup and shard-writing of Chapter 14.2 — and re-weighting is a sampler config edit, not a re-download. Always include the null hypothesis (web_only): if the fancy mix does not beat 100% FineWeb-Edu on your fixed validation set, you have learned something valuable and cheap. And read the per-domain columns, not just the score — a mix that wins overall while doubling your math loss is the wrong mix for a capstone whose downstream agent must do arithmetic. Beyond static weights, the online methods (DoReMi’s group-DRO proxy, Aioli’s on-the-fly mixing-law estimation) and the broader curriculum question are covered in Data Mixing, Domain Weighting & Curriculum.

The Over-Training Decision: Compute-Optimal ≠ Deployment-Optimal¶

Chinchilla answers “what minimizes training loss for a training-compute budget.” That is almost never the real objective. Stack-100M will be trained once and served indefinitely — quantized to int4 and run on a laptop (Ch. 14.11). For a model you deploy, the quantity to minimize is lifetime compute, training plus all future inference:

subject to hitting a target loss \(L^\star\), where \(c(N)\) is the full per-token training cost from the table above. Every generated token costs \(\approx 2N\) FLOPs and the KV cache scales with \(N\) (see The Anatomy of LLM Inference and Inference Economics). The larger \(D_{\text{infer}}\) is, the more it pays to shrink \(N\) and grow \(D_{\text{train}}\) to buy back the loss you gave up. This is the inference-aware regime formalized by Sardana, Frankle et al. (Beyond Chinchilla-Optimal, 2024), and it is why Llama-3-8B saw ~15T tokens (~1900 tok/param), two orders of magnitude past Chinchilla.

Here is the decision made concrete with the ladder’s own numbers:

import numpy as np
from stacklm.scaling.ladder import family
from stacklm.scaling.flops import flops_per_token

def lifetime_optimal(L_target, D_infer, E, A, alpha, B, beta):
    """Pick a REAL config (and its D_train) that hits L_target while minimizing
    TRAIN+INFER FLOPs. We search the deep-and-thin family so the training cost is
    the honest c(N) -- blocks + attention + head -- not 6ND."""
    best = None
    for d in range(128, 2049, 64):
        cfg = family(d)
        N   = cfg.nonembed_params()
        budget = L_target - E - A * N**(-alpha)      # loss left for the data term
        if budget <= 0:
            continue                                 # this N alone overshoots L_target
        D_train = (B / budget) ** (1.0 / beta)
        total = flops_per_token(cfg)["total"] * D_train + 2 * N * D_infer
        if best is None or total < best["total"]:
            best = dict(d=d, N=N, D_train=D_train, total=total, tpp=D_train/N)
    return best

E, A, alpha, B, beta = 2.45, 124.0, 0.33, 234.0, 0.30    # the fitted (here, GT) law
for D_infer in (1e10, 1e12, 1e14):                       # light -> heavy serving
    r = lifetime_optimal(2.94, D_infer, E, A, alpha, B, beta)
    print(f"D_infer={D_infer:.0e}: d={r['d']:4d}  N*={r['N']/1e6:6.1f}M  "
          f"D_train={r['D_train']/1e9:8.2f}B  tok/param={r['tpp']:6.0f}")
# D_infer=1e+10: d= 640  N*= 146.0M  D_train=    9.28B  tok/param=    64
# D_infer=1e+12: d= 448  N*=  49.5M  D_train=   68.04B  tok/param=  1376
# D_infer=1e+14: d= 384  N*=  31.5M  D_train=  465.72B  tok/param= 14805
# As planned serving load rises, the optimal model SHRINKS and tok/param CLIMBS
# -- the quantitative engine behind "over-train a small model for deployment."

One honest caveat before you crank tokens/param to the moon: the Chinchilla form assumes every token is fresh. High-quality data is finite, and past a few epochs of repetition returns collapse (Muennighoff et al.’s data-constrained law; see Scaling Laws and Data Cleaning, Deduplication & Quality Filtering). At 20B tokens on our deduplicated ~20B-token mix we are inside one epoch, so the fresh-token law holds — but the 1376- and 14805-tok/param rows above are fantasy at our corpus size, and that is exactly why the plan invests so heavily in dedup and synthetic Cosmopedia data rather than simply looping the corpus.

Key Takeaways¶

Further Reading¶

• Hoffmann, Borgeaud, Mensch, et al. Training Compute-Optimal Large Language Models (Chinchilla). arXiv 2022. — The three fitting methods this chapter miniaturizes.
• Kaplan, McCandlish, Henighan, et al. Scaling Laws for Neural Language Models. arXiv 2020. — The original power-law form and the LR-schedule confound to avoid.
• Pearce & Song. Reconciling Kaplan and Chinchilla Scaling Laws. TMLR 2024. — Why the parameter-counting convention, at small scale, is the primary source of the two papers’ disagreement.
• Yang, Hu, et al. Tensor Programs V: Zero-Shot Hyperparameter Transfer (muP). arXiv 2022. — The width-scaling rules behind this chapter’s per-rung learning rates.
• Sardana, Frankle, et al. Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Laws. arXiv 2024. — Formalizes lifetime-compute and inference-aware over-training.
• Besiroglu, Erdil, Barnett, et al. Chinchilla Scaling: A Replication Attempt (Epoch AI). 2024. — Why the parametric constants are fragile and the extrapolated loss is the thing to trust.
• Muennighoff, Rush, Barak, et al. Scaling Data-Constrained Language Models. NeurIPS 2023. — The repetition/data-wall limit on how far you can over-train.
• Xie, Pham, et al. DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining. NeurIPS 2023. — Proxy-model domain reweighting, the ancestor of this chapter’s mixture screen.
• Liu, Zheng, et al. RegMix: Data Mixture as Regression for Language Model Pre-training. ICLR 2025. — Many small runs plus a regression surface to predict the best unseen mixture.
• Magnusson, Tai, et al. DataDecide: How to Predict Best Pretraining Data with Small Experiments. ICML 2025. — What does and does not transfer from 4M-parameter proxy runs.
• Liu, Chang, et al. MobileLLM: Optimizing Sub-billion Parameter Language Models. 2024. — The deep-and-thin small-model result the ladder inherits.

Exercises¶

1. (Conceptual.) This chapter fits on non-embedding parameters; Chinchilla fit on total parameters. Suppose a colleague fits \(L(N_{\text{total}}, D) = E + A/N_{\text{total}}^{\alpha} + B/D^{\beta}\) on our 17 ladder points. Explain concretely why this rotates the fitted exponent \(\alpha\), why the problem is specific to small models, and what you would have to change about the compute variable to make the total-parameter convention self-consistent.

2. (Conceptual.) The warning admonition says every rung must get “its own schedule that decays to zero at its token budget \(D\).” Suppose instead you train S4 once for its largest sweep budget (0.578B tokens) on a single WSD schedule, and to obtain S4’s loss at 0.2B and 0.4B tokens you read the value off that one curve at those step counts. Why are those numbers biased, what is the name of the historical mistake this reproduces, and how does the bias interact with the batch-size rule?

3. (Quantitative.) Using the exact arithmetic in LadderConfig.nonembed_params() and flops_per_token(), compute rung S2 (\(d=256\), \(L=13\), head_dim\(=64\), n_kv_heads\(=1\), \(s=2048\), \(V=32768\)) by hand: (a) its SwiGLU intermediate width, the per-block attention and MLP parameter counts, and the total non-embedding count; (b) its tied embedding and the embedding’s share of total parameters; © the three FLOPs-per-token terms and the ratio of the total to \(6N\).

4. (Quantitative.) Cost the flagship run two ways: Stack-100M (\(N_{\text{nonembed}} = 84{,}541{,}440\), \(N_{\text{total}} = 101{,}318{,}656\), \(L\)=30, \(s\)=2048, \(V\)=32768) at the plan’s \(D = 2.0\times10^{10}\) tokens. Compute (a) tokens per parameter under both conventions; (b) \(C\) under \(6N_{\text{nonembed}}D\) and under the full model, and verify the full model equals Chapter 14.1’s \((6N_{\text{total}} + 6Lsd)D\); © single-A100 wall-clock across the 0.30–0.45 MFU band (peak 312 TFLOP/s); (d) the number of optimizer steps at the plan’s 524,288-token batch.

5. (Implementation.) The prediction example claims that for the flagship budget \(C = 1.593\times10^{19}\) FLOPs, the compute-optimal configuration is \(d=768\), 40 layers, \(N^\star \approx 250\text{M}\), \(D^\star \approx 7.86\text{B}\) (~31 tok/param). Write compute_optimal_allocation(C, ...) that finds this by searching real configurations in the deep-and-thin family with \(D\) forced by the full FLOP model. Run it with the ground-truth law, confirm the numbers, then re-run with a \(6ND\) constraint and explain how much the answer moves.

from stacklm.scaling.ladder import family
from stacklm.scaling.flops import flops_per_token

def compute_optimal_allocation(C, E, A, alpha, B, beta, use_6nd=False):
    """Minimize L(N, D) = E + A N^-alpha + B D^-beta subject to the FLOP
    constraint, searching real configs. D is forced by the constraint."""
    best = None
    for d in range(128, 2049, 64):
        cfg = family(d)
        N = cfg.nonembed_params()
        D = C / (6.0 * N) if use_6nd else C / flops_per_token(cfg)["total"]
        L = E + A * N**(-alpha) + B * D**(-beta)
        if best is None or L < best["L"]:
            best = dict(d=d, layers=cfg.n_layers, N=N, D=D, L=L, tpp=D / N)
    return best

E, A, alpha, B, beta = 2.45, 124.0, 0.33, 234.0, 0.30     # the ground-truth law
C_full = 7.9666e8 * 2.0e10                                # 1.593e19
r = compute_optimal_allocation(C_full, E, A, alpha, B, beta)
print(f"d={r['d']} L={r['layers']} N*={r['N']/1e6:.1f}M D*={r['D']/1e9:.2f}B "
      f"tok/param={r['tpp']:.1f} L*={r['L']:.4f}")
# -> d=768 L=40 N*=249.7M D*=7.86B tok/param=31.5 L*=2.9115

6. (Quantitative, hard.) Derive the allocation exponent under a non-\(6ND\) cost model. Let \(c(N)\) be the true per-token training cost and \(\gamma(N) = \mathrm{d}\log c/\mathrm{d}\log N\). (a) Minimize \(L = E + AN^{-\alpha} + BD^{-\beta}\) subject to \(c(N)D = C\) and derive the stationarity condition. (b) Differentiate it to show \(a \equiv \mathrm{d}\log N^\star/\mathrm{d}\log C = \beta/(\alpha + \gamma\beta + \mathrm{d}\log\gamma/\mathrm{d}\log N)\). © Using the measured \(\gamma = 0.759\) and \(\mathrm{d}\log\gamma/\mathrm{d}\log N = 0.092\) at \(C = 9\times10^{16}\), evaluate \(a\) and compare with the numerically measured 0.463 and with the \(6ND\) answer. (d) Explain, in one sentence each, why the real cost model raises the optimal \(N\) at a given \(C\) but slightly lowers the exponent.

7. (Quantitative.) The chapter rejects batch_tokens = 2**19 for all runs. (a) For the cheapest sweep run — S1 at 12 tokens/param — compute \(D\) and the number of optimizer steps a 0.5M-token batch would give. (b) With WSD warmup at 5% of steps and decay over the final 20%, how many steps land in each phase? Is the measurement usable? © Apply the chapter’s batch_tokens() rule and recompute. (d) Verify that the same rule returns \(2^{19}\) at the flagship’s token count.

8. (Design.) You have USD 2 of A100 time (about 1.1 GPU-hours at USD 1.75/hr) to decide between the plan’s default mixture and one alternative that doubles the math fraction. Design the experiment: which rung, how many tokens, what validation set, what you would report, and what result would make you reject the alternative even if its aggregate loss is lower. Then state the one thing this experiment cannot tell you.