14.12 Retrospective: Cost Accounting, Reproducibility, and the Path to 1B¶

Eleven chapters ago we set out to do something specific: take the entire LLM stack — every idea in this book — and compress it into one project small enough to run on a single rented GPU for about the price of a nice dinner, yet real enough that nothing was faked. We trained a byte-level BPE tokenizer, fit our own scaling law on a ladder of tiny models, pretrained Stack-100M on ~20B tokens with a Muon+AdamW hybrid under a WSD schedule, mid-trained it for long context and quality, aligned it with SFT / DPO / GRPO, distilled a narrow tool-using ReAct agent into it, evaluated it honestly, and quantized it to int4 to run on a laptop. The plane is in the air. This chapter lands it.

Landing means four concrete things, and this chapter is those four things: (1) an audit of the prediction — did the scaling law we fit in Chapter 14.5 actually forecast the flagship run’s loss, and how far is such a law allowed to reach; (2) a full cost accounting — every GPU-hour and dollar, stage by stage, derived rather than asserted, and re-priced onto 2026 hardware; (3) a reproducibility discipline — the seeds, config hashes, data manifests, environment pins, checkpoint formats, and release artifacts that separate “I trained a model once” from “anyone can rebuild this”; and (4) an honest scale-up analysis — what actually breaks when you take this exact recipe from 100M to 1B parameters, which libraries you switch to, what it costs, and what you change (data, memory, parallelism, learning-rate and batch scaling, and the Mixture-of-Experts option) to get there.

This chapter builds directly on Scaling Laws: Kaplan, Chinchilla & Beyond for the FLOP arithmetic, The Roofline Model & Performance Engineering for the utilization analysis, Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP and Distributed Training II: Tensor, Pipeline, Sequence & Expert Parallelism for the parallelism transitions, and Checkpointing, Fault Tolerance & Long-Running Jobs for the reproducibility machinery. It stays strictly consistent with the canonical Stack-100M config (vocab_size=32768, d_model=512, n_layers=30, GQA 8:2, SwiGLU intermediate=1408, tied embeddings, ≈101M params) fixed in the capstone plan.

14.12.1 Where We Landed¶

Before we count the money, let us name the artifact precisely, because the cost only means something relative to what it bought.

At the end of the pipeline the reader owns four checkpoints, all derived from one 101M-parameter base:

stacklm-100m/
├── base/           # 18B-token pretrained base (WSD stable phase)   ~2.9-3.1 nats/tok held-out
├── mid/            # +2B tokens: quality anneal -> long-context (8192) -> capability injection
├── chat/           # SFT + DPO on the ChatML-style template
├── agent/          # distilled ReAct traces -> narrow auto-research tool-user
└── agent-int4/     # round-to-nearest int4 weights, ~55-65 MB on disk, runs on a CPU laptop

None of these is a frontier model, and the capstone never pretended otherwise. A 100M model is a narrow instrument: within a scaffolded, retrieval-grounded domain it can produce coherent, useful, grounded text; outside that scaffold it hallucinates freely. What is remarkable is not the model’s raw quality but that every layer of the stack that produces a GPT-4-class system is present here in miniature and actually ran — the tokenizer, the scaling-law fit, the optimizer, the distributed-ready loop, the alignment stack, the agent harness, the quantized deployment.

The single most important lesson the capstone teaches — the one that makes a 100M model in 2026 vastly better than GPT-2 (117M) in 2019 at the same size — is deliberate over-training. Chinchilla-optimal for our ~84.5M non-embedding parameters is ~1.7B tokens; we trained on ~20B, roughly 200 tokens/param and ~12× past compute-optimal. That is economically irrational for a model you train and throw away, and completely rational for a model you will serve, because you pay training compute once and save inference cost forever. Keep that asymmetry in mind — it reappears in the cost table (over-training is most of the bill) and in the scale-up section (it only gets more extreme at 1B).

14.12.2 Did Our Scaling Law Hold?¶

The capstone’s differentiator was that we did not import someone else’s scaling constants. In Chapter 14.5 we trained a four-rung ladder ({3.9M, 8.8M, 19M, 44M} non-embedding parameters), fit \(L(N,D) = E + A/N^{\alpha} + B/D^{\beta}\) to ~17 measured losses, cross-checked with IsoFLOP profiles, and used the result to predict the flagship’s loss before spending the flagship’s money. A retrospective that never asks whether the prediction landed is not a retrospective. So: did it?

# stacklm/scaling_check.py -- close the loop on the Ch. 14.5 ladder fit.
"""Compare the loss the ladder PREDICTED against the loss the flagship ACHIEVED.

Two rules carried over from Ch. 14.5, both easy to get wrong:
  1. N is the NON-EMBEDDING parameter count (the 30 blocks: 84.5M), because the
     ladder fit N against the work the 6ND rule counts. Feeding 101.4M here
     silently corrupts the comparison.
  2. The law predicts loss on the PRETRAIN mix. Evaluate `base/` on a held-out
     shard of that same 70/15/10/5 mix -- not on the mid-training anneal mix,
     whose entropy floor E is different.
"""
LADDER_FIT = dict(E=2.45, A=124.0, alpha=0.33, B=234.0, beta=0.30)  # Ch. 14.5
N_NONEMBED = 84_500_000      # 30 blocks x ~2.82M; excludes the tied 16.8M embedding
FIT_BAND = 0.10              # +/- nats: the honest resolution of a 4-rung ladder


def predicted_loss(n_nonembed: int, n_tokens: float,
                   E: float, A: float, alpha: float,
                   B: float, beta: float) -> float:
    """Chinchilla parametric form, evaluated at (N, D)."""
    return E + A * n_nonembed ** (-alpha) + B * n_tokens ** (-beta)


for label, D in [("Chinchilla-optimal (~20 tok/param)", 1.69e9),
                 ("flagship, stable phase only",        1.80e10),
                 ("flagship, full campaign",            2.00e10)]:
    L = predicted_loss(N_NONEMBED, D, **LADDER_FIT)
    print(f"{label:36s} D={D:.2e}  L_pred = {L:.3f} +/- {FIT_BAND:.2f} nats")
# Chinchilla-optimal (~20 tok/param)   D=1.69e+09  L_pred = 3.150 +/- 0.10 nats
# flagship, stable phase only          D=1.80e+10  L_pred = 2.949 +/- 0.10 nats
# flagship, full campaign              D=2.00e+10  L_pred = 2.940 +/- 0.10 nats

The ladder’s forecast for the flagship base checkpoint is therefore ≈ 2.94 nats/token, with an honest ±0.1-nat band (2.84–3.04), which is exactly the range Chapter 14.11 documents for the held-out perplexity measurement (a loss of ~2.9 nats corresponds to PPL ≈ 18). The prediction and the artifact agree to within the band the fit deserves. That is the outcome you should expect, and it is worth being precise about why, because the failure modes are more instructive than the success.

Why it was allowed to work. The ladder was trained under the identical recipe — same tokenizer, same 32,768 vocab, same 70/15/10/5 mix, same seq_len=2048, same Muon+AdamW hybrid, same WSD shape. Every one of those choices is baked into \(E\), \(A\), and \(B\); change any of them and the constants move. And the extrapolation in \(N\) is short: from a top rung of 44M to a target of 84.5M is a factor of 1.92, or 0.28 decades. Extrapolating a factor of two is engineering; extrapolating a factor of a hundred is a prayer.

Where the strain actually is. It is not in \(N\) — it is in \(D\). The ladder’s rungs were trained near compute-optimal (tens of tokens per parameter); the flagship is deliberately run at ~200 tokens/parameter. So the \(B/D^{\beta}\) term is being evaluated an order of magnitude outside the token range that constrained \(\beta\). Three consequences follow, and all three are things to check rather than assume:

1. The data term is the extrapolated one. At 20B tokens the data penalty \(B/D^{\beta}\) has fallen to ~0.19 nats while the capacity penalty \(A/N^{\alpha}\) sits at ~0.30 nats. The model is capacity-limited, not data-limited — which is the correct diagnosis of an over-trained model, and the quantitative reason further over-training has diminishing returns while a bigger model would not.
2. Repetition is not counted by the law. \(L(N,D)\) assumes \(D\) unique tokens. If your 20B-token budget contains 2 epochs of a 10B-token corpus, the law over-predicts the benefit. Muennighoff et al. (2023) quantify the discount; below ~4 epochs it is small, which is why the plan’s mix is sized for roughly single-epoch coverage.
3. \(E\) is mix-specific. The mid-training anneal deliberately shifts to a cleaner, denser mix, which lowers measured loss partly because the text is more predictable, not because the model got better. Comparing mid/ against a law fit on the pretrain mix will make the model look better than it is. Always evaluate base/ on the pretrain mix’s held-out shard, and treat mid-training’s gain as a separate, separately-measured number.

Finally, the caveat the loss curve cannot tell you: the law predicts loss, not capability. Nothing in \(L(N,D)\) says whether the model can do two-digit arithmetic or follow a ReAct format. Downstream abilities are threshold-y and must be measured, which is what the probes in Chapter 14.11 are for. A scaling law is a budgeting instrument, not an evaluation.

14.12.3 Cost Accounting: Anatomy of the ~USD 100 Model¶

The headline is “the ~USD 100 model,” but a headline is not an accounting. Let us derive the number from first principles — including the FLOPs the usual rule of thumb omits — then reconcile it against a realistic itemized bill, then re-price the whole thing onto 2026 hardware, including the parts nobody advertises: the teacher API calls, the object storage, and the re-run reality tax of failed launches and OOM debugging.

The compute floor: 6ND, and the FLOPs 6ND forgets¶

Training FLOPs for a dense transformer follow the 6ND rule (2 FLOPs/param/token forward, 4 backward), covered in depth in Scaling Laws. For the full pretraining campaign (18B tokens in the stable phase plus ~2B in mid-training):

