14.8 Mid-Training: Quality Annealing, Long-Context Extension & Capability Injection¶

By the end of Chapter 14.7 our Stack-100M model has consumed roughly 18 billion tokens during the long stable phase of a Warmup-Stable-Decay (WSD) schedule, at a constant learning rate, on the broad 70/15/10/5 FineWeb-Edu / Cosmopedia / code / math mix. The train loss has flattened somewhere in the ballpark of the low-3s nats/token, and we have a checkpoint — ckpt_stable.pt — that is unfinished on purpose. We deliberately did not decay the learning rate to zero, because the most valuable compute in the whole project is still in front of us.

This chapter spends it. Mid-training is the phase that sits between raw pretraining and post-training — a term the OLMo 2 report (Allen Institute for AI, 2024–2025) helped crystallize, though the idea is now near-universal across open models. We make three moves, all sharing one checkpoint and one continued-training loop:

1. WSD decay-phase annealing — run the short, sharp LR-decay phase on a higher-quality data mix (more Cosmopedia, instruction-flavored text, more math and code). Because the decay phase is where most of the committed loss reduction lands, upgrading the data here buys a large, cheap quality jump.
2. Long-context extension — continue training at seq_len = 8192 (a 4× jump from the pretrain 2048) using RoPE base rescaling (NTK-aware / YaRN-style), so the model can actually attend across a document instead of a paragraph.
3. Capability injection — concentrate math and code in the final sub-phase, so the narrow tool-using agent we build in Chapter 14.10 has an arithmetic-and-structure floor to stand on.

This is where “narrow but real” specialization begins. We are still doing self-supervised next-token prediction — no SFT, no chat template yet (that is Chapter 14.9). Mid-training changes what data the model sees and how far it can see, not what objective it optimizes.

Two of those three moves look like config changes and are not. Long-context extension in particular has two failure modes that are invisible until you check for them: your packed shards may contain no positions past 2048, in which case the RoPE rescale trains nothing; and the dense document mask that worked at 2048 does not fit at 8192, in which case you either OOM or silently fall off the FlashAttention path. This chapter cashes both of those out with code, because they are exactly the things that get hand-waved.

This chapter builds directly on Continual & Domain-Adaptive Pretraining, Learning Rate Schedules, Warmup, Batch Size & Hyperparameters, Long-Context Pretraining & Context Extension, Positional Encodings: Sinusoidal, Learned, RoPE & ALiBi, and Data Mixing, Domain Weighting & Curriculum. We cross-link them rather than re-derive; here we apply them to one concrete model.

Why a Separate Mid-Training Phase Exists¶

The instinct of a first-time trainer is to pick one data mix and one cosine schedule, train to convergence, and call it done. That is the Karpathy-nanoGPT recipe, and it is a fine baseline. But two facts about modern small models push us to split the run.

Fact 1: the decay phase is disproportionately valuable. In a WSD schedule (MiniCPM, Hu et al., 2024; used in spirit by DeepSeek’s multi-step decay), the learning rate is warmed up, held constant for the bulk of training (the stable phase), then decayed sharply over a short final window (the decay phase). Empirically — and this is the whole reason WSD exists — the model looks unconverged at the end of the stable phase and then drops steeply the moment you begin decaying. Loss that was stuck falls by a visible margin over the last few percent of tokens. The decay phase is where the model commits the representations it has been loosely exploring at high LR.

That timing is a gift. If the last few percent of tokens matter most, then the quality of those specific tokens matters most too. Feeding premium data precisely into the decay window is the highest-leverage data decision in the entire pipeline. This is the core insight of mid-training and the reason WSD and annealing are described together in Chapter 14.6.

Fact 2: quality and context length are cheap to change late, expensive to change early. Pretraining on 8192-token contexts from step zero would have quadrupled the attention FLOP term for all 20B tokens — and for a deep-and-thin model like Stack-100M that term is not a rounding error (we do the arithmetic later in this chapter: at 8192 it is more than half of all training FLOPs). That is a large tax paid over the entire run for a capability we only need at the end. Similarly, high-quality synthetic data (Cosmopedia v2) and curated math/code are scarcer and more expensive per token than filtered web text; spending 20B tokens of it would be wasteful when ~2B tokens of it, placed correctly, captures most of the benefit. Mid-training is the economically correct place to pay for both.

The mid-training budget¶

We slice a ~2B-token window off the ~20B-token budget for all three moves combined — about 10% of total tokens, and exactly the decay_frac = 0.10 leg that Ch. 14.7’s StackConfig reserved (3,815 of 38,147 steps). A representative split (illustrative; tune per Chapter 14.5):

That is 3,813 optimizer steps, which is Ch. 14.7’s 3,815-step decay leg up to floor division. The LR decays monotonically across all three sub-phases — mid-training is one continuous WSD decay, just with the data mix and sequence length changing underneath it. (The exact boundary multipliers, 0.2253 and 0.0513, fall out of the \(1-\sqrt{t}\) shape; Exercise 3 derives them.)

Move 1 — WSD Decay-Phase Annealing¶

The annealing mix¶

We shift the sampling weights toward denser, cleaner, more instruction-shaped text. The pretrain mix was tuned for breadth; the anneal mix is tuned for the model’s final impression.

The instruction-flavored slice is deliberately small and is not SFT: it is raw text that happens to be shaped like questions-and-answers or explanations, so the model warms to the register without us yet imposing a chat template or masking losses. Full instruction tuning is Chapter 14.9.

Building the mixture loader¶

The data plumbing — streaming shards, document-aware packing, the PackedMemmapDataset memmap reader — was all built in Chapter 14.2. Mid-training only re-weights the mixture, and the cleanest way to express a weighted mixture over per-source datasets in PyTorch is a ConcatDataset plus a WeightedRandomSampler: it samples at the sequence level, so one micro-batch can contain web, code, and math side by side (which decorrelates gradients across domains), and it needs no custom sampler code.

# capstone/stacklm/mid/mixture.py
"""Weighted source mixtures over the Ch. 14.2 packed shards.

Each source lives in its own shard directory, packed at the sequence length the
sub-phase will train at (`data/mid/<source>_<seq_len>/`). Sampling is per
*sequence*, not per micro-batch, so a single forward pass mixes domains.
"""
import torch
from torch.utils.data import ConcatDataset, DataLoader, WeightedRandomSampler

from stacklm.data import PackedMemmapDataset          # Ch. 14.2

# Sub-phase A: the annealing mix. Keys are shard-directory names written by the
# tokenize+pack pass (Ch. 14.2) and recorded in the data manifest.
ANNEAL_MIX = {
    "fineweb_edu":   0.40,
    "cosmopedia_v2": 0.30,
    "starcoder":     0.15,
    "finemath":      0.10,
    "instruct_flav": 0.05,
}

# Sub-phases B and C are defined in the same module; their sources are motivated
# in "Repacking for sub-phase B" and "Move 3" below.
LONGCTX_MIX = {
    "starcoder_repo":   0.35,
    "books_pg19":       0.25,
    "arxiv_proofpile2": 0.15,
    "fineweb_edu_long": 0.15,
    "cosmopedia_v2":    0.10,   # deliberately SHORT: the anti-drift anchor
}
CAPABILITY_MIX = {
    "finemath":         0.30,
    "starcoder_repo":   0.30,
    "cosmopedia_v2":    0.25,
    "fineweb_edu_long": 0.15,
}