That number is a floor, and at Stack-100M’s shape it is a surprisingly loose one. The 6ND rule counts only the parameterized matmuls; it ignores the attention score computation, which has no parameters at all. Per token, per layer, causal attention costs \(2 s\, d_{\text{model}}\) FLOPs in the forward pass (\(QK^\top\) over the causal half plus the \(AV\) product, summed across heads), and backward is ~2× forward, so:

For Stack-100M (\(L=30\), \(s=2048\), \(d=512\), \(N=1.014\times10^8\)) that ratio is \(30 \times 2048 \times 512 / 1.014\times10^{8} = \mathbf{0.31}\). Attention adds 31% on top of the 6ND count. Two things follow immediately:

• Deep-and-thin models pay this tax hardest. For a standard block the ratio simplifies to roughly \(s/(12\,d_{\text{model}})\) — it grows with sequence length and shrinks with width. Our narrow \(d=512\) is why it is 31% here, while a 1B model at \(d=2048\) and the same sequence length pays only ~10%. The deep-thin recipe that buys quality per parameter (MobileLLM, Liu et al., 2024) costs efficiency per FLOP. That is a real trade, not a free win.
• Long context multiplies it. At the mid-training length \(s=8192\) the ratio quadruples to ~1.24, i.e. attention now costs more than all the matmuls combined. This is the quantitative version of Chapter 14.8’s warning that pretraining at 8192 from step zero would have been ruinous.

Note that GQA does not help here: 2 KV heads shrink the KV cache 4×, but K and V are still broadcast to 8 query heads, so the attention FLOPs are unchanged. GQA is a memory optimization, not a compute one.

MFU, HFU, and which one your loop logs¶

An NVIDIA A100 (80GB) has a bf16 tensor-core peak of ~312 TFLOP/s dense. You never get peak; you get Model FLOPs Utilization (MFU) — sustained useful FLOP/s divided by peak. Three conventions get called “MFU” and they differ by tens of percent, so state yours:

The flagship config runs without activation checkpointing (Chapter 14.7’s table), so for us HFU = model-FLOPs MFU, and the only correction that matters is the 1.31× attention factor. Concretely: a loop logging MFU(6ND) = 0.45 is really pushing \(0.45 \times 1.31 \approx 0.59\) of A100 peak. That is at the top of what this shape sustains, and it requires torch.compile, a FlashAttention-backed SDPA, and a micro-batch large enough to keep the GEMMs saturated.

Putting the pieces together for the 18B-token stable phase. Hardware FLOPs \(= 6ND \times 1.31 = (6 \times 1.014\times10^{8} \times 1.8\times10^{10}) \times 1.31 \approx 1.43\times10^{19}\). At a sustained \(u = 0.59\) of A100 peak:

which corresponds to ~231,500 tokens/s and MFU(6ND) ≈ 45% — matching the worked example in Chapter 14.7 almost exactly. Run the same arithmetic at \(u = 0.45\) (an un-compiled loop, an unfused attention path, or a smaller micro-batch) and you get 28.4 GPU-hours at MFU(6ND) ≈ 34%. So the honest planning band is 22–29 GPU-hours for the stable phase, and the bill below uses the well-tuned end while the re-run tax and the price axis absorb the rest.

The following helper turns any measured throughput into a cost — feed it the tokens/sec your loop actually logs, not the theoretical peak.

# stacklm/cost.py
"""Turn a training run's measured throughput into GPU-hours and dollars.

Everything here is derived from numbers the training loop already logs
(tokens/sec, total tokens) -- no theoretical peaks required for the bill.
"""
from dataclasses import dataclass

# Stack-100M canonical constants (from the capstone plan; keep in sync with stacklm.config)
N_PARAMS = 101_400_000          # ~101.4M total params (tied embedding counted once)
N_LAYERS, D_MODEL = 30, 512
A100_BF16_PEAK_FLOPS = 312e12   # A100-80GB bf16 tensor-core peak (dense)


def training_flops(n_params: int, n_tokens: float) -> float:
    """Dense-transformer training FLOPs via the 6ND rule (matmuls only)."""
    return 6.0 * n_params * n_tokens


def attention_flops(n_tokens: float, seq_len: int,
                    n_layers: int = N_LAYERS, d_model: int = D_MODEL) -> float:
    """The FLOPs 6ND forgets: causal QK^T and AV, fwd+bwd, summed over layers.

    Per token per layer: 2*s*d forward (causal halves the s^2 term), x3 for
    fwd+bwd => 6*L*s*d per token. Parameter-free, so 6ND misses it entirely.
    """
    return 6.0 * n_layers * seq_len * d_model * n_tokens


def hardware_flops(n_tokens: float, seq_len: int, n_params: int = N_PARAMS,
                   recompute: bool = False) -> float:
    """What the GPU actually executes: matmuls + attention (+ recomputation)."""
    total = training_flops(n_params, n_tokens) + attention_flops(n_tokens, seq_len)
    return total * (8.0 / 6.0) if recompute else total   # full ckpt re-runs the fwd


def gpu_hours_from_throughput(n_tokens: float, tokens_per_sec: float) -> float:
    """The honest number: wall-clock GPU-hours from *measured* throughput."""
    return n_tokens / tokens_per_sec / 3600.0


def mfu_6nd(tokens_per_sec: float, n_params: int = N_PARAMS,
            peak_flops: float = A100_BF16_PEAK_FLOPS) -> float:
    """MFU under the 6ND convention -- what stacklm.train logs. Under-reports."""
    return 6.0 * n_params * tokens_per_sec / peak_flops


def peak_fraction(tokens_per_sec: float, seq_len: int,
                  peak_flops: float = A100_BF16_PEAK_FLOPS, **kw) -> float:
    """Fraction of the accelerator's peak actually being used (attention counted)."""
    return hardware_flops(tokens_per_sec, seq_len, **kw) / peak_flops


@dataclass
class Stage:
    name: str
    gpu_hours: float
    usd_per_gpu_hour: float = 1.80   # illustrative A100-80GB spot price; RE-PRICE IT
    extra_usd: float = 0.0           # non-GPU line items: teacher API, storage, egress

    @property
    def usd(self) -> float:
        return self.gpu_hours * self.usd_per_gpu_hour + self.extra_usd


if __name__ == "__main__":
    tps = 231_500                    # sustained tokens/sec logged by the pretrain loop
    hrs = gpu_hours_from_throughput(18e9, tps)
    print(f"stable phase: {hrs:.1f} GPU-hr  "
          f"MFU(6ND)={mfu_6nd(tps):.1%}  peak_frac={peak_fraction(tps, 2048):.1%}")
    # -> stable phase: 21.6 GPU-hr  MFU(6ND)=45.1%  peak_frac=59.2%

2026 hardware tiers: the same FLOPs, six different bills¶

The capstone’s flagship tier is a single A100-80GB because that is the cheapest thing that comfortably fits the run — but by 2026 it is two generations old, and the cost chapter owes you the comparison. Peak figures below are vendor-published dense bf16 tensor-core numbers (halve any “with sparsity” marketing figure); prices are illustrative mid-2026 spot bands from the usual GPU-rental market and move fast enough that you must re-price before you quote a number.

Three consequences for this project, and one for the next one:

• The newest chip is often the cheapest, not just the fastest. Repricing the 18B-token stable phase onto an H100 at a (deliberately conservative) \(u = 0.45\) — small models get less efficient on bigger chips because they are more launch-bound and less able to fill the tensor cores — gives \(1.43\times10^{19}/(0.45 \times 9.89\times10^{14}) \approx 9.0\) GPU-hours. At USD 2.50/hr that is USD 22.4 against the A100’s USD 38.9 at USD 1.80/hr, and it finishes in 9 hours instead of 22. Exercise 7 works out the break-even rental price.
• FP8 is a ≥1B lever, not a 100M one. Hopper and Blackwell roughly double peak throughput in FP8, and per-tensor/per-block scaling recipes (see Mixed Precision, bf16 & FP8 Training, and torchao’s float8 integration used by torchtitan) make it usable for pretraining. But the realized end-to-end gain is well under the 2× peak ratio, and at 100M you are not GEMM-limited enough to collect most of it, while you are taking on real numerical risk. Turn it on at 1B+, where the GEMMs are fat enough to pay you back.
• Consumer cards are a legitimate tier, not a consolation prize. The flagship’s ~22 GB peak (§14.12.5) leaves a 24 GB 4090 no room for fragmentation or the torch.compile workspace, so Chapter 14.7’s consumer tier drops the micro-batch 4× to 8 sequences and raises gradient accumulation to 32, holding the ~524k-token global batch fixed. Wall-clock roughly triples; the marginal dollar cost of a card you already own is electricity. The owned-hardware scenario in §14.12’s cost sweep — the GPU column collapsing to electricity — assumes exactly this.
• Memory tier, not model size, decides when you shard. Hold that thought — it is why the “7B does not fit one GPU” folklore is now hardware-generation-dependent (§14.12.5).

The itemized bill¶

Every stage of the capstone, priced at an illustrative USD 1.80/GPU-hour A100-80GB spot rate. GPU-hours are the sustained wall-clock the loop would log at the well-tuned end of the band, not the theoretical floor. Non-GPU line items (teacher-model API for agent distillation, object storage for the ~200 GB of tokenized shards and checkpoints, egress) are called out explicitly.

Two of those lines deserve to be derived rather than accepted, because they are the ones a reader would otherwise have to take on faith:

The mid-training line. Chapter 14.8 slices ~2B tokens into ~1.2B of anneal at seq_len=2048, ~0.6B of long-context extension at 8192, and ~0.2B of capability injection at 8192. Applying the attention correction per sub-phase: the 2048 portion costs \(6ND \times 1.31 = 9.6\times10^{17}\) hardware FLOPs, and the 8192 portion costs \(6ND \times 2.24 = 1.09\times10^{18}\) — the shorter sub-phase is the more expensive one, entirely because of the quadratic attention term. Sum ≈ \(2.05\times10^{18}\) FLOPs, and at the reduced \(u \approx 0.50\) you sustain once the micro-batch drops 4× to fit 8192-token sequences, that is 3.6 GPU-hours. The table’s number was never a guess.

The ladder line. Chapter 14.5 prices the four-rung ladder at 4.11 GPU-hours and the six-candidate data-mixture screen at a further 0.80 — 4.9 GPU-hours together, or ~14–16% of the flagship’s compute. That is the cheapest insurance in the project: it is what told you 20B tokens was the right budget, and 70/15/10/5 the right mix, before you spent the other ~86%.

Four things this table teaches that a bare “USD 98” hides:

1. Pretraining is ~48% of the pre-tax bill and over-training is ~92% of that. A Chinchilla-optimal run (1.7B tokens) would have cost ~2.0 GPU-hours, ~USD 3.7. We deliberately spent the other ~USD 42 to buy a permanently cheaper-to-serve model. That is the deployment-economics trade made visible.
2. Alignment + agent is cheap; the teacher API is the surprise. All of SFT+DPO+GRPO+distill is ~USD 19, and USD 8 of that is not your GPU at all — it is API calls to a large teacher to generate ReAct trajectories you then filter and distill. At 1B and beyond, teacher/data-generation cost often exceeds your own training cost (§14.12.5 puts a number on it).
3. The reality tax is real, and 25% is the optimistic version. No first run of a 30-layer model at high LR survives cleanly. You will hit an OOM from a mis-set gradient-accumulation count, a loss spike from an un-clipped attention logit, a corrupted shard. A practitioner who has done this before pays ~25%; a first-timer pays closer to 50%, which pushes the same bill to ~USD 115. That gap — not the FLOPs — is why the sticker says “~USD 100.”
4. The bill is a function of three inputs, only one of which is the model. Tokens, MFU, and USD/GPU-hr. Two of the three are yours to control; the third is a market.

14.12.4 Reproducibility: Making the Run Rebuildable¶

A result you cannot reproduce is an anecdote, not an experiment. At 100M on one GPU, reproducibility is achievable to the bit if you are disciplined; at 1B across a cluster it becomes merely approximate (non-deterministic all-reduce ordering, hardware variance, cuBLAS split-k) — which is all the more reason to nail the discipline now while you can verify it.

The checklist has six pillars: seeds, config, data, environment, checkpoints, and release. We give real code for each — it lives in stacklm/repro.py and is called from stacklm.train — and, just as importantly, we name the open-source tool that does each job properly at scale, because hand-rolling all six is a teaching device, not a recommendation.

Seeds and determinism¶

# stacklm/repro.py  (part 1: determinism)
import os, random
import numpy as np
import torch


def seed_everything(seed: int = 1337, deterministic: bool = True) -> None:
    """Seed every RNG the pipeline touches. Returns nothing; mutates global state.

    We seed FOUR independent generators: Python's `random`, NumPy (used by the
    tokenizer and data sampler), the CPU torch RNG, and all CUDA device RNGs.
    Missing any one silently reintroduces nondeterminism (e.g. dropout masks,
    data-order shuffles, or Muon's Newton-Schulz init).
    """
    os.environ["PYTHONHASHSEED"] = str(seed)   # affects dict/set iteration order
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        # Trade a little throughput for reproducible kernels. cuBLAS needs the
        # workspace env var set BEFORE the first CUDA context is created.
        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
        torch.use_deterministic_algorithms(True, warn_only=True)
        torch.backends.cudnn.benchmark = False   # disable autotuner (nondeterministic pick)


def rng_state_dict() -> dict:
    """Snapshot every RNG so a resumed run continues the SAME random stream.

    NOTE the numpy conversion: `np.random.get_state()` returns a tuple containing
    a raw ndarray, which `torch.load(weights_only=True)` will refuse. We store it
    as a torch tensor so the whole checkpoint stays loadable under the SAFE
    unpickler (see `load_checkpoint` below). This is not fussiness -- it is
    the difference between a checkpoint you can hand to a stranger and one you
    cannot.
    """
    kind, keys, pos, has_gauss, cached = np.random.get_state()
    return {
        "python": random.getstate(),
        "numpy": (kind, torch.from_numpy(keys.copy()), int(pos),
                  int(has_gauss), float(cached)),
        "torch": torch.get_rng_state(),
        "cuda": torch.cuda.get_rng_state_all(),
    }


def load_rng_state(state: dict) -> None:
    random.setstate(state["python"])
    kind, keys, pos, has_gauss, cached = state["numpy"]
    np.random.set_state((kind, keys.numpy().astype(np.uint32), pos, has_gauss, cached))
    torch.set_rng_state(state["torch"])
    if torch.cuda.is_available():
        torch.cuda.set_rng_state_all(state["cuda"])

The subtle one is RNG state (not just the seed). If you resume a run and re-seed from the same integer, you replay the random stream from the start — reusing data batches you already consumed. You must save and restore the RNG state at the resume step, exactly as in Checkpointing, Fault Tolerance & Long-Running Jobs.

Config: one frozen object, one hash¶

Every knob that affects the result — architecture, optimizer, schedule, data mix — lives in one frozen dataclass. Hash its canonical serialization and stamp the hash onto every checkpoint and log line. Two runs with the same config hash are the same experiment; a changed hash is a new one.

# stacklm/repro.py  (part 2: config + data + env provenance)
import hashlib, json, subprocess
from dataclasses import asdict, is_dataclass


def _canonical_json(obj) -> str:
    """Deterministic JSON: sorted keys, no whitespace jitter. Dataclasses -> dict."""
    if is_dataclass(obj):
        obj = asdict(obj)
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), default=str)


def config_hash(config) -> str:
    """Short, stable fingerprint of the full run configuration."""
    blob = _canonical_json(config).encode("utf-8")
    return hashlib.sha256(blob).hexdigest()[:12]   # 12 hex chars disambiguates plenty

The libraries that do this properly. A frozen dataclass is the right mental model; at more than a handful of experiments you want a config framework. Hydra (with OmegaConf) gives you composable YAML groups plus command-line overrides and writes the fully-resolved config into each run directory — hash that resolved file, not your defaults. pydantic-settings is the lighter-weight option when you want validation and environment-variable binding without Hydra’s multirun machinery. Both preserve the invariant that matters: exactly one serialized object determines the run.

And log the hash somewhere you will look. Print it as line 1, then push it to an experiment tracker so the loss curve and the config travel together:

# stacklm/tracking.py -- one place where the config hash meets the loss curve.
def start_tracking(config, backend: str = "wandb"):
    """Open a run in an experiment tracker keyed by the config hash.

    Using the config hash as the run id is the trick that makes this useful:
    a resumed job reattaches to the SAME run instead of forking a new curve,
    and two runs with identical configs collide loudly instead of silently
    becoming 'experiment 47' and 'experiment 63'.
    """
    h = config_hash(config)
    if backend == "wandb":
        import wandb
        return wandb.init(project="stacklm-100m", id=h, resume="allow",
                          config=asdict(config), tags=[f"cfg:{h}"])
    if backend == "mlflow":
        import mlflow
        mlflow.set_experiment("stacklm-100m")
        run = mlflow.start_run(run_name=h)
        mlflow.log_params(asdict(config))
        return run
    raise ValueError(backend)   # 'aim' and plain TensorBoard are equally fine

Weights & Biases and MLflow are the two defaults; Aim is a good self-hosted alternative and plain TensorBoard (torch.utils.tensorboard) is entirely sufficient for a single-GPU capstone. What matters is not which one, but that the config hash, the git commit, the corpus hash, and the loss curve are queryable together six months later.

The data manifest¶

The dirtiest source of “it doesn’t reproduce” is data drift: a shard silently re-tokenized, a file truncated by a failed download, the mix reweighted. A data manifest pins the content of every shard by SHA-256, not its path, so you can prove the bytes are identical.