for _m in (ANNEAL_MIX, LONGCTX_MIX, CAPABILITY_MIX):
    assert abs(sum(_m.values()) - 1.0) < 1e-9, "mixture weights must sum to 1"


def build_mixture_loader(mix: dict, seq_len: int, micro_bs: int, *,
                         root: str = "data/mid", seed: int = 1234,
                         num_workers: int = 4):
    """Return an INFINITE iterator of batch dicts drawn from `mix`.

    Each yielded dict has `input_ids`, `position_ids`, `seq_ids`, `targets`
    (shapes (micro_bs, seq_len - 1)) -- exactly what `Stack100M.forward` and the
    document-aware mask consume. All four are threaded into the model: dropping
    `position_ids` is a silent bug (see the training loop below).
    """
    datasets, weights = [], []
    for name, w in mix.items():
        d = PackedMemmapDataset(f"{root}/{name}_{seq_len}")
        # Guard against the single most common mid-training bug: reading shards
        # that were packed at the PRETRAIN length while claiming to train long.
        assert d.seq_len == seq_len, (
            f"{name} shards are packed at {d.seq_len}, not {seq_len}; "
            f"re-run the repack pass (see 'Repacking for sub-phase B')")
        datasets.append(d)
        # Per-sequence probability ∝ source weight, spread evenly inside a source.
        weights.extend([w / len(d)] * len(d))

    concat = ConcatDataset(datasets)
    g = torch.Generator().manual_seed(seed)
    sampler = WeightedRandomSampler(weights, num_samples=len(concat),
                                    replacement=True, generator=g)
    loader = DataLoader(concat, batch_size=micro_bs, sampler=sampler,
                        drop_last=True, num_workers=num_workers,
                        pin_memory=True, persistent_workers=num_workers > 0)

    def infinite():
        while True:
            yield from loader
    return infinite()

Resuming the WSD decay from the stable checkpoint¶

The learning-rate mechanics are the crux. During the stable phase the effective schedule multiplier was pinned at 1.0 (peak LR). Mid-training runs the decay leg: the multiplier goes 1.0 → 0 over the mid-training token budget, using a 1−sqrt shape (the MiniCPM default; sharper early, gentle tail). We compute it against the mid-training step count, not the global step count, because the stable phase already happened.

This function is not new code — it already ships in capstone/stacklm/optim/schedule.py from Chapter 14.6, and mid-training imports it rather than redefining it. It is reproduced here because its shape is the thing you need in your head for the rest of the chapter.

# capstone/stacklm/optim/schedule.py  (Ch. 14.6 — reproduced, not redefined)
import math

def wsd_decay_multiplier(mid_step: int, num_decay_steps: int,
                         final_frac: float = 0.0, shape: str = "1-sqrt") -> float:
    """LR multiplier for the WSD *decay* leg, indexed from the start of mid-training.

    mid_step        : steps taken *since* resuming from ckpt_stable (0-indexed).
    num_decay_steps : total mid-training steps (all three sub-phases together).
    final_frac      : LR floor as a fraction of peak (we use ~0.0; some use 0.1).
    shape           : "1-sqrt" (MiniCPM) or "linear" or "cosine".
    Returns a value in [final_frac, 1.0]; multiply by EACH GROUP's peak LR.
    """
    t = min(mid_step, num_decay_steps) / max(1, num_decay_steps)   # progress in [0, 1]
    if shape == "1-sqrt":
        decayed = 1.0 - math.sqrt(t)          # steep at first, long gentle tail
    elif shape == "linear":
        decayed = 1.0 - t
    elif shape == "cosine":
        decayed = 0.5 * (1.0 + math.cos(math.pi * t))
    else:
        raise ValueError(shape)
    return final_frac + (1.0 - final_frac) * decayed

# Quick sanity check of the 1-sqrt shape:
if __name__ == "__main__":
    N = 10_000
    for frac in (0.0, 0.25, 0.5, 0.75, 1.0):
        s = int(frac * N)
        print(f"progress {frac:>4.0%}  ->  lr_mult {wsd_decay_multiplier(s, N):.4f}")
    # progress   0%  ->  lr_mult 1.0000
    # progress  25%  ->  lr_mult 0.5000
    # progress  50%  ->  lr_mult 0.2929
    # progress  75%  ->  lr_mult 0.1340
    # progress 100%  ->  lr_mult 0.0000

Notice the shape: half the LR is gone by the 25% mark and the last three-quarters of the decay window runs at under 30% of peak. That long low-LR tail is where the premium data gets committed. Feeding math and code into that tail (sub-phase C) is why we schedule capability injection last.

mult = wsd_decay_multiplier(mid_step, total_decay_steps)
for opt, peak in ((muon, MUON_PEAK_LR), (adamw, ADAMW_PEAK_LR)):
    for g in opt.param_groups:
        g["lr"] = peak * mult

Why 1−sqrt, and how to watch the anneal working¶

The 1-sqrt shape is the MiniCPM default for a concrete reason: it drops the LR aggressively at the start of decay — where the model still has enough plasticity to re-organize — then spends a long, patient tail near the floor consolidating on the premium tokens. Linear decay spreads the drop evenly and tends to leave a little loss on the table; cosine is smoother but its slow early descent wastes some of the high-plasticity window. If you only run one anneal, 1-sqrt is the safe pick. The one knob worth ablating (Ch. 14.5) is final_frac: decaying to a hard 0.0 squeezes out the last fraction of loss but leaves you a checkpoint that is finished and awkward to continue-train; decaying to ~0.1 of peak keeps the door open for a second anneal or for post-training that prefers a non-zero starting LR. We decay to the floor for the release checkpoint and keep a pre-final checkpoint (ckpt_mid_longctx.pt) around as the resumable one.

You should watch the anneal rather than trust it. The single most informative diagnostic is a held-out loss curve on a fixed, frozen validation set sampled from the pretrain distribution, evaluated every few hundred steps throughout mid-training. Because the validation set never changes, the curve is directly comparable to the stable-phase plateau, and you will see the characteristic WSD “elbow”: a visible downward bend the moment decay begins, steepening through the low-LR tail. If that elbow does not appear, something is wrong — the usual suspects, in order of frequency, are: the resume loaded weights but not the optimizer state; the LR multiplier is being computed against the global step instead of mid_step; or the two optimizer groups were collapsed onto one LR (see the warning above). A second, cheaper signal is the gradient norm: it should shrink smoothly as the LR falls; a rising grad norm during decay means the premium mix is too far out-of-distribution (dial back the instruction-flavored slice) or the batch is too small for the long sequences.

Move 2 — Long-Context Extension to 8192¶

Stack-100M was pretrained at max_seq_len = 2048 with RoPE base \(\theta = 10000\) (see the config in Chapter 14.4). We now want it to attend over 8192 tokens. Three things must change together, and skipping any one of them silently wastes the sub-phase: the positional geometry (RoPE base), the data (shards that actually contain positions past 2048), and the attention mask (a dense (B, 1, T, T) mask does not survive 8192). We take them in that order.