def data_manifest(shard_paths: list[str]) -> dict:
    """Content-addressed manifest of the tokenized corpus.

    For each .bin memmap shard we record its SHA-256, byte size, and token count
    (uint16 tokens => 2 bytes each, valid because vocab_size=32768 <= 65536). The
    top-level `corpus_hash` fingerprints the WHOLE dataset: change any shard, or
    the order, and it changes.
    """
    shards, running = [], hashlib.sha256()
    for path in sorted(shard_paths):                 # sorted => order-independent of FS listing
        h = hashlib.sha256()
        size = 0
        with open(path, "rb") as f:
            for chunk in iter(lambda: f.read(1 << 20), b""):  # 1 MB chunks
                h.update(chunk)
                size += len(chunk)
        digest = h.hexdigest()
        running.update(digest.encode())
        shards.append({"path": path, "sha256": digest,
                       "bytes": size, "tokens": size // 2})
    return {
        "corpus_hash": running.hexdigest()[:16],
        "n_shards": len(shards),
        "total_tokens": sum(s["tokens"] for s in shards),
        "shards": shards,
    }

At capstone scale (~200 GB of shards) this file is your data version control. Past that, use a tool: DVC tracks large files by content hash in git without putting the bytes in git, git-lfs does the crude version of the same, and if your corpus lives on the Hub, datasets already computes a deterministic fingerprint for every transformed dataset (ds._fingerprint) plus a per-file sha256 in the repo metadata — record those and you get the same guarantee for free. The invariant to preserve is: the corpus is identified by its content, never by a path or a date.

Environment pinning¶

Same code + same data + different PyTorch build can still diverge (a changed default in a fused kernel, a new cuDNN heuristic, a different Triton autotune result). Capture the exact environment as data, alongside the run — and capture the lockfile, not a prose reminder to make one.

def _sha256_file(path: str, n: int = 16) -> str | None:
    try:
        with open(path, "rb") as f:
            return hashlib.sha256(f.read()).hexdigest()[:n]
    except OSError:
        return None


def environment_fingerprint(lockfile: str = "uv.lock") -> dict:
    """Everything outside your code that can change the result."""
    def _git(*args):
        try:
            return subprocess.check_output(["git", *args], text=True).strip()
        except Exception:
            return None
    return {
        "python": __import__("platform").python_version(),
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "cudnn": torch.backends.cudnn.version(),
        "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",
        "git_commit": _git("rev-parse", "HEAD"),
        "git_dirty": bool(_git("status", "--porcelain")),   # True => uncommitted: DANGER
        # The dependency set, captured as DATA rather than as a comment:
        "lockfile": lockfile,
        "lockfile_sha256": _sha256_file(lockfile),
        # Set by the launcher, e.g. `docker inspect --format='{{index .RepoDigests 0}}'`
        # or an Apptainer/Singularity image hash. The only true hermetic pin.
        "container_digest": os.environ.get("STACKLM_IMAGE_DIGEST"),
    }

A git_dirty=True at training time is a red flag: it means the code that produced the checkpoint was never committed, so the run is not reproducible from any commit. Fail loudly on it for the flagship run. In 2026 the ergonomic choice for the lockfile is uv (uv lock → uv sync --frozen), which resolves and installs a fully-pinned environment in seconds; poetry.lock and a pip freeze requirements file are equivalent in kind, weaker in guarantee. And the only hermetic pin is a container digest — a Docker or Apptainer image referenced by sha256: rather than by tag, since tags are mutable and pytorch:latest is a moving target that will silently change your kernels.

Checkpoint hygiene, and the format that does not execute code¶

Weld provenance into the checkpoint so a checkpoint directory is self-describing: not just weights but the config hash, the corpus hash, the git commit, the RNG state, and the step.

There is a security dimension here that the naive torch.save/torch.load pattern gets wrong. torch.load deserializes Python pickles, which can execute arbitrary code; since PyTorch 2.6 the default is weights_only=True, and passing weights_only=False opts out of that protection. A checkpoint you downloaded is untrusted input. The fix is to split the checkpoint: weights in safetensors (a zero-copy, mmap-able, executes-nothing tensor container), trainer state in a weights_only-safe torch blob, and provenance in plain JSON you can read without importing anything.

# stacklm/repro.py  (part 3: self-describing, safe-by-default checkpoints)
import shutil
from pathlib import Path
from safetensors.torch import save_file, load_file


def build_provenance(config, shard_paths: list[str]) -> dict:
    """Assemble the full provenance record stamped into every checkpoint & log."""
    return {
        "config_hash": config_hash(config),
        "config": asdict(config) if is_dataclass(config) else dict(config),
        "data_manifest": data_manifest(shard_paths),
        "env": environment_fingerprint(),
    }


def weights_for_safetensors(model) -> dict:
    """Model weights, de-aliased, ready for `safetensors.save_file`.

    safetensors REFUSES to serialize two keys that alias the same storage --
    `RuntimeError: Some tensors share memory ...`. Stack-100M ties
    `lm_head.weight` to `tok_emb.weight` (16.8M params, 16.6% of the model), so
    a naive `model.state_dict()` trips exactly that guard.

    Two traps worth naming, because they hide the bug until the worst moment:
      * `.contiguous()` does NOT break aliasing. It returns `self` when the
        tensor is already contiguous, so both keys keep the same `data_ptr`.
      * `.cpu()` DOES break it -- but only on a CUDA tensor, where it allocates
        a fresh host copy per key. So the naive version silently writes 33 MB of
        duplicate embedding on the GPU run and raises on the CPU toy/CI run.

    The correct fix is the one HF `save_pretrained` uses via `_tied_weights_keys`:
    drop the duplicate key and re-create the tie on load. Ch. 14.11 §7.4 leans on
    the same refusal as a free correctness check on the quantized export.
    """
    sd = model.state_dict()
    if getattr(model.cfg, "tie_embeddings", False):
        sd = {k: v for k, v in sd.items() if k != "lm_head.weight"}
    out = {k: v.detach().cpu().contiguous() for k, v in sd.items()}
    by_ptr: dict[int, list[str]] = {}
    for k, v in out.items():
        by_ptr.setdefault(v.data_ptr(), []).append(k)
    dupes = [ks for ks in by_ptr.values() if len(ks) > 1]
    assert not dupes, f"unexpected aliased tensors would break safetensors: {dupes}"
    return out


def save_checkpoint(path, model, optimizer, step, provenance, dataloader=None):
    """Write a self-describing checkpoint DIRECTORY:

        ckpt_step12000/
        |- model.safetensors   weights only; no pickle, mmap-able, framework-neutral
        |- trainer.pt          optimizer + RNG + dataloader state (weights_only-safe)
        |- provenance.json     config, config_hash, corpus_hash, env, git commit

    Written to a .tmp sibling and atomically renamed, so a crash mid-write can
    never leave a half-written 'latest' behind (checkpoint hygiene rule #1).
    """
    path = Path(path)
    tmp = path.with_suffix(".tmp")
    if tmp.exists():
        shutil.rmtree(tmp)
    tmp.mkdir(parents=True)

    # 1. Weights, with the tied duplicate dropped (re-tied by load_checkpoint).
    save_file(weights_for_safetensors(model), tmp / "model.safetensors",
              metadata={"step": str(step), "config_hash": provenance["config_hash"],
                        "tied_keys": "lm_head.weight"})

    # 2. Trainer state. BOTH the Muon and the AdamW param-group states live here.
    trainer = {"step": step, "optimizer": optimizer.state_dict(), "rng": rng_state_dict()}
    if dataloader is not None and hasattr(dataloader, "state_dict"):
        trainer["dataloader"] = dataloader.state_dict()      # StatefulDataLoader
    torch.save(trainer, tmp / "trainer.pt")

    # 3. Provenance, as human-readable JSON.
    (tmp / "provenance.json").write_text(json.dumps(provenance, indent=2, default=str))

    if path.exists():
        shutil.rmtree(path)
    tmp.rename(path)                                          # atomic on POSIX


def load_checkpoint(path, model, optimizer, expect_config_hash=None, dataloader=None):
    """Load a checkpoint directory, refusing a config mismatch and never unpickling
    anything the safe loader would not accept."""
    path = Path(path)
    provenance = json.loads((path / "provenance.json").read_text())
    got = provenance["config_hash"]
    if expect_config_hash is not None and got != expect_config_hash:
        raise RuntimeError(
            f"Config hash mismatch: checkpoint={got} expected={expect_config_hash}. "
            "You are resuming a run with a DIFFERENT architecture/optimizer config."
        )
    # strict=False because `lm_head.weight` was intentionally not serialized.
    missing, unexpected = model.load_state_dict(load_file(path / "model.safetensors"),
                                                strict=False)
    if getattr(model.cfg, "tie_embeddings", False):
        model.lm_head.weight = model.tok_emb.weight     # re-create the tie, explicitly
        missing = [k for k in missing if k != "lm_head.weight"]
    assert not missing and not unexpected, (missing, unexpected)

    trainer = torch.load(path / "trainer.pt", map_location="cpu", weights_only=True)
    optimizer.load_state_dict(trainer["optimizer"])
    load_rng_state(trainer["rng"])
    if dataloader is not None and "dataloader" in trainer:
        dataloader.load_state_dict(trainer["dataloader"])
    return trainer["step"]

Release: the artifact other people can use¶

The pillar most tutorials skip. A run is reproducible-in-principle when the six items above are pinned; it is reproducible-in-practice when someone else can pip install, download, and re-measure. Three concrete steps:

# 1. Publish weights + tokenizer + config, tagged with the provenance you already have.
python - <<'PY'
from huggingface_hub import HfApi
api = HfApi()
api.create_repo("your-org/stack-100m-base", repo_type="model", exist_ok=True)
api.upload_folder(folder_path="stacklm-100m/base",
                  repo_id="your-org/stack-100m-base",
                  commit_message="base @ cfg 4f2a9c1b8e30, corpus a17c93de5b204f61")
PY

# 2. Re-measure the headline evals against a PINNED harness commit, not `main`.
uv pip install "lm_eval @ git+https://github.com/EleutherAI/lm-evaluation-harness@<commit-sha>"
lm_eval --model hf \
        --model_args pretrained=./stacklm-100m/chat,dtype=bfloat16 \
        --tasks arc_easy,hellaswag,piqa --batch_size 16 \
        --output_path evals/results.json          # commit this file next to the weights

3. Write the model card, and mean it. README.md in the model repo carries the intended use and — the part that matters legally as much as scientifically — data provenance and licensing. FineWeb and FineWeb-Edu are released under ODC-By; Cosmopedia is synthetic and carries its generator’s terms; code from The Stack / StarCoder carries per-repository licenses and an opt-out mechanism you are expected to respect. Derived weights inherit obligations from all of them. State the mix, the token count, the eval numbers with the harness commit that produced them, and the honest capability ceiling from Chapter 14.11. A model card that omits the training mix is not a model card.

The one-page checklist¶

Before you press go on a flagship run, this must all be true. It is the difference between science and a story.

REPRODUCIBILITY PREFLIGHT  (stacklm)
[ ] seed_everything(1337, deterministic=True) called before model init
[ ] RNG *state* + dataloader state saved in checkpoint (not just the integer seed)
[ ] config frozen in one object (dataclass / Hydra-resolved YAML); hash on log line 1
[ ] tracker run opened with id=config_hash (W&B / MLflow / Aim / TensorBoard)
[ ] data_manifest.json written; corpus_hash matches the intended mix
[ ] held-out eval shard provably EXCLUDED from the manifest (contamination, Ch. 14.11)
[ ] git commit clean (git_dirty == False); commit hash logged
[ ] uv.lock committed AND its sha256 recorded; container digest recorded
[ ] checkpoints: safetensors weights (tied key dropped) + weights_only-safe trainer
    state + provenance.json
[ ] checkpoints atomic (tmp+rename), retention policy set (keep_last + milestones)
[ ] save -> load -> save round-trip is byte-identical, run on CPU (where the tied
    embedding actually aliases) not only on GPU
[ ] eval harness pinned to a commit sha; results.json committed beside the weights
[ ] a 50-step CPU toy run reproduces bit-identically across two invocations

That last line is the cheap insurance: a hermetic 50-step CPU run (tiny vocab, tiny model, synthetic data — the same toy path the book’s CI smoke-tests) that must produce identical loss across two invocations. If it does not, your determinism is broken and no amount of cluster time will fix it later.

14.12.5 What Breaks at 1B — and What to Change¶

The capstone’s proudest claim is that the entire 100M pipeline fits on one GPU. The instructive question is: run the exact same code targeting 1B parameters — what breaks first? We take the failures in the order you hit them.

Data: the wall you hit before the compute wall¶

A 1B model at Chinchilla-optimal wants ~20B tokens; over-trained for deployment (the capstone’s whole philosophy) it wants 200B–1T tokens. Our ~20B-token corpus is now 10–50× too small. You cannot fix this by looping the same data: repetition helps only up to ~4 epochs and its value collapses toward zero by ~16 epochs (Muennighoff et al., Scaling Data-Constrained Language Models, 2023). So the first thing that breaks at 1B is data volume, and the fix is more and cleaner data:

• Volume, with a real pipeline. Move from a FineWeb-Edu sample to the full FineWeb / FineWeb-Edu (trillions of tokens). At that scale you stop hand-rolling dedup and use the tool that actually built FineWeb: datatrove (HuggingFace), whose pipeline blocks — MinhashDedupSignature → MinhashDedupBuckets → MinhashDedupCluster → MinhashDedupFilter — run MinHash-LSH near-duplicate removal as a sharded, resumable, Slurm- or Ray-schedulable job rather than a laptop script. NVIDIA NeMo Curator (GPU-accelerated fuzzy dedup and classifier filtering) and AI2 dolma (the toolkit behind the Dolma corpus) are the credible alternatives. All three implement the same conceptual pipeline as Data Cleaning, Deduplication & Quality Filtering — the difference is that they survive a trillion documents.
• Quality and mix. The marginal token matters more at 1B; retune the mix (the 70/15/10/5 FineWeb-Edu / Cosmopedia / code / math split) using the domain-weighting methods of Data Mixing, Domain Weighting & Curriculum, and lean harder on synthetic data (Cosmopedia-style, see Synthetic Data for Pre- and Post-Training) to buy quality where raw web runs thin.

The uncomfortable truth: at 1B, data engineering — not model code — is where most of your time and a growing share of your budget goes. We will put a number on “growing share” two subsections from now.

Memory: the accounting the folklore gets wrong¶

Does 1B even fit on one 80 GB GPU? Do the arithmetic properly, because both the usual per-parameter rule and the usual conclusion are wrong for this project — and because the standard activation formula, applied naively, is wrong for our architecture by nearly 2×.

First error: 16 bytes/param assumes pure AdamW. Stack-100M uses the Muon+AdamW hybrid fixed in the plan — Muon for the 2D hidden matrices, AdamW for embeddings, norms, and 1D parameters. Muon keeps one momentum buffer, not Adam’s \(m\) and \(v\), so its groups cost 12 bytes/param, not 16.

Second error: weights+optimizer is not what fills the GPU. Activations are — and at Stack-100M’s 32,768-token vocabulary the single biggest activation is not in the transformer blocks at all, it is the loss head. Here is the accounting that actually predicts an OOM:

# stacklm/scaleup.py -- what really fills an 80 GB GPU?

def state_bytes_per_param(optimizer: str) -> int:
    """Steady-state bytes per parameter for mixed-precision training.

    Common to all: bf16 weights(2) + bf16 grads(2) + fp32 master copy(4) = 8 B.
      adamw      + fp32 m and v        -> +8  => 16 B/param
      muon       + ONE fp32 momentum   -> +4  => 12 B/param   <-- our 2D hidden matrices
      adamw8bit  + int8 m,v (bnb)      -> +2  => 10 B/param
      adamw_bf16 + bf16 m,v, no master -> ...  =>  8 B/param   (see Ch. 4.10)
    """
    return {"adamw": 16, "muon": 12, "adamw8bit": 10, "adamw_bf16": 8}[optimizer]


def optimizer_state_gb(groups: dict[str, float]) -> float:
    """groups maps optimizer name -> n_params in that group. Returns GB."""
    return sum(n * state_bytes_per_param(o) for o, n in groups.items()) / 1e9


KORTHIKANTI_COEFF = 34   # MHA + 4h GELU MLP + dropout, as published
STACK_COEFF       = 20   # Stack-100M's actual block shape (see docstring)


def block_activation_gb(n_layers: int, d_model: int, seq_len: int, micro_batch: int,
                        n_heads: int, coeff: float = STACK_COEFF,
                        flash: bool = True, recompute: bool = False) -> float:
    """Per-layer TRUNK activation memory, Korthikanti et al. (2022), 16-bit storage:

           bytes_per_layer = s*b*h * (coeff + 5*a*s/h)

    The 5*a*s/h term is the materialized s x s attention matrix. A fused
    FlashAttention kernel never writes it, so with `flash=True` it vanishes --
    which is the entire reason 8192-token mid-training fits at all.

    `coeff` is the paper's 34 ONLY for its reference block: multi-head attention
    (K and V as wide as Q), a 4h GELU MLP, and dropout masks. Stack-100M is none
    of those, so the coefficient drops:
      - no dropout in pretraining          -> -2 s*b*h  (two saved masks)
      - GQA 8:2, so K and V are h/4 each   -> -3 s*b*h  (K, V storage 4x smaller)
      - SwiGLU at 2.75h with torch.compile fusing silu(g)*u and recomputing it in
        backward -> the MLP costs ~13-16 s*b*h instead of the paper's 19
    That lands at roughly **16-24**, i.e. ~20 as a planning number -- which is
    exactly what reconciles this function with Ch. 14.7's measured 15-25 GB
    trunk budget for the same config. Pass 34 for a conservative upper bound.
    Full activation checkpointing stores only each layer's input: 2*s*b*h.
    """
    s, b, h, a = seq_len, micro_batch, d_model, n_heads
    if recompute:
        per_layer = 2 * s * b * h
    else:
        per_layer = s * b * h * (coeff + (0.0 if flash else 5 * a * s / h))
    return n_layers * per_layer / 1e9


def logits_gb(micro_batch: int, seq_len: int, vocab_size: int,
              chunk_size: int = 4096, chunked: bool = True) -> float:
    """The loss head -- the term Korthikanti's per-layer formula does not model,
    and the largest single allocation in the Stack-100M job.

    Naive (Ch. 14.4's forward): for every one of B*T tokens you hold V logits in
    bf16 (2 B) + an fp32 copy for cross-entropy (4 B) + the saved log_softmax
    output (4 B) + a second fp32 copy for the z-loss logsumexp (4 B) = 14 B.
    Chunked (Ch. 14.7's `chunked_lm_loss`): only `chunk_size` rows are live at
    once, holding fp32 logits + their gradient (8 B), and the logits are
    recomputed in backward rather than stored.
    """
    rows = chunk_size if chunked else micro_batch * seq_len
    return rows * vocab_size * (8 if chunked else 14) / 1e9


# --- Stack-100M at the flagship config (Ch. 14.7: micro-batch 32 x 2048, no recompute)
STACK_100M = {"muon": 84.5e6,     # 2D hidden matrices across the 30 blocks
              "adamw": 16.83e6}   # tied embedding + norms + 1D params
L, D, S, MB, A, V = 30, 512, 2048, 32, 8, 32768

state = optimizer_state_gb(STACK_100M)
trunk = block_activation_gb(L, D, S, MB, A)
print(f"weights + optimizer   : {state:6.2f} GB")
print(f"trunk, our shape      : {trunk:6.1f} GB   (coeff 20; Ch. 14.7 measures 15-25)")
print(f"  Korthikanti bound   : {block_activation_gb(L, D, S, MB, A, coeff=34):6.1f} GB")
print(f"  without FlashAttn   : {block_activation_gb(L, D, S, MB, A, flash=False):6.1f} GB (!!)")
print(f"  full recompute      : {block_activation_gb(L, D, S, MB, A, recompute=True):6.1f} GB")
print(f"loss head, naive      : {logits_gb(MB, S, V, chunked=False):6.1f} GB  <-- biggest single alloc")
print(f"loss head, chunked    : {logits_gb(MB, S, V):6.2f} GB")
print(f"PEAK as shipped       : {state + trunk + logits_gb(MB, S, V):6.1f} GB")
# weights + optimizer   :   1.28 GB
# trunk, our shape      :   20.1 GB   (coeff 20; Ch. 14.7 measures 15-25)
#   Korthikanti bound   :   34.2 GB
#   without FlashAttn   :  181.2 GB (!!)
#   full recompute      :    2.0 GB
# loss head, naive      :   30.1 GB  <-- biggest single alloc
# loss head, chunked    :   1.07 GB
# PEAK as shipped       :   22.5 GB

Read those numbers again, because three of them overturn folklore. Weights and optimizer state are 1.28 GB — under 6% of the memory in play. Without FlashAttention the run does not fit on any GPU that exists. And the naive loss head alone, 30.1 GB, is larger than the entire 30-block trunk: at \(B\cdot T = 65{,}536\) and \(V = 32{,}768\) the logits tensor is 64× wider than the residual stream. At a large vocabulary and a narrow \(d_{\text{model}}\), it is the loss head — not attention — that caps your batch size, which is why Chapter 14.7’s chunked loss head buys more memory than activation checkpointing does. Peak with both fixes in place is ~22.5 GB, matching Chapter 14.7’s measured budget, and that is why the flagship tier leaves recompute off at 80 GB and turns it on at 24 GB.

Now scale the same three functions to 1B and beyond:

for name, groups in [
    ("Stack-100M", {"muon": 84.5e6,  "adamw": 16.8e6}),
    ("1B  dense",  {"muon": 0.9e9,   "adamw": 0.1e9}),
    ("7B  dense",  {"muon": 6.3e9,   "adamw": 0.7e9}),
    ("70B dense",  {"muon": 63e9,    "adamw": 7e9}),
]:
    print(f"{name:11s} weights+opt = {optimizer_state_gb(groups):7.1f} GB")
# Stack-100M  weights+opt =     1.3 GB
# 1B  dense   weights+opt =    12.4 GB
# 7B  dense   weights+opt =    86.8 GB
# 70B dense   weights+opt =   868.0 GB

# ...and the activations that decide whether it ACTUALLY fits (1B: L=24, d=2048, a=16)
head = logits_gb(0, 0, 32768)                     # chunked: batch-independent, 1.07 GB
for mb in (4, 8, 16, 32):
    lo = block_activation_gb(24, 2048, 2048, mb, 16)                  # coeff 20
    hi = block_activation_gb(24, 2048, 2048, mb, 16, coeff=34)        # conservative
    print(f"1B micro-batch {mb:2d}: trunk {lo:5.1f}-{hi:5.1f} GB + 12.4 state + "
          f"{head:.1f} head = {lo+12.4+head:5.1f}-{hi+12.4+head:5.1f} GB")
# 1B micro-batch  4: trunk   8.1- 13.7 GB + 12.4 state + 1.1 head =  21.5- 27.2 GB
# 1B micro-batch  8: trunk  16.1- 27.4 GB + 12.4 state + 1.1 head =  29.6- 40.9 GB
# 1B micro-batch 16: trunk  32.2- 54.8 GB + 12.4 state + 1.1 head =  45.7- 68.2 GB
# 1B micro-batch 32: trunk  64.4-109.5 GB + 12.4 state + 1.1 head =  77.9-123.0 GB  <- OOM

So 1B fits a single 80 GB GPU, but the headroom is activations, not weights — micro-batch 16 is comfortable under either coefficient, and 32 only survives at the optimistic end, which is not a bet to take on a multi-day run. The naive “12 GB of state, 68 GB to spare” reading is off by a factor of four in the thing that matters. Note also how the shape shifts: at 1B the naive loss head at micro-batch 16 is \(32{,}768 \times 32{,}768 \times 14\,\text{B} \approx 15\) GB against a 32–55 GB trunk, so the head is no longer the dominant term. The loss head dominating is a small-and-narrow, large-vocab phenomenon — a Stack-100M problem that mostly solves itself as \(d_{\text{model}}\) grows.

Note also what happened to the folklore threshold: 7B needs ~87 GB of weights+optimizer state, so it does not fit an 80 GB A100 or H100 — but it does fit an H200 (141 GB), a B200 (192 GB), or an MI300X (192 GB) with room for a real batch. “When do I need FSDP?” is a hardware-generation question, not a model-size question. The rule that survives is the ratio: shard when weights + optimizer + activations exceed your device memory, and re-derive it every time you change tiers.

Levers, in the order you should reach for them, all covered in Memory-Efficient Training: Checkpointing, Offloading & LoRA Math: chunk the loss head (free, and biggest at small \(d\)), raise gradient accumulation (free), turn on activation checkpointing (33% FLOP surcharge), switch AdamW groups to 8-bit (bitsandbytes.optim.AdamW8bit, ~6 GB saved at 7B), keep optimizer states in bf16, then — only then — shard.

Time, and the parallelism ladder (with the libraries you would actually use)¶

The recipe does not break on memory at 1B. It breaks on time: 1B over 200B tokens is \(6 \times 10^{9} \times 2\times10^{11} = 1.2\times10^{21}\) FLOPs, ~99× the capstone pretrain, i.e. thousands of GPU-hours on one card. That is when you go multi-GPU, and the progression is exactly the one taught in Part III — but the libraries have moved, and naming the 2019 ones is no longer sufficient.

1. DDP / data parallelism first. Replicate the model on each of \(g\) GPUs, split the batch, all-reduce gradients. Near-linear speedup while the model fits per GPU (true at 1B). torch.nn.parallel.DistributedDataParallel launched with torchrun. This is your first move and it turns thousands of GPU-hours of wall-clock into hundreds.

2. FSDP2 when it stops fitting — and note it is fully_shard, not the old wrapper. Shard parameters, gradients, and optimizer state across data-parallel ranks so each holds \(1/g\) of the per-parameter bytes (ZeRO-3 in DeepSpeed’s vocabulary). The current PyTorch API is torch.distributed.fsdp.fully_shard (“FSDP2”), which shards per parameter using DTensor and composes cleanly with torch.compile and tensor parallelism; the older FullyShardedDataParallel module wrapper is the one every stale tutorial shows and the one you should not start from. At 1B you can use FSDP2 purely to buy activation headroom; you do not need it.

   # 100M -> 1B: the distributed diff really is about this small.
   import torch
   from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
   
   mp = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)
   for block in model.blocks:        # shard each transformer block: overlaps
       fully_shard(block, mp_policy=mp)   # all-gather(i+1) with compute(i)
   fully_shard(model, mp_policy=mp)  # then the root, for embeddings/head
   

3. Save with DCP, not torch.save. A sharded run’s state_dict is spread across ranks, and a plain save bakes in the world size — so resuming on 16 GPUs from an 8-GPU checkpoint fails. torch.distributed.checkpoint writes a resharding-aware, parallel checkpoint:

   import torch.distributed.checkpoint as dcp
   from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
   
   msd, osd = get_state_dict(model, optimizer)
   dcp.save({"model": msd, "optim": osd}, checkpoint_id=f"ckpt/step{step}")
   # ...later, on a DIFFERENT number of ranks:
   dcp.load({"model": msd, "optim": osd}, checkpoint_id=f"ckpt/step{step}")
   set_state_dict(model, optimizer, model_state_dict=msd, optim_state_dict=osd)
   

4. Tensor + pipeline parallelism only when a layer is too big or too slow (tens of B+). Split individual matmuls across GPUs (tensor) and stages of layers across GPUs (pipeline), the Megatron-style approach of Distributed Training II and Megatron-LM, DeepSpeed & Parallelism in Practice. You do not reach for this at 1B — pipeline bubbles and tensor-parallel communication are pure overhead you only accept when nothing else works. Naming it here is the point: knowing you don’t need it at 1B is as valuable as knowing you will at 30B.

What does 1B actually cost?¶

A cost chapter that spends a third of its length on “the path to 1B” owes you the number. Run the same machinery as §14.12.3 on a 1B model (24 layers, \(d=2048\), so the attention correction is \(1 + Lsd/N = 1 + 24\times2048\times2048/10^{9} = 1.10\), a 10% adder rather than our 31%) over 200B tokens. Hardware FLOPs \(= 1.2\times10^{21} \times 1.10 = 1.32\times10^{21}\):

(The FP8 row assumes a realized 1.2–1.6× end-to-end speedup, not the 2× peak ratio — that gap is scaling-factor overhead and the non-GEMM work FP8 does not touch. Wall-clock figures already include ~10% for data-parallel communication and stragglers.)

Now the part that is genuinely different at 1B — the non-GPU column:

Two readings of that table, and the second is the thesis of the whole chapter:

• 1B is ~55× the cost, not ~98×. The FLOPs went up 98×, but three generations of hardware and a cheaper FLOP absorbed part of it. The compute wall moves under you; plan against current price-performance, never against the number in a two-year-old blog post.
• The non-GPU share went from 13% to 43%. At 100M, “cost” means GPU-hours and the data is basically free. At 1B, the dedup pipeline and the synthetic-data bill together (~USD 2,300) rival the entire pretraining run (~USD 1,850). This is the quantitative form of “at 1B, data engineering is where the budget goes” — and it is why the tooling recommendations in the data subsection are not incidental. Extrapolate the trend one more decade and you can see why frontier labs are data organizations that happen to own GPUs.

Learning rate and batch size: what to rescale¶

You cannot copy Stack-100M’s hyperparameters to 1B unchanged; width and batch both grew. Two rules, both from Learning Rate Schedules, Warmup, Batch Size & Hyperparameters:

• Learning rate vs. width. Optimal LR shrinks as the model widens. The principled tool is μP (Maximal Update Parametrization): tune LR on a small proxy, then transfer it to the large model with width-dependent scaling so the same tuned value stays optimal — the same “fit small, predict large” discipline as §14.12.2, applied to hyperparameters instead of loss, and well supported in torchtitan and Megatron-Core. A cheaper heuristic is \(\eta \propto 1/\sqrt{d_{\text{model}}}\); going from \(d{=}512\) to \(d{=}2048\) roughly halves the peak LR. The AdamW groups (embeddings, norms) and the Muon groups (2D hidden matrices) rescale differently — Muon’s orthogonalized update is naturally more scale-robust, one reason the hybrid was chosen (see Optimizers: SGD, Adam, Adafactor, Lion, Muon & Shampoo) — but you still retune.
• Batch size and warmup. Larger models have a larger critical batch size — the point past which more parallelism stops helping per-step progress. At 1B you can raise the global batch from ~0.5M tokens toward ~1–2M tokens (more data-parallel ranks make this nearly free) and extend warmup proportionally. The WSD schedule transfers cleanly: keep the long stable phase, keep the short decay phase as your mid-training quality anneal — the schedule’s shape is scale-invariant even though the numbers move.

Keep MuonClip / QK-clip on. Attention-logit blow-ups get worse at scale and high LR, and QK-clip is exactly the stability fix (Kimi K2, Moonshot AI, 2025) that made Muon usable at scale — do not drop it when you can least afford instability. See Training Stability, Loss Spikes & Debugging Large Runs.

The MoE fork in the road¶

At 1B dense you spend 1B params of compute per token. Mixture-of-Experts (MoE) breaks that link: route each token to a few experts so total params (capacity) far exceed active params (compute/token). This is the single highest-leverage architecture change on the path beyond 1B, and the capstone’s own plan lists it as the scale-up option.

The modern design point is fine-grained experts with shared experts, from DeepSeekMoE (Dai et al., 2024): split each expert into many smaller ones for finer specialization, and keep a few always-on shared experts to absorb common patterns so the routed experts can specialize. Qwen3-MoE (Qwen Team, 2025) is a recent production instance of the same family. The trade you are making:

A sketch of the fine-grained + shared-expert MoE FFN that would replace Stack-100M’s SwiGLU block — the full treatment (load-balancing loss, capacity factor, expert/all-to-all parallelism) is in Mixture-of-Experts (MoE) Architectures:

# stacklm/moe.py  -- DeepSeekMoE-style FFN: many fine-grained + few shared experts
import torch, torch.nn as nn, torch.nn.functional as F


class SwiGLUExpert(nn.Module):
    """One small SwiGLU expert (same activation as the dense Stack-100M MLP)."""
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up   = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class DeepSeekMoEFFN(nn.Module):
    """Fine-grained routed experts + always-on shared experts (Dai et al., 2024).

    - `n_routed` many small experts; top-`k` are activated per token.
    - `n_shared` experts run for EVERY token (capture common structure).
    - `router_bias` is DeepSeek-V3's auxiliary-loss-FREE load balancer: a
      per-expert bias added to the routing logits for SELECTION only (never for
      the gate value), nudged up for under-loaded experts and down for
      over-loaded ones between steps. It balances load without an auxiliary
      loss term fighting the language-modeling objective.

    Active params/token = shared + k*routed, far below the total capacity of
    (n_shared + n_routed) experts. Here: active 3 of 17 experts resident.
    """
    def __init__(self, d_model=512, d_ff=352, n_routed=16, n_shared=1, k=2):
        super().__init__()
        self.k = k
        self.router = nn.Linear(d_model, n_routed, bias=False)   # token -> expert affinities
        self.register_buffer("router_bias", torch.zeros(n_routed))  # updated by the balancer
        self.routed = nn.ModuleList([SwiGLUExpert(d_model, d_ff) for _ in range(n_routed)])
        self.shared = nn.ModuleList([SwiGLUExpert(d_model, d_ff) for _ in range(n_shared)])

    def forward(self, x):                        # x: (B, T, d_model)
        B, T, D = x.shape
        flat = x.reshape(B * T, D)
        out = torch.zeros_like(flat)
        for e in self.shared:                    # shared experts: always on
            out = out + e(flat)

        scores = self.router(flat)                       # (B*T, n_routed)
        _, topi = (scores + self.router_bias).topk(self.k, dim=-1)   # SELECT with bias
        gates = F.softmax(scores.gather(-1, topi), dim=-1)           # WEIGHT without it
        for slot in range(self.k):               # accumulate the k routed experts
            idx = topi[:, slot]                  # which expert each token chose
            g = gates[:, slot:slot + 1]
            for e_id, expert in enumerate(self.routed):
                mask = idx == e_id
                if mask.any():                   # teaching sketch; see the note below
                    out = out.index_put((mask.nonzero(as_tuple=True)[0],),
                                        g[mask] * expert(flat[mask]), accumulate=True)
        return out.reshape(B, T, D)

With those numbers, active compute per token is n_shared + k = 1 + 2 = 3 experts (3 × 352 ≈ 1056 FFN width, below the dense 1408), while total resident capacity is n_shared + n_routed = 17 experts — roughly 6× the capacity per active expert for less-than-dense compute. That is the whole pitch of MoE in one line, and Exercise 3 does the parameter accounting exactly.

The catch — and why MoE is a fork, not a free lunch: MoE trades compute for memory and communication. All experts must live in memory even though each token uses few, and multi-GPU MoE needs all-to-all communication to route tokens to the GPUs holding their experts (expert parallelism, in Distributed Training II; the serving-side consequences are in Serving Mixture-of-Experts). At 1B on one GPU, dense is simpler and probably right. MoE is the lever you pull when you want 7B-worth of capacity at ~1–2B-worth of inference cost — the DeepSeek/Qwen3 playbook.

14.12.6 Landing the Plane¶

You have now built the whole stack once — small, real, and end to end. Not a tutorial that stops at “and the rest is left as an exercise,” but every stage a frontier lab runs, executed in miniature on hardware you can rent for an afternoon. The tokenizer is yours; the scaling law is one you fit and then checked against the outcome; the optimizer, the schedule, the alignment stack, the agent, the quantized deployment — you touched all of it, and you know where every dollar went and why.

The gap between Stack-100M and a frontier model is real and enormous, and this book has been honest about it at every step. But the shape of the machine is the same at every scale, and the path to 1B is mostly data engineering and parallelism discipline, not new ideas. Someone who has built the machine once at 100M is far better equipped to reason about a 100B run than someone who has only read about one. That transfer — from a run you can hold in your head to systems you cannot — is the entire point of the capstone.

Go run it. Then over-train it. Then, when you are ready, scale it up.

14.12.7 Further Reading: The Works Behind Part XIV¶

Every technique in Stack-100M traces to a real, load-bearing paper. This is the curated list of the works actually cited across the capstone — read these and you have read the sources of the modern small-model recipe. (Illustrative magnitudes throughout the capstone are “on the order of”; these citations are the verifiable ground truth.)

Small-model architecture & the deep-thin recipe

• Liu et al. MobileLLM: Optimizing Sub-billion Parameter Language Models (2024) — the deep-and-thin insight at fixed parameter budget.
• Su et al. RoFormer: Rotary Position Embedding (RoPE, 2021); Kazemnejad et al. The Impact of Positional Encoding on Length Generalization (NoPE, 2023); HuggingFace SmolLM3 (2025) — RoPE + NoPE-on-a-subset for length generalization.
• Ainslie et al. GQA: Grouped-Query Attention (2023); Shazeer GLU Variants Improve Transformer (SwiGLU, 2020); Zhang & Sennrich Root Mean Square Layer Normalization (RMSNorm, 2019); Press & Wolf Using the Output Embedding to Improve Language Models (tied embeddings, 2017).
• DeepSeek-AI DeepSeek-V2 (MLA, 2024) and DeepSeek-V3 (MTP, FP8 training, auxiliary-loss-free load balancing, 2024); Gloeckle et al. Better & Faster Large Language Models via Multi-token Prediction (2024) — the efficiency options.

Data, scaling, and over-training

• Penedo et al. The FineWeb Datasets (FineWeb / FineWeb-Edu, HuggingFace, 2024); HuggingFace Cosmopedia — the data recipe, and (via datatrove) the pipeline that produced it.
• Hoffmann et al. Training Compute-Optimal Large Language Models (Chinchilla, 2022); Kaplan et al. Scaling Laws for Neural Language Models (2020) — the scaling-law foundation.
• Besiroglu et al. Chinchilla Scaling: A Replication Attempt (2024) — why the fitted constants are fragile and the extrapolated loss is the thing to trust, the basis for §14.12.2’s error bars.
• Sardana et al. Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Laws (2024) — the over-training economics that justify 200 tokens/param.
• Muennighoff et al. Scaling Data-Constrained Language Models (2023) — the data wall and epoch limits that bite at 1B.

Optimizer, schedule, and stability

• Jordan et al. Muon: An Optimizer for Hidden Layers in Neural Networks (2024) and Moonshot AI Kimi K2 (MuonClip / QK-clip, 2025) — the Muon+AdamW hybrid and its stability fix.
• Hu et al. MiniCPM (WSD schedule, 2024); OLMo team OLMo 2 (2024) — WSD and the mid-training phase.

Systems, memory, and cost

• Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models (2022) — the activation-memory formula behind §14.12.5.
• Dao et al. FlashAttention (2022) and FlashAttention-2 (2023) — why the \(s^2\) activation term disappears and 8192-token mid-training fits at all.
• Chowdhery et al. PaLM (2022) — the MFU definition that counts attention FLOPs, the convention §14.12.3 reconciles against.

Alignment, agents, and deployment

• Rafailov et al. Direct Preference Optimization (DPO, 2023); Shao et al. DeepSeekMath (GRPO, 2024); Yao et al. ReAct (2022) — SFT/DPO/GRPO and the agent loop.
• Frantar et al. GPTQ (2022); Lin et al. AWQ (2023) — the quantization behind the int4 laptop deployment.
• Dai et al. DeepSeekMoE (2024), Gale et al. MegaBlocks (2022), and Qwen Team Qwen3 Technical Report (2025) — the MoE path beyond 1B and the kernels that make it fast.

Distributed training (the scale-up ladder)

• Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2019); Shoeybi et al. Megatron-LM (2019); Narayanan et al. Efficient Large-Scale LM Training on GPU Clusters (2021) — the DP → FSDP → tensor/pipeline progression, and the achieved-vs-peak FLOPs reporting methodology this chapter’s MFU section extends.