The extrapolation problem, concretely¶

RoPE rotates each 2-D slice of a query/key by an angle \(m\theta_k\) that grows linearly with position \(m\), where \(\theta_k = \theta^{-2k/d}\) for pair index \(k = 0,1,\dots,d/2-1\) and \(d\) is the head dimension. For Stack-100M, \(d = 64\), so there are \(32\) frequency pairs. The dot product \(q_m \cdot k_n\) depends only on the relative offset \(m - n\) — that is RoPE’s defining property, derived from scratch in Positional Encodings: Sinusoidal, Learned, RoPE & ALiBi.

The problem: during pretraining, \(m\) never exceeded \(2048\). The low-frequency pairs (large \(k\), tiny \(\theta_k\)) barely complete a fraction of a rotation within that window. If we naively run the model at position \(8191\), those slow dimensions swing into angular ranges the model has never seen, and attention scores go pathological. The high-frequency (local) pairs are fine — they complete many turns even within 2048 — but the long-range pairs break.

RoPE base rescaling (NTK-aware)¶

The fix is to stretch the frequency ladder so that positions up to 8192 land in the same angular range the model already learned for positions up to 2048. NTK-aware scaling (the “increase the base” trick, popularized on LocalLLaMA by bloc97, 2023) does this by raising the RoPE base \(\theta\):

A larger base makes every \(\theta_k\) smaller, so each dimension rotates more slowly — exactly compensating for the longer positions. The \(d/(d-2)\) exponent is the NTK correction that leaves the highest-frequency pair essentially untouched (preserving local resolution) while stretching the low-frequency pairs (which needed the range). YaRN (Peng et al., 2023) refines this per-wavelength and adds an attention-temperature (“length scaling”) correction that divides the pre-softmax logits by a factor growing with \(\log\) of the scale, compensating for the entropy growth of attention over longer contexts; we cross-link its full treatment in Long-Context Pretraining & Context Extension and use the simpler base rescale here, because we are going to continue-train — a small amount of training at 8192 repairs any residual mismatch far more cheaply than getting the zero-shot formula perfect. (If you were extending without continued training — the zero-shot case — YaRN is clearly the better choice, and it is what transformers exposes at the config level: rope_scaling={"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 2048} alongside linear, dynamic, longrope, and llama3 variants.)

Here is the rescale, wired to rebuild the model’s RoPE cache in place. Stack100M holds one pair of rope_cos / rope_sin buffers and indexes them by the batch’s position_ids before passing cos/sin into every block, so there is a single cache to rebuild — rebuild_rope does it and updates cfg in the same breath, which is what makes the new geometry land in every subsequent checkpoint. Because Stack-100M uses NoPE on every 4^th layer (SmolLM3, HuggingFace 2025; Kazemnejad et al., 2023), those layers ignore cos/sin entirely and are untouched by rescaling — they already generalize across length by construction, which is precisely why we included them.

# capstone/stacklm/model/rope.py  (Ch. 14.4 — signatures reproduced for reference)
import torch

def build_rope_cache(head_dim: int, max_seq: int, theta: float,
                     device=None, dtype=torch.float32):
    """Precompute cos/sin tables of shape (max_seq, head_dim). Mid-training changes
    only `max_seq` and `theta`; the code path is identical to pretraining."""
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float()
                                / head_dim))                  # (head_dim/2,)
    t = torch.arange(max_seq, device=device).float()           # positions
    freqs = torch.outer(t, inv_freq)                           # (max_seq, head_dim/2)
    emb = torch.cat((freqs, freqs), dim=-1)                    # (max_seq, head_dim)
    return emb.cos().to(dtype), emb.sin().to(dtype)

def ntk_rescaled_base(base: float, head_dim: int,
                      old_len: int, new_len: int) -> float:
    """NTK-aware RoPE base rescaling: theta' = theta * s^(d/(d-2))."""
    s = new_len / old_len
    return base * (s ** (head_dim / (head_dim - 2)))

# capstone/stacklm/mid/continue_training.py  (this chapter)
import torch
from ..model.rope import ntk_rescaled_base


def _unwrap(model):
    """`torch.compile` wraps the module; reach the real one for cfg/buffer surgery."""
    return model._orig_mod if hasattr(model, "_orig_mod") else model


@torch.no_grad()
def extend_context(model, new_seq_len: int, device="cpu") -> float:
    """Rescale RoPE base for a longer context and rebuild the cache. Returns the
    new theta. NoPE layers are unaffected (they never consult the cache).

    Call this ONCE, at the start of sub-phase B. `rebuild_rope` swaps the model's
    single (rope_cos, rope_sin) buffer pair and updates cfg.max_seq_len /
    cfg.rope_theta, so the new geometry is recorded in every later checkpoint.
    """
    m = _unwrap(model)
    cfg = m.cfg
    new_base = ntk_rescaled_base(cfg.rope_theta, cfg.head_dim,
                                 old_len=cfg.max_seq_len, new_len=new_seq_len)
    m.rebuild_rope(new_seq_len, new_base, device=device)
    return new_base

Repacking for sub-phase B: where the long documents actually come from¶

This is the step most write-ups skip, and it is the one that decides whether sub-phase B does anything at all.

Recall two facts from Chapter 14.2. (1) The corpus was packed into fixed 2048-token windows (SEQ_LEN = 2048), so PackedMemmapDataset("data/fineweb_edu") physically cannot hand you an 8192-token row. (2) Every document’s position ids restart at 0 inside the packed window. Put those together and the consequence is stark:

With per-document position resets, the largest position id the model ever sees equals the length of the longest document, not the length of the packed window. If your longest document is 900 tokens, rescaling RoPE for 8192 and training 0.6B tokens on it teaches the model nothing about positions 2048–8191. The rescale and the compute are both wasted.

So sub-phase B needs two things: shards repacked at seq_len = 8192, and documents long enough to fill them. The first is a re-run of Ch. 14.2’s build_shards with a different seq_len. The second requires actually sourcing long documents, because the capstone’s pretrain mix (PLAN §2) has none: Cosmopedia v2 is short-form synthetic textbook and story generations, FineWeb-Edu’s document-length distribution is heavily skewed short, and FineMath samples are short by construction. Only code has natural length, and only if you stop treating a file as the unit.

Three sources, in decreasing order of yield:

1. Repo-level code. Concatenate all files within a single repository, in a deterministic order, into one “document.” This is the standard trick — StarCoder2 (Lozhkov et al., 2024) and DeepSeek-Coder (Guo et al., 2024) both do repo-level pretraining precisely because it produces genuinely long, genuinely coherent sequences with real cross-file dependencies.
2. Books. PG-19 (Project Gutenberg books, introduced with the Compressive Transformer, Rae et al., 2019; on the Hub as deepmind/pg19) is public-domain, long by construction (whole books), and clean.
3. Papers. The arXiv subset of EleutherAI/proof-pile-2 (Llemma, Azerbayev et al., 2024) gives long, math-dense documents — a nice double-duty with sub-phase C.

Plus a length filter over the sources you already have, which will recover some long FineWeb-Edu documents. Do not guess the yield; measure it on your own shards, because it depends entirely on your dump and your filter settings.

# capstone/scripts/repack_long.py
"""Build the sub-phase-B shards: seq_len=8192, long documents only.

Run once, between sub-phase A and sub-phase B:
    python capstone/scripts/repack_long.py --out data/mid --seq-len 8192
"""
import argparse
from collections import defaultdict

import numpy as np

from stacklm.data import (DataMixEntry, PackedMemmapDataset, build_shards,
                          stream_source)                       # Ch. 14.2
from stacklm.tokenizer import StackTokenizer                   # Ch. 14.3

MIN_DOC_TOKENS = 4096          # half the target window; see the assertion below


def length_filtered(docs, tok, min_tokens: int = MIN_DOC_TOKENS):
    """Keep only documents that can actually exercise positions past 2048.

    Tokenizing twice (here and in `pack_documents`) is wasteful; at 20B tokens you
    would instead carry a `n_tokens` field through the Ch. 14.2 pipeline, or use
    HuggingFace `datatrove`'s TokensCounter + a LambdaFilter to do it in one pass.
    """
    for doc in docs:
        if len(tok.encode(doc["text"])) >= min_tokens:
            yield doc


def repo_level_documents(files, sep: str = "\n\n# ==== file: {path} ====\n\n"):
    """Concatenate a repository's files into ONE document (StarCoder2 / DeepSeek-Coder).

    `files` is a stream of dicts with `repo_name`, `path`, `content`. Sorting by
    path makes the concatenation deterministic (and puts headers near sources,
    which is what a human reading the repo would do).
    """
    by_repo = defaultdict(list)
    for f in files:
        by_repo[f["repo_name"]].append(f)
    for repo, fs in by_repo.items():
        fs.sort(key=lambda f: f["path"])
        body = "".join(sep.format(path=f["path"]) + f["content"] for f in fs)
        yield {"text": body, "source": "starcoder_repo", "repo": repo}


def verify_positions(shard_dir: str, seq_len: int, floor: int = 4096,
                     n_sample: int = 512, seed: int = 0):
    """The check that decides whether sub-phase B is real or theatre.

    Ch. 14.2 does NOT store position ids on disk (`store_positions=False` is the
    default; they are recomputed from `input_ids == bos_id` by
    `segments_from_bos`), so we read them back through the dataset that the
    trainer itself will use -- which also proves the shards are readable and
    packed at the right length. Sampling a few hundred rows is enough: we only
    need ONE window whose document reaches past `floor`.
    """
    ds = PackedMemmapDataset(shard_dir)
    assert ds.seq_len == seq_len, f"{shard_dir} packed at {ds.seq_len}, not {seq_len}"
    rng = np.random.default_rng(seed)
    rows = rng.choice(len(ds), size=min(n_sample, len(ds)), replace=False)
    hi = max(int(ds[int(i)]["position_ids"].max()) for i in rows)
    print(f"  max position id in {shard_dir}: {hi} (window {seq_len})")
    assert hi > floor, (
        f"{shard_dir} contains no document longer than {floor} tokens: RoPE "
        f"rescaling would train on positions the data never reaches.")


# The sub-phase-B sources, as Ch. 14.2 `DataMixEntry` records (name, hf_path,
# weight, domain). The first three are genuinely long; the fourth is a
# length-FILTERED slice of a pretrain source; the last is the deliberately SHORT
# anti-drift anchor. `weight` is the sub-phase-B mixture weight consumed later
# by `build_mixture_loader`.
LONG_SOURCES = [
    (DataMixEntry("starcoder_repo",   "bigcode/starcoderdata",     0.35, "code"),  True),
    (DataMixEntry("books_pg19",       "deepmind/pg19",             0.25, "web"),   False),
    (DataMixEntry("arxiv_proofpile2", "EleutherAI/proof-pile-2",   0.15, "math"),  False),
    (DataMixEntry("fineweb_edu_long", "HuggingFaceFW/fineweb-edu", 0.15, "web"),   False),
    (DataMixEntry("cosmopedia_v2",    "HuggingFaceTB/cosmopedia",  0.10, "synthetic"), False),
]


def main(out_root: str, seq_len: int, tokenizer_path: str):
    tok = StackTokenizer.load(tokenizer_path)             # Ch. 14.3, vocab 32768
    for entry, repo_level in LONG_SOURCES:
        raw = stream_source(entry)                        # Ch. 14.2 streaming reader
        docs = repo_level_documents(raw) if repo_level else raw
        if entry.name != "cosmopedia_v2":     # the short-form anchor stays unfiltered
            docs = length_filtered(docs, tok)
        out = f"{out_root}/{entry.name}_{seq_len}"
        n = build_shards(docs, tok, out, seq_len=seq_len,
                         tokens_per_shard=100_000_000)
        verify_positions(out, seq_len,
                         floor=4096 if entry.name != "cosmopedia_v2" else 0)
        print(f"{entry.name}: {n} shard(s) -> {out}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data/mid")
    ap.add_argument("--seq-len", type=int, default=8192)
    ap.add_argument("--tokenizer", default="artifacts/tokenizer.json")
    a = ap.parse_args()
    main(a.out, a.seq_len, a.tokenizer)

Nothing about the on-disk format changes: the shards stay uint16 (the largest position id, 8191, is well under 65535) and stay position-free, because segments_from_bos reconstructs both seq_ids and position_ids from the tokens themselves. That is why verify_positions goes through PackedMemmapDataset rather than globbing for a .pos.bin file that the default writer never produces.

With those shards in hand, sub-phase B’s mixture is honest about which sources are long:

That last row is not an oversight. Training only on long documents is a known way to degrade short-context quality, and both Llama 3’s long-context stage (Grattafiori et al., 2024) and ProLong (Gao et al., Princeton, 2024) explicitly keep short-form data mixed in while upsampling long documents. The 10% Cosmopedia slice is the anchor that keeps the model from drifting off the distribution sub-phase A just spent 1.2B tokens sharpening. If you read one paper alongside this section, make it ProLong: it is the closest open, reproducible reference for exactly this step — data recipe, length schedule, and evaluation. Exercise 2 works out how many epochs each of these five sources actually sees, which is the number that decides whether your long-context stage is memorizing Gutenberg.

Masking at 8192: FlexAttention and varlen kernels¶

Now the systems half. Stack-100M implements document-aware attention by materializing a boolean mask. Here is the Ch. 14.4 builder, reduced to the training path (its full form also handles the rectangular KV-cache case at decode time):

# capstone/stacklm/model/transformer.py  (Ch. 14.4 — training path of `_build_mask`)
def _build_mask(self, seq_ids, T, kv_len, start_pos, device):
    """Bool mask (B, 1, T, kv_len); True = attend. None = plain-causal fast path."""
    if seq_ids is None and kv_len == T and start_pos == 0:
        return None                                     # SDPA's is_causal=True
    q_pos = torch.arange(start_pos, start_pos + T, device=device)
    kv_pos = torch.arange(kv_len, device=device)
    m = (q_pos[:, None] >= kv_pos[None, :])[None, None]          # (1, 1, T, kv_len)
    if seq_ids is not None:                                       # no cross-document
        same = seq_ids[:, -T:, None] == seq_ids[:, None, :kv_len] # (B, T, kv_len)
        m = m & same[:, None]
    return m