For the annotated, book-wide version of this list see Key Papers: An Annotated Reading List.

Exercises¶

1. (Conceptual) The chapter calls training Stack-100M on ~20B tokens (~200 tokens/param) instead of the Chinchilla-optimal ~1.7B “economically irrational for a model you train and throw away, and completely rational for a model you will serve.” Explain the asymmetry that makes both halves of that sentence true, and name where in the cost table the over-training decision shows up.

2. (Quantitative) The cost.py example assumes the stable phase sustained 231,500 tokens/sec, which the chapter says gives 21.6 GPU-hours at MFU(6ND) 45.1% over 18B tokens. Suppose instead your kernels are less well tuned and the loop sustains only 200,000 tokens/sec on the same 18B-token run. Using the chapter’s formulas and constants (\(N = 1.014\times10^8\), A100 bf16 peak \(= 3.12\times10^{14}\) FLOP/s, USD 1.80/GPU-hr), compute (a) the pretrain GPU-hours, (b) MFU under the 6ND convention, © the true fraction of A100 peak being used at seq_len=2048, and (d) the pretrain dollar cost. By what fraction does the slower run raise the pretrain bill?

3. (Quantitative) Take the DeepSeekMoEFFN config from the chapter: d_model=512, d_ff=352, n_routed=16, n_shared=1, k=2. Each SwiGLUExpert has three bias-free linear layers (w_gate, w_up: \(d_{\text{model}}\times d_{\text{ff}}\) each; w_down: \(d_{\text{ff}}\times d_{\text{model}}\)). Compute (a) parameters per expert, (b) total resident expert parameters, © active expert parameters per token, and (d) compare © against the dense Stack-100M SwiGLU MLP (intermediate=1408). Does the MoE block use less compute per token than the dense block? Then (e): if you replaced all 30 blocks’ MLPs this way, what happens to the whole model’s parameter count, and what does that do to the memory analysis in §14.12.5?