At micro_bs = 32, T = 2048 that mask is \(32 \times 2048^2 = 1.34\times10^8\) bools = 134 MB — annoying but survivable. At micro_bs = 8, T = 8192 it is \(8 \times 8192^2 = 5.37\times10^8\) bools = 537 MB, allocated fresh on every forward, and it must persist through the backward pass. Worse, passing any dense attn_mask to F.scaled_dot_product_attention disqualifies the FlashAttention backend — PyTorch falls back to the memory-efficient or math kernel, so you simultaneously lose the speed and lose the “never materializes the \(T \times T\) score matrix” property that made 8192 affordable in the first place. Chapter 14.2 flagged this as a promise to cash later; here is where it comes due.

There are two production answers in 2026, and you should know both.

Option 1 — FlexAttention (PyTorch-native). torch.nn.attention.flex_attention (PyTorch 2.5+) lets you express the mask as a predicate on indices, mask_mod(b, h, q_idx, kv_idx) -> bool, and compiles it into a fused, block-sparse Triton kernel. create_block_mask evaluates the predicate once per \(128 \times 128\) block and stores only which blocks are non-empty — so memory goes from \(O(T^2)\) to \(O((T/128)^2)\), and fully masked blocks are never computed at all. Document masking is the canonical example, and it makes short documents faster, not slower. (The module also ships and_masks / or_masks combinators, so you can build the causal predicate and the same-document predicate separately and compose them; we inline both for clarity.)

# capstone/stacklm/model/doc_attention.py
import torch
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

# Compile once at import; `dynamic=False` keeps one specialization per shape.
flex_attention = torch.compile(flex_attention, dynamic=False)
_create_block_mask = torch.compile(create_block_mask, dynamic=False)


def document_block_mask(seq_ids: torch.Tensor):
    """Block-sparse causal + intra-document mask from Ch. 14.2's `seq_ids`.

    seq_ids : (B, T) int tensor; tokens of the same packed document share an id.
    Returns a BlockMask consumable by flex_attention. Broadcast over heads (H=None).
    """
    B, T = seq_ids.shape

    def mask_mod(b, h, q_idx, kv_idx):
        causal = q_idx >= kv_idx                       # no peeking ahead
        same_doc = seq_ids[b, q_idx] == seq_ids[b, kv_idx]
        return causal & same_doc

    return _create_block_mask(mask_mod, B=B, H=None, Q_LEN=T, KV_LEN=T,
                              device=seq_ids.device)


def doc_attention(q, k, v, block_mask):
    """q, k, v: (B, n_heads, T, head_dim) -- GQA heads already expanded (Ch. 14.4).
    Scaling is 1/sqrt(head_dim) by default, matching SDPA."""
    return flex_attention(q, k, v, block_mask=block_mask)

Wiring it in is a two-line change to Attention.forward: build the BlockMask once per micro-batch in Stack100M.forward (not once per layer — that is 30× the mask-construction cost for an identical result) and thread it down in place of attn_mask, then call doc_attention instead of F.scaled_dot_product_attention. At T = 8192 the block mask is \(64 \times 64 = 4096\) block entries per batch element — kilobytes instead of half a gigabyte.

Option 2 — varlen FlashAttention. The flash-attn library (Dao et al.) exposes flash_attn_varlen_func, which takes unpadded concatenated tokens plus a cumulative-sequence-length index and runs a separate causal attention per document, never computing a cross-document score at all. This is what Megatron-LM and most production stacks use for packed pretraining.

# Requires: pip install flash-attn --no-build-isolation   (CUDA, Ampere or newer)
import torch
from flash_attn import flash_attn_varlen_func

def varlen_attention(q, k, v, doc_lens: torch.Tensor):
    """q, k, v: (total_tokens, n_heads, head_dim) -- the batch is FLATTENED and
    unpadded. doc_lens: (n_docs,) int32 lengths, summing to total_tokens.

    cu_seqlens is the standard 'cumulative sequence lengths' index: the exclusive
    prefix sum with a leading 0, i.e. document i occupies [cu[i], cu[i+1]).
    """
    cu = torch.cat([torch.zeros(1, dtype=torch.int32, device=doc_lens.device),
                    doc_lens.cumsum(0).to(torch.int32)])
    max_len = int(doc_lens.max())
    return flash_attn_varlen_func(q, k, v,
                                  cu_seqlens_q=cu, cu_seqlens_k=cu,
                                  max_seqlen_q=max_len, max_seqlen_k=max_len,
                                  causal=True)

Which to pick? FlexAttention, for this capstone: it is PyTorch-native (no extra CUDA build), it composes with torch.compile and with the rest of Stack-100M unchanged, and it keeps the (B, T) batched layout the Ch. 14.2 dataset already produces. Varlen FlashAttention is marginally faster and is the right answer inside a framework that already flattens batches (Megatron-LM, and xFormers’ BlockDiagonalCausalMask.from_seqlens for the same idea with a different API) — but adopting it means rewriting the data path to emit cu_seqlens. Use it when you graduate to a multi-GPU framework; use FlexAttention today.

Validating that the long-context actually works¶

Perplexity averaged over a sequence can fall even when the model still cannot use position 8000 — a couple of well-predicted local tokens hide a broken tail. Two cheap probes catch this.

First, loss-versus-position: bin the per-token loss by its position within the 8192-window and plot the curve. A healthy extension shows loss that keeps decreasing (or at worst flattening) as position grows — the model is using more context to predict better. A curve that rises past ~2048 means the rescale-plus-training has not taken and the far positions are still noise. Crucially, run this on the long-document subset only: on document-masked packed windows full of short documents, position within the window is not position within the document, and the curve tells you nothing. Exercise 6 implements the binning.

Second, a tiny needle-in-a-haystack check: plant a short unique fact (“the passcode is 4713”) at a random depth in a long filler document and ask the model, via next-token prediction, to complete “the passcode is”. If it recovers the needle only when it sits in the first 2048 tokens, you have geometry without capability — train sub-phase B longer or on longer documents.

The 2026 standard long-context benchmarks are RULER (Hsieh et al., NVIDIA, 2024), which generalizes needle-in-a-haystack into 13 synthetic tasks with controllable length, and HELMET (Yen et al., Princeton, 2024–2025), which evaluates on realistic downstream long-context applications. Both are the right tools for a frontier model and both will read near floor for a 100M model — quote them as the destination, not as evidence. The honest, full evaluation belongs in Chapter 14.11; loss-versus-position and the needle probe are the two you run during mid-training to know the elbow is real.

Move 3 — Capability Injection¶

The final sub-phase (C) runs at the LR floor and pushes the mixture hard toward the capabilities our downstream narrow agent needs: arithmetic, symbolic manipulation, and code structure.

This is a small budget (~0.2B tokens) at a low, still-nonzero LR — enough to sharpen number-and-symbol handling without over-fitting or wrecking general fluency. It is honest to call this what it is: not turning Stack-100M into a math model, but giving a 100M model just enough arithmetic and code grounding that the RLVR run in Chapter 14.9 (narrow GRPO on verifiable integer arithmetic, following GRPO from DeepSeekMath, Shao et al., 2024) and the ReAct agent (Yao et al., 2022) in Chapter 14.10 have a base to reinforce. At 100M params you cannot inject a capability the model has no capacity for; you can only make sure the capacity you have is pointed at the right target. That is “narrow but real.”

Capability injection reuses the exact same loop and schedule — it is simply sub-phase C’s mixture, applied while the LR finishes its decay to the floor. Because it runs at seq_len = 8192 on the sub-phase-B shards, it needs no new repack: only a mixture dict.

Why this ordering¶

The three moves are sequenced deliberately, and the order is not arbitrary. Annealing runs first, at 2048, because the general quality jump wants the highest LR of the decay leg and the largest, most diverse token budget — you consolidate broad representations before you specialize. Long-context extension runs second because it is a targeted change to a subsystem (the positional geometry) that benefits from a settled model: rescaling RoPE on a mid-anneal checkpoint whose attention patterns have already sharpened means the few hundred million long-context tokens teach the heads to use the new range rather than fighting a still-shifting representation. Capability injection runs last, at the LR floor, because it is the narrowest and most over-fitting-prone mix; keeping it in the low-LR tail limits how far it can pull the model off the general manifold while still committing arithmetic and code structure.

Reverse any two of these and you pay for it. Putting the narrow math/code mix first would spend the highest-LR, highest-plasticity window dragging the model off the general manifold, and the later broad phase would have to undo that at a lower LR where it has less leverage. A useful mental model: broad → long → narrow, riding one decreasing learning rate down.

The Mid-Training Loop: One Continuous Decay¶

Now we assemble everything into a single continued-training routine that resumes from ckpt_stable.pt and runs all three sub-phases back to back. It reuses Stack100M, the Muon+AdamW hybrid optimizer, and the checkpoint helpers established in Chapters 14.4, 14.6, and 14.7 — mid-training adds no new model code, only orchestration.

The driver below is the shipped run_mid_training in capstone/stacklm/mid/continue_training.py.

# capstone/stacklm/mid/continue_training.py  (the phase driver)
from dataclasses import dataclass, field

import torch

from ..optim import build_optimizers                   # Muon+AdamW hybrid, Ch. 14.6
from ..optim.schedule import wsd_decay_multiplier      # Ch. 14.6
from ..train.loop import autocast_ctx                  # bf16 on CUDA, fp32 on CPU


@dataclass
class SubPhase:
    """One mid-training sub-phase. `steps` overrides the token-budget arithmetic
    (used by CI, which runs a handful of steps instead of ~1e9 tokens)."""
    name: str
    tokens: int
    seq_len: int
    mix: dict = field(default_factory=dict)
    extend_now: bool = False       # rescale RoPE at the start of this sub-phase
    steps: int = 0                 # 0 => derive from `tokens`