4. (Implementation + quantitative) The chapter argues that Korthikanti’s coefficient 34 is an upper bound for Stack-100M, and that the term the formula omits entirely — the loss head — is the largest allocation in the job. Make that concrete. (a) Extend stacklm/scaleup.py with a peak_memory_gb(cfg, micro_batch, coeff, chunked_head) helper that sums weights+optimizer, trunk activations, and the loss head, and use it to reproduce the flagship’s ~22.5 GB and the ~50 GB figure Ch. 14.7 reports before the head is chunked. (b) Using the same helper, find the largest micro-batch that keeps a 1B model (\(L{=}24\), \(h{=}2048\), \(a{=}16\), \(s{=}2048\), \(V{=}32{,}768\), 12.4 GB of state) under 80 GB at both coeff=20 and coeff=34, and say which number you would actually plan against. © Compute total pretraining FLOPs for 1B over 200B tokens including the attention correction, express it as a multiple of the capstone’s \(1.22\times10^{19}\), and estimate the single-GPU A100 wall-clock at MFU(6ND) 0.45.

# stacklm/scaleup.py  (append)
def peak_memory_gb(n_layers, d_model, seq_len, micro_batch, n_heads, vocab_size,
                   state_gb, coeff=STACK_COEFF, chunked_head=True,
                   flash=True, recompute=False) -> float:
    """Weights+optimizer + trunk activations + loss head, in GB.

    The three terms have completely different scaling: state is constant in
    batch, the trunk is linear in micro_batch, and the loss head is linear in
    micro_batch when unchunked and CONSTANT when chunked. That last fact is
    why chunking the head is what unlocks a large micro-batch.
    """
    trunk = block_activation_gb(n_layers, d_model, seq_len, micro_batch,
                                n_heads, coeff, flash, recompute)
    head = logits_gb(micro_batch, seq_len, vocab_size, chunked=chunked_head)
    return state_gb + trunk + head

S100 = dict(n_layers=30, d_model=512, seq_len=2048, n_heads=8,
            vocab_size=32768, state_gb=1.28)
print(peak_memory_gb(micro_batch=32, **S100))                      # 22.5 GB
print(peak_memory_gb(micro_batch=32, chunked_head=False, **S100))  # 51.5 GB

5. (Implementation) Extend stacklm/cost.py with an inference-side accounting to make the “worked example” reproducible in code. Implement serving_flops(n_params, n_requests, tokens_per_request=256) using the \(2ND\) inference rule, verify it reproduces the chapter’s \(\approx 5.2\times10^{19}\) FLOPs for a billion 256-token requests, and add breakeven_requests(...) returning how many 256-token requests it takes for cumulative serving FLOPs to equal the 20B-token pretraining budget. Report that number and explain why it does not depend on model size.

# stacklm/cost.py  (append)
def serving_flops(n_params: int, n_requests: int,
                  tokens_per_request: int = 256) -> float:
    """Inference compute via the 2ND rule: 2 FLOPs/param/generated token.

    Decode-only approximation: it ignores prefill (which is prompt-length
    dependent) and the attention-over-KV term, both of which push the real
    number UP -- so this is a lower bound on serving cost, which makes the
    over-training argument conservative.
    """
    return 2.0 * n_params * n_requests * tokens_per_request