def steps_for(sub: SubPhase, global_batch_tokens: int) -> int:
    return sub.steps or (sub.tokens // global_batch_tokens)


def run_mid_training(model, phases, loader_fn, *, device="cpu",
                     global_batch_tokens=524_288, micro_batch_tokens=65_536,
                     muon_lr=6e-3, adamw_lr=3e-3, grad_clip=1.0,
                     optimizers=None, use_seq_ids=True, log_every=50, seed=1234,
                     checkpoint_fn=None):
    """Walk `phases` back to back under ONE WSD decay leg.

    `loader_fn(sub, micro_batch_size)` returns an *infinite* iterator of batch
    dicts (`input_ids`, `position_ids`, `seq_ids`, `targets`) at `sub.seq_len`.
    Pass `optimizers=[muon, adamw]` to CONTINUE the stable phase's optimizer state
    (Muon momentum, AdamW moments) restored from `ckpt_stable.pt`.
    """
    torch.manual_seed(seed)
    device = torch.device(device)
    model.to(device).train()

    if optimizers is None:
        muon, adamw = build_optimizers(model, muon_lr=muon_lr, adamw_lr=adamw_lr)
        optimizers = [muon, adamw]
    else:
        muon, adamw = optimizers[0], optimizers[-1]
    peaks = {id(muon): muon_lr, id(adamw): adamw_lr}     # two groups, two peaks

    # The decay spans ALL sub-phases: it is never reset at a boundary.
    total_decay_steps = sum(steps_for(p, global_batch_tokens) for p in phases)
    mid_step, history = 0, []

    for sub in phases:
        if sub.extend_now:
            new_base = extend_context(model, sub.seq_len, device=device)
            print(f"  [{sub.name}] RoPE base -> {new_base:.0f}, seq_len -> {sub.seq_len}")

        # Micro-batch shrinks with sequence length to hold activation memory
        # roughly constant; grad-accum makes up the global batch.
        micro_bs = max(1, micro_batch_tokens // sub.seq_len)
        accum = max(1, (global_batch_tokens // sub.seq_len) // micro_bs)
        it = loader_fn(sub, micro_bs)

        for _ in range(steps_for(sub, global_batch_tokens)):
            mult = wsd_decay_multiplier(mid_step, total_decay_steps)
            for opt in optimizers:
                for g in opt.param_groups:
                    g["lr"] = peaks[id(opt)] * mult      # SAME multiplier, own peak
                opt.zero_grad(set_to_none=True)

            loss_acc = 0.0
            for _micro in range(accum):
                batch = {k: v.to(device) for k, v in next(it).items()}
                with autocast_ctx(device):
                    # z-loss is inside the model (cfg.z_loss_coef, Ch. 14.4) and
                    # stays on through the low-LR tail: it keeps logits scaled.
                    _, loss = model(batch["input_ids"], targets=batch["targets"],
                                    position_ids=batch["position_ids"],
                                    seq_ids=batch["seq_ids"] if use_seq_ids else None)
                    loss = loss / accum
                loss.backward()
                loss_acc += loss.item()

            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            for opt in optimizers:
                opt.step()
            history.append(loss_acc)
            if log_every and mid_step % log_every == 0:
                print(f"  [{sub.name}] mid_step {mid_step:>5} loss {loss_acc:.4f} "
                      f"lr_muon {peaks[id(muon)] * mult:.2e} seq {sub.seq_len}")
            mid_step += 1

        if checkpoint_fn is not None:
            checkpoint_fn(sub.name, model, optimizers, mid_step)

    return {"loss_history": history, "mid_steps": mid_step,
            "total_decay_steps": total_decay_steps}

And the full-run entry point that supplies the three sub-phases, the mixtures, and the resume:

# capstone/scripts/midtrain.py
"""Full-run mid-training: resume ckpt_stable.pt and run the WSD decay across
anneal @2048 -> long-context @8192 -> capability injection @8192.

    python capstone/scripts/midtrain.py \
        --stable artifacts/ckpt_stable.pt --out artifacts/ --device cuda
"""
import argparse
from dataclasses import asdict

from stacklm.config import StackConfig
from stacklm.model import Stack100M
from stacklm.optim import build_optimizers
from stacklm.train.loop import load_checkpoint, save_checkpoint     # Ch. 14.7
from stacklm.mid import SubPhase, run_mid_training
from stacklm.mid.mixture import (build_mixture_loader, ANNEAL_MIX,
                                 LONGCTX_MIX, CAPABILITY_MIX)

# 32 x 2048 x 8 accumulation steps = 524,288 tokens/step, identical to the
# pretraining loop's effective batch (Ch. 14.6/14.7). A POWER OF TWO on purpose:
# it divides exactly by every (micro_bs, seq_len) pair we will run, so the
# realized batch equals the nominal one at 2048 and at 8192 alike.
GLOBAL_BATCH_TOKENS = 524_288
MICRO_BATCH_TOKENS  = 65_536     # 32 x 2048 at pretrain length; 8 x 8192 when long
MUON_PEAK_LR        = 6e-3       # Muon group's stable-phase peak (Ch. 14.6)
ADAMW_PEAK_LR       = 3e-3       # AdamW group's peak = muon_lr / 2 (Ch. 14.6)

# The three moves of mid-training, in order. Token budgets are illustrative
# (~2B total = ~10% of the 20B pretrain budget); tune per Ch. 14.5.
PHASES = [
    SubPhase("anneal", tokens=1_200_000_000, seq_len=2048, mix=ANNEAL_MIX),
    SubPhase("longctx", tokens=600_000_000, seq_len=8192, extend_now=True,
             mix=LONGCTX_MIX),
    SubPhase("capability", tokens=200_000_000, seq_len=8192, mix=CAPABILITY_MIX),
]


def main(stable_ckpt: str, out_dir: str, device: str):
    # ---- 1. Resume the *pre-decay* model AND optimizer state ------------------
    model = Stack100M(StackConfig()).to(device)
    optimizers = list(build_optimizers(model, muon_lr=MUON_PEAK_LR,
                                       adamw_lr=ADAMW_PEAK_LR))   # (muon, adamw)
    step, extra = load_checkpoint(stable_ckpt, model, optimizers,
                                  map_location=device)
    # `extra` is the payload Ch. 14.7 stores alongside the tensors: tokens_seen,
    # the data cursor, and the training config. Trust it over this file.
    muon_peak = extra.get("muon_lr", MUON_PEAK_LR)
    adamw_peak = extra.get("adamw_lr", ADAMW_PEAK_LR)
    print(f"resumed {stable_ckpt} @ global step {step} "
          f"({extra.get('tokens_seen', 0)/1e9:.1f}B tokens, LR still at peak); "
          f"peaks muon={muon_peak} adamw={adamw_peak}")

    # ---- 2. One loader factory per sub-phase, from that sub-phase's mixture ---
    def loader_fn(sub, micro_bs):
        return build_mixture_loader(sub.mix, sub.seq_len, micro_bs,
                                    root="data/mid", seed=1234 + len(sub.name))

    def checkpoint_fn(name, model, opts, mid_step):
        save_checkpoint(f"{out_dir}/ckpt_mid_{name}.pt", model, opts,
                        step=step + mid_step,
                        extra={**extra, "phase": name,
                               "cfg": asdict(model.cfg),   # records the NEW rope geometry
                               "mid_step": mid_step})
        print(f"[{name}] done -> ckpt_mid_{name}.pt")

    # ---- 3. Walk the three sub-phases under one continuous decay --------------
    result = run_mid_training(
        model, PHASES, loader_fn, device=device,
        global_batch_tokens=GLOBAL_BATCH_TOKENS,
        micro_batch_tokens=MICRO_BATCH_TOKENS,
        muon_lr=muon_peak, adamw_lr=adamw_peak,
        optimizers=optimizers,          # CONTINUE Muon momentum + AdamW moments
        log_every=50, checkpoint_fn=checkpoint_fn)

    save_checkpoint(f"{out_dir}/ckpt_mid_final.pt", model, optimizers,
                    step=step + result["mid_steps"],
                    extra={**extra, "cfg": asdict(model.cfg)})
    print(f"mid-training complete: {result['mid_steps']} steps "
          f"-> ckpt_mid_final.pt (ready for SFT in Ch. 14.9)")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--stable", default="artifacts/ckpt_stable.pt")
    ap.add_argument("--out", default="artifacts")
    ap.add_argument("--device", default="cuda")
    a = ap.parse_args()
    main(a.stable, a.out, a.device)

A few design points worth flagging:

• One decay across three phases. total_decay_steps spans all sub-phases and mid_step runs continuously, so the LR falls smoothly from peak to floor through the sequence-length change and the mixture changes. We do not reset the schedule at each boundary — that would create three little decay cliffs instead of one clean anneal. This is the single most important line in the file.
• Two peaks, one multiplier. peaks[id(opt)] * mult scales Muon and AdamW by the same schedule while preserving the 2:1 ratio Ch. 14.6 fixed between them.
• position_ids are threaded alongside seq_ids. Omit them and Stack100M.forward falls back to torch.arange(T) — contiguous positions that do not reset per document. That silently changes the geometry from what pretraining used, and it invalidates the entire “largest position id = longest document” argument above, because RoPE would then see 0..8191 in every window regardless of what the documents look like. The packer emits position_ids; use them.
• seq_ids too. Document-aware masking is not optional and does not lapse at mid-training; it is exactly as load-bearing at 8192 as at 2048 — more so, because a packed 8192-window holds four times as many unrelated documents.
• We resume the optimizer state, not just the weights. Muon’s momentum buffers and AdamW’s moments carry over from the stable phase — which is why optimizers is built before load_checkpoint and passed into run_mid_training rather than being rebuilt inside it. Throwing that state away injects a transient the fresh decay does not need. This is the concrete payoff of saving ckpt_stable.pt with optimizer state in Ch. 14.7.
• Micro-batch shrinks, global batch holds. At 2048, micro_bs = 32 and accum = 256 // 32 = 8; at 8192, micro_bs = 8 and accum = 64 // 8 = 8. Both yield exactly 524,288 tokens per optimizer step, so the LR/batch relationship (Ch. 3.10) is unchanged across the sequence-length jump.
• z-loss stays on. The small logsumexp penalty lives inside Stack100M.forward via cfg.z_loss_coef (Ch. 14.4); it keeps logits well-scaled through the low-LR tail, and there is nothing to re-enable.

What CI actually covers¶

Per the PLAN’s deliverable standard, capstone/smoke_test.py exercises this code at toy scale, CPU-only and hermetic — a 2-layer d_model=64 model, an in-process synthetic corpus, and a handful of steps. Concretely it runs, in this order:

1. mid_train(model, ds, steps=4, micro_batch_size=4, grad_accum=2, peak_lr=1.5e-3, extend_to=2*max_seq_len) — the single-phase entry point — and asserts the losses are finite and model.cfg.max_seq_len doubled.
2. run_mid_training(model, phases, toy_loader, ...) over three SubPhases (anneal → longctx with extend_now=True → capability), asserting that mid_steps == total_decay_steps (the decay was never reset), that all three boundary checkpoints fired in order, that cfg.max_seq_len grew again, and that cfg.rope_theta strictly increased — the machine-checkable statement of “we rescaled the base, once, upward.”

What CI does not cover, and cannot: the real 8192 shapes, the FlexAttention kernel (CUDA-only), and the actual data mixtures. Those are documented here with full-run magnitudes, exactly as the PLAN’s “CI proves it runs; the prose documents the expected numbers” standard requires. Do not read a green smoke test as evidence that your long-context extension worked — that is what verify_positions and the loss-versus-position curve are for.

What Mid-Training Costs and Buys¶

You should expect the held-out loss to fall visibly across sub-phase A — on the order of a couple tenths of a nat below where the stable phase plateaued — with most of the drop concentrated in the low-LR tail. Long-context sub-phase B will raise the average loss slightly (8192-token prediction on books and whole repositories is genuinely harder than 2048-token snippets, and the mix itself changed), which is expected and correct, not a regression; the loss-versus-position curve is the metric that shows the extension worked even as the scalar average ticks up. Capability sub-phase C nudges arithmetic and code perplexity down at the cost of a hair of general-web perplexity — the trade we are deliberately making. Report these as illustrative movements; never quote a fabricated benchmark. The honest evaluation lives in Chapter 14.11.

Key Takeaways¶

Further Reading¶

• OLMo 2 Team, 2 OLMo 2 Furious (Allen Institute for AI, 2024–2025): the open report that named and detailed the mid-training / annealing stage.
• Hu et al., MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies (2024): the WSD schedule and decay-phase data annealing.
• Gao et al., How to Train Long-Context Language Models (Effectively) (ProLong, Princeton NLP, 2024): the long-context continued-training data recipe — code repositories and books, with short data retained.
• Penedo et al., The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale (HuggingFace, 2024): FineWeb-Edu and the Cosmopedia synthetic-data recipe.
• Peng et al., YaRN: Efficient Context Window Extension of Large Language Models (2023): per-wavelength RoPE scaling and the attention-temperature correction.
• Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021): the original RoPE.
• Chen et al., Extending Context Window of Large Language Models via Positional Interpolation (2023): position interpolation, the baseline NTK/YaRN build on.
• Ibrahim et al., Simple and Scalable Strategies to Continually Pre-train Large Language Models (2024): LR re-warming/re-decaying, replay, and the cost of resuming a decayed checkpoint.
• Lozhkov et al., StarCoder 2 and The Stack v2 (BigCode, 2024) and Guo et al., DeepSeek-Coder (2024): repo-level code pretraining, the source of sub-phase B’s longest documents.
• Rae et al., Compressive Transformers for Long-Range Sequence Modelling (2019): introduces PG-19, the book-length benchmark corpus.
• Hsieh et al., RULER (NVIDIA, 2024) and Yen et al., HELMET (Princeton, 2024): modern long-context evaluation beyond needle-in-a-haystack.
• Kazemnejad et al., The Impact of Positional Encoding on Length Generalization in Transformers (2023), and the SmolLM3 report (HuggingFace, 2025): NoPE and length generalization.
• Dao et al., FlashAttention and the PyTorch FlexAttention blog post (2024): the two kernel-level answers to document masking on packed long sequences.

Exercises¶

1. You do not get to resume from a pre-decay checkpoint: you inherit a fully-decayed, publicly released ~100M base model and are asked to anneal it on your premium mix for ~1B tokens. Design the learning-rate schedule you would use (peak, warmup, shape, floor), name one data-mixture change you would make that this chapter’s sub-phase A does not need, and state what the whole thing costs you relative to the pre-decay path — citing the result that quantifies it.

2. Sub-phase B draws 0.6B tokens from LONGCTX_MIX. Suppose your repack pass produced these shard sizes, in 8192-token sequences: starcoder_repo 40,000; books_pg19 9,000; arxiv_proofpile2 12,000; fineweb_edu_long 6,000; cosmopedia_v2 200,000. (a) Compute the tokens drawn from each source and the number of epochs each sees. (b) Which source is repeated most, what does that risk at the LR floor, and name two fixes that do not require sourcing new data.

3. Use the budgets from the chapter: sub-phase A = 1.2B tokens, B = 0.6B, C = 0.2B, with GLOBAL_BATCH_TOKENS = 524_288. The LR decays as one continuous 1-sqrt WSD leg across all three phases (final_frac = 0), with MUON_PEAK_LR = 6e-3 and ADAMW_PEAK_LR = 3e-3 (Ch. 14.6). Compute (a) total_decay_steps, and (b) the learning rate of each optimizer group at the A→B boundary and at the B→C boundary.

4. A different model, Stack-Big, has head dimension \(d = 128\), was pretrained at \(L_{\text{old}} = 2048\) with RoPE base \(\theta = 10000\), and you want to extend it to \(L_{\text{new}} = 16384\). (a) Compute the NTK-rescaled base \(\theta'\). (b) The fastest frequency pair (\(k=0\)) is supposed to be left essentially untouched — show that it is exactly unchanged by any base rescale, and say in one sentence why that matters.

5. The driver shrinks the micro-batch as the sequence grows so the global batch stays fixed: micro_bs = max(1, MICRO_BATCH_TOKENS // seq_len), seqs_per_global = GLOBAL_BATCH_TOKENS // seq_len, accum = seqs_per_global // micro_bs, with MICRO_BATCH_TOKENS = 65_536. (a) Compute micro_bs, accum, and tokens-per-optimizer-step at seq_len = 2048 and 8192 for GLOBAL_BATCH_TOKENS = 524_288. (b) Now redo it for the plausible-looking choice GLOBAL_BATCH_TOKENS = 512_000, and explain which downstream number this breaks worst — the optimizer’s behavior, or the project’s bookkeeping.

6. Implement the loss-versus-position diagnostic. Write bin_loss_by_position(per_token_loss, num_bins) that takes a (B, T) tensor of per-position next-token losses and returns the mean loss in each of num_bins equal-width position bins. Then say why the diagnostic must be run on the long-document subset.

import torch

def bin_loss_by_position(per_token_loss: torch.Tensor, num_bins: int = 16) -> torch.Tensor:
    """Mean next-token loss binned by position within the sequence.

    per_token_loss : (B, T) tensor of per-position NLL (targets already shifted),
                     e.g. F.cross_entropy(logits, y, reduction="none") reshaped to (B, T).
    Returns         : (num_bins,) tensor; entry j = mean loss over positions in bin j.
    Bin j spans positions [j*T/num_bins, (j+1)*T/num_bins). Average over the batch
    first, then scatter-add into bins so every position contributes equally.
    """
    B, T = per_token_loss.shape
    per_pos = per_token_loss.mean(dim=0)                       # (T,) average over batch
    pos = torch.arange(T, device=per_token_loss.device)
    bin_idx = (pos * num_bins) // T                            # 0 .. num_bins-1
    sums = torch.zeros(num_bins, device=per_token_loss.device)
    counts = torch.zeros(num_bins, device=per_token_loss.device)
    sums.index_add_(0, bin_idx, per_pos)
    counts.index_add_(0, bin_idx, torch.ones_like(per_pos))
    return sums / counts.clamp(min=1)                          # (num_bins,)

7. (Systems) At micro_bs = 8, T = 8192, n_heads = 8, compute (a) the memory of the dense (B, 1, T, T) boolean document mask that Stack100M._build_mask produces, and (b) the number of \(128\times128\) blocks a FlexAttention BlockMask tracks per batch element. © A packed 8192-window holds documents averaging 1024 tokens. Estimate the factor by which block-sparse document masking reduces attention FLOPs versus a plain causal mask, and explain why the saving is not free at sub-phase B specifically.