def breakeven_requests(n_params: int = N_PARAMS,
                       pretrain_tokens: float = 2.0e10,
                       tokens_per_request: int = 256) -> float:
    """How many requests until cumulative serving FLOPs == pretraining FLOPs."""
    pretrain = training_flops(n_params, pretrain_tokens)        # 6ND
    per_request = 2.0 * n_params * tokens_per_request           # 2ND, one request
    return pretrain / per_request


if __name__ == "__main__":
    c = serving_flops(N_PARAMS, 1_000_000_000, 256)
    print(f"serve 1e9 reqs x256 tok: {c:.2e} FLOPs")   # -> 5.19e+19
    print(f"breakeven: {breakeven_requests():,.0f} requests")

6. (Implementation) Two checkpoint-hygiene tasks from §14.12.4. (a) Write a pytest test that would have caught the tied-embedding bug: build a tiny Stack100M on CPU with tie_embeddings=True, save and reload a checkpoint, and assert the tie survives the round-trip. Explain why the test must run on CPU. (b) Implement prune_checkpoints(ckpt_dir, keep_last=3, milestones=()) consistent with the directory-per-checkpoint layout: given a directory containing sub-directories named ckpt_step{N}, delete every checkpoint except the keep_last highest steps and any step in milestones; return the list of deleted names. Be careful about the .tmp sibling an in-flight atomic write leaves behind.

# tests/test_checkpoint_roundtrip.py
def test_tied_embedding_survives_roundtrip(tmp_path):
    cfg = StackConfig(vocab_size=64, d_model=32, n_layers=2,
                      n_heads=2, n_kv_heads=1, tie_embeddings=True)
    model = Stack100M(cfg)                       # CPU: the tie is a real alias
    assert model.lm_head.weight.data_ptr() == model.tok_emb.weight.data_ptr()

    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    prov = {"config_hash": config_hash(cfg)}
    save_checkpoint(tmp_path / "ckpt_step1", model, opt, 1, prov)

    # the duplicate key must NOT be on disk
    from safetensors import safe_open
    with safe_open(tmp_path / "ckpt_step1/model.safetensors", framework="pt") as f:
        assert "lm_head.weight" not in f.keys()

    fresh = Stack100M(cfg)
    load_checkpoint(tmp_path / "ckpt_step1", fresh, opt)
    assert fresh.lm_head.weight.data_ptr() == fresh.tok_emb.weight.data_ptr()
    torch.testing.assert_close(fresh.tok_emb.weight, model.tok_emb.weight)

# stacklm/repro.py  (part 4: retention policy)
import re, shutil
from pathlib import Path

_CKPT_RE = re.compile(r"^ckpt_step(\d+)$")     # anchored: excludes ckpt_step900.tmp


def prune_checkpoints(ckpt_dir: str, keep_last: int = 3,
                      milestones=()) -> list[str]:
    """Enforce checkpoint retention: keep the `keep_last` newest checkpoints
    plus any `milestones` steps (e.g. end of stable phase, end of decay);
    delete the rest. Returns the names removed.

    Prevents the 'fills a disk overnight' footgun: ~1.28 GB/checkpoint
    (84.5M x 12 B Muon + 16.8M x 16 B AdamW) every 500 steps is unsustainable.
    """
    root = Path(ckpt_dir)
    milestones = set(milestones)
    found = []                                   # (step, name)
    for entry in root.iterdir():
        m = _CKPT_RE.match(entry.name)
        if m and entry.is_dir():                 # dirs only: skips stray files
            found.append((int(m.group(1)), entry.name))
    found.sort()                                 # ascending by step

    newest = {step for step, _ in found[-keep_last:]} if keep_last > 0 else set()
    keep = newest | milestones

    removed = []
    for step, name in found:
        if step not in keep:
            shutil.rmtree(root / name)           # rmtree, not unlink: it's a dir
            removed.append(name)
    return removed

7. (Quantitative) Your cloud offers A100-80GB at USD 1.30/GPU-hr and H100-SXM at USD 2.50/GPU-hr. The 18B-token stable phase requires \(1.43\times10^{19}\) hardware FLOPs (6ND plus the 31% attention correction). Assume you sustain \(u = 0.59\) of peak on the A100 (312 TFLOP/s dense bf16) but only \(u = 0.45\) on the H100 (989 TFLOP/s dense bf16), because the smaller model leaves the bigger chip more launch-bound. Compute (a) GPU-hours and cost on each, (b) the speedup, and © the H100 rental price at which the two become equally expensive. What does this say about how to choose a tier?

8. (Conceptual + quantitative) In §14.12.2 the Ch. 14.5 ladder fit (\(E{=}2.45\), \(A{=}124\), \(\alpha{=}0.33\), \(B{=}234\), \(\beta{=}0.30\)) predicts \(L \approx 2.94\) nats/token for the flagship. Suppose your run instead lands at 3.35 nats/token on the held-out pretrain-mix shard. (a) Is that inside the fit’s honest band? (b) Decompose the predicted loss into its three terms and use the decomposition to argue which failure hypotheses the discrepancy is and is not consistent with. © Name three checks you would run, in order, before concluding the scaling law was wrong.