Blame for source/src/deg/train.py

"""deg trainer: datagen records + jev v1 decisions -> one model (SPEC.md §12).

Losses (each a mean over its items in the micro-batch):

* ``decide``  KL(t || q) over decide outcomes. ``t`` = the record's ``decide_probs`` (Clef
  probabilities where Clef agreed with the filled value), else one-hot of the filled value;
  jev questions use their teacher ``probs`` (else ``hard``).
* ``extract`` start/end pointer CE. Every ``span_alternatives`` entry is a positive: with
  ``span_probs`` the target is that soft distribution, otherwise the loss is the marginal
  ``-log sum_pos p``. A field absent from the filled value targets the null outcome.
* ``count``   CE of the array length (masked to the array's bounds); index-addressed arrays.
* ``anchor``  anchor mode (``--array-mode anchor``, default): per anchored array instance,
  BCE of every in-region context token against the gold item starts (summed, divided by
  1 + the number of items, so the 0.5 threshold stays a probability). Gold item starts
  come from the datagen spans: an item's anchor is the first token of the earliest span
  among its leaves that no sibling item shares (else its earliest span). An array whose
  items do not all get distinct in-region anchors is left out of training (logged).
* ``anchor_count`` (``--anchor-count-weight``, default 0 = off): the count head's CE on
  every anchored array of a whole context (gold item count, capped at the head's
  capacity), so count-then-peaks decoding (``array_decode``) has a trained count.
* ``lm``      next-token CE on ``x-generate`` text (+ EOS), teacher-forced on gold values;
  weight ``--lm-weight`` (default 1). The LM has no head of its own (the tied LM head is
  frozen), so it learns only through the backbone's share of the clipped gradient.
LoRA (``--lora-rank R``, SPEC.md §12.7): the backbone is frozen (= the base model) and
only named low-rank adapters on its projections train, with the heads as before.
``--adapters shared`` = one adapter for every decide/extract fork; ``decide,extract`` =
one per capability, chosen by the fork's kind (a decide fork runs the decide adapter, an
extract/count/anchor fork the extract adapter), so each adapter is trained only by its
own losses (no gradient interference between capabilities in the backbone).
``--lora-context base`` (default) reads the context ONCE with the base weights (no
gradient through the trunk at all; generate forks continue that same trunk at
inference); ``--lora-context <adapter>`` reads it with that adapter (generation then
re-reads the context adapter-off). x-generate is never trained under LoRA: it decodes
with the base (adapter off).
* ``lm_kl_base`` (``--lm-kl-base W``, default 0 = off): anti-forgetting on the same
  x-generate target positions: KL(p_base || p_deg) of the next-token distributions,
  where p_base is the frozen base Qwen3.5 (:func:`deg.model.attach_base`, no grad) on
  the identical trunk + fork tokens (self-distillation toward the untrained model's
  language ability; deg-suite v1 showed training made generated text unfaithful).

Filler augmentation (``--filler-frac P``): a drawn short datagen record is, with
probability P, rebuilt with its document at a random offset inside unrelated text from
the v3 corpus (``--filler-corpus``; train-split docs only), total length log-uniform up to
``--filler-max-tokens`` (gold spans shifted); a ``--filler-negative-frac`` share of them is
filler only, training the schema's top-level anchored arrays as empty (all negative).

Each record is one trunk (its context) and many forks (one per field / array / generated
text); forks continue the trunk state, so gradients reach the shared context encoding.
Choice options are shuffled (with the target permuted) so option order carries no signal.

Single GPU::

    CUDA_VISIBLE_DEVICES=3 python -m deg.train --out /data/deg/runs/smoke --steps 300

Multi-GPU (manual gradient all-reduce, one data shard per rank)::

    torchrun --nproc-per-node 4 -m deg.train --out /data/deg/runs/full ...
"""

from __future__ import annotations

import argparse
import bisect
import dataclasses
import gc
import json
import math
import os
import random
import re
import subprocess
import threading
import time
from collections import defaultdict, deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

import torch
import torch.nn.functional as F

from deg.infer import (
    Anchor,
    ArrayInstance,
    LeafInstance,
    anchor_bounds,
    anchor_count_bounds,
    array_query,
    base_generate_header,
    count_bounds,
    field_query,
    generate_header,
    is_anchored,
    make_anchors,
    walk,
)
from deg import ground as dground
from deg.media import is_media_record, text_view
from deg.memplan import GIB, MemCoef, MemPlanner, group_forks, static_bytes
from deg.model import (
    ARRAY_MODES,
    BACKBONES,
    GEN_LAYOUTS,
    KINDS,
    GemmaBackbone,
    backbone_spec,
    MAX_COUNT,
    DegModel,
    EncodedContext,
    EncodedQuery,
    Query,
    State,
    load_model,
    read_meta,
    region_bounds,
    span_rel_bias,
    token_span,
)
from deg.request import FillOptions, render_context
from deg.zero import ZeroMaster, broadcast_module, sharded_static_bytes
from deg.schema import FieldPlan, compile_schema

DEG_RECORDS = Path("/data/deg/datagen/full-off/records.jsonl")
FILLER_CORPUS = Path("/data/deg/external/v3/docs.jsonl")
JEV_TRAIN = Path("/data/jev/dataset/train.jsonl")
JEV_DEV = Path("/data/jev/dataset/dev.jsonl")


# --------------------------------------------------------------------------------- data


@dataclass
class Item:
    kind: str  # choice | score | noul | count | anchor | extract | generate
    seq: tuple[int, ...]
    query: EncodedQuery | None = None
    target: list[float] | None = None  # decide soft target / count one-hot
    bounds: tuple[int, int] | None = None  # count
    anchors: list[int] = field(default_factory=list)  # anchor: gold item start tokens
    starts: list[int] = field(default_factory=list)  # extract positives (token idx; -1=null)
    ends: list[int] = field(default_factory=list)
    weights: list[float] | None = None
    prompt_len: int = 0  # generate
    ptr: str = ""
    aux: bool = False  # count: the auxiliary count of an anchored array (anchor_count loss)
    media: bool = False  # extract: the gold pointer is a media region (deg.ground)


@dataclass
class Example:
    source: str  # deg | jev
    uid: str
    ctx: EncodedContext
    items: list[Item]
    anchored_arrays: int = 0
    skipped_arrays: int = 0  # anchored arrays without complete gold anchors (not trained)
    window: tuple[int, int] | None = None  # context chars [a, b) of a cropped window
    dropped: int = 0  # targets left out because the window does not show their gold
    subsampled: int = 0  # forks left out of this draw by :func:`cap_forks`
    rec: dict[str, Any] | None = field(default=None, repr=False)  # source (filler aug)
    aug: str = ""  # "" | "filler" (record inside filler text) | "negative" (filler only)

    @property
    def tokens(self) -> int:
        return len(self.ctx.ids) + sum(len(it.seq) for it in self.items)


def get_at(value: Any, parts: tuple) -> tuple[bool, Any]:
    for part in parts:
        if isinstance(part, int):
            if not isinstance(value, list) or part >= len(value):
                return False, None
            value = value[part]
        else:
            if not isinstance(value, dict) or part not in value:
                return False, None
            value = value[part]
    return True, value


def value_index(fld: FieldPlan, value: Any) -> int:
    if fld.mode == "decide.noul":
        return 0 if value else 1
    if fld.mode == "decide.choice":
        assert fld.options is not None
        return next(i for i, o in enumerate(fld.options)
                    if type(o) is type(value) and o == value)
    assert fld.levels is not None
    return len(fld.levels) if value is None else fld.levels.index(value)


def _onehot(n: int, i: int) -> list[float]:
    out = [0.0] * n
    out[i] = 1.0
    return out


def _shuffled(q: Query, target: list[float], rng: random.Random, p: float
              ) -> tuple[Query, list[float]]:
    if q.kind != "choice" or rng.random() >= p:
        return q, target
    order = list(range(len(q.options)))
    rng.shuffle(order)
    return (Query(q.kind, q.header, tuple(q.options[i] for i in order)),
            [target[i] for i in order])


def _query_item(model: DegModel, q: Query, target: list[float] | None, ptr: str,
                **kw: Any) -> Item:
    eq = model.renderer.query(q)
    return Item(q.kind, eq.ids, eq, target, ptr=ptr, **kw)


def primary_spans(rec: dict[str, Any]) -> dict[str, tuple[int, int]]:
    """Concrete pointer -> its gold char span: ``spans``, else the first
    ``span_alternatives`` entry (records may carry only alternatives for a field)."""
    out = {k: (int(v[0][0]), int(v[0][1])) for k, v in (rec.get("span_alternatives") or {})
           .items() if v}
    out.update({k: (int(v[0]), int(v[1])) for k, v in (rec.get("spans") or {}).items()})
    return out


def gold_item_anchors(ptr: str, n: int, spans: dict[str, tuple[int, int]],
                      ectx: EncodedContext, shift: int = 0, limit: int | None = None
                      ) -> list[int | None]:
    """Gold start token of each of the ``n`` items of the array instance at ``ptr`` (gold
    order): the first token of the earliest span among the item's leaves that no sibling
    item shares (a shared span, e.g. a table's column header, does not say where an item
    is), else of its earliest span; ``None`` when the item has no span. ``ectx`` may encode
    a window of ``limit`` chars starting at char ``shift`` of the document the spans index
    (the anchor span is still chosen on the whole document); an anchor outside the window
    is ``None``."""
    out: list[int | None] = []
    for sp in gold_item_anchor_spans(ptr, n, spans):
        if sp is None or (limit is not None and not 0 <= sp[0] - shift < limit):
            out.append(None)
            continue
        a, b = sp[0] - shift, sp[1] - shift
        try:
            out.append(token_span(ectx.offsets, a, b if limit is None else min(b, limit))[0])
        except (StopIteration, ValueError):
            out.append(None)
    return out


def gold_item_anchor_spans(ptr: str, n: int, spans: dict[str, tuple[int, int]]
                           ) -> list[tuple[int, int] | None]:
    """The char span locating each item (see :func:`gold_item_anchors`)."""
    per_item: list[list[tuple[int, int]]] = []
    for i in range(n):
        pre = f"{ptr}/{i}"
        per_item.append(sorted({v for k, v in spans.items()
                                if (k == pre or k.startswith(pre + "/")) and v[1] > v[0]}))
    users: dict[tuple[int, int], int] = defaultdict(int)
    for own in per_item:
        for sp in own:
            users[sp] += 1
    out: list[tuple[int, int] | None] = []
    for own in per_item:
        unshared = [sp for sp in own if users[sp] == 1]
        out.append((unshared or own)[0] if own else None)
    return out


# An item of a cropped window whose anchor lies outside the window: walked (indices stay
# the gold's) but nothing inside it trains.
OUTSIDE = Anchor(-1, -1, "")


def _outside(items: tuple) -> bool:
    return any(len(loc) > 2 and loc[2] is OUTSIDE for loc in items)


def record_text(rec: dict[str, Any]) -> str:
    return rec["context"] if isinstance(rec["context"], str) else render_context(rec["context"])


def judged_ok(rec: dict[str, Any], ptr: str) -> bool:
    """The field's open-judge verdict (gen-v1 ``judge``) is faithful 2 AND responsive 2
    and not dropped (array items take the array's verdict)."""
    judge = rec.get("judge") or {}
    key = ptr
    while key not in judge and re.search(r"/\d+$", key):
        key = re.sub(r"/\d+$", "", key)
    j = judge.get(key)
    return bool(j) and j.get("faithful") == 2 and j.get("responsive") == 2 \
        and j.get("kept", True) is not False


def deg_example(model: DegModel, rec: dict[str, Any], rng: random.Random,
                  shuffle_p: float = 0.5, window: tuple[int, int] | None = None,
                  anchor_count: bool = False, base_layout: bool | None = None,
                  judge_filter: bool = True, media: Any = None) -> Example:
    """One training example over the record's context, or over the context chars
    ``window = (a, b)`` only. A window that is not the whole context is **cropped**: it
    trains only what its own text shows. Extract fields keep the gold spans inside the
    window (none inside: the field is dropped, never trained as null); absent-gold (null)
    extract fields, decide and generate fields (their evidence has no span, so it may lie
    outside) and count heads are dropped; an anchored array trains its anchor fork on the
    in-window gold item starts when at least one lies inside (else it is dropped), and the
    items whose anchor lies outside are dropped with their whole subtree. Every dropped
    target is counted in ``Example.dropped``.

    ``anchor_count``: every anchored array of a whole context also trains the count head
    (:func:`anchor_count_item`; the ``anchor_count`` loss), so count-then-peaks decoding
    has a trained count; a cropped window never does (a count needs the whole array).

    ``base_layout``: generate forks in the hybrid's layout (KNOWN VALUES + the base header
    of :func:`deg.infer.base_generate_header`, no ``VALUE:`` line), as the base-model
    generator reads them (:mod:`deg.gentrain`); None = the model's own
    ``gen_layout`` (``chat``: deg writes with its own weights in that layout).

    ``judge_filter``: a record with an open-judge verdict map (gen-v1 ``judge``) trains
    only the x-generate fields the judge kept (:func:`judged_ok`); the others are counted
    in ``dropped`` (off: every field; :mod:`deg.gentrain` filters and counts itself).

    ``media`` (a :class:`deg.media.MediaEncoder`): a media record (``<media:ID>``
    placeholders, :func:`deg.media.is_media_record`) renders its items as soft-token spans
    and trains on its text coordinates (:func:`deg.media.text_view`); whole contexts only.
    An extract leaf of a media record whose value is in a media item has a gold
    ``grounds[ptr]`` (``{"media", "box": [x0, y0, x1, y1]}`` source pixels or ``{"media",
    "time": [s, e]}`` seconds; None = not in the media): its pointer targets the region's
    top-left / first and bottom-right / last soft token (:func:`deg.ground.
    ground_target`). A present value with neither a text span nor a ground entry is not
    trained (never as null)."""
    parts = None
    if is_media_record(rec):
        if media is None:
            raise ValueError("media record without a media encoder (train: --media-cache)")
        if window is not None:
            raise ValueError("media records train whole contexts only (no windows)")
        parts = media.parts(rec)
        rec = text_view(rec)
    if base_layout is None:
        base_layout = model.gen_layout == "chat"
    plan = compile_schema(rec["schema"])
    full_text = record_text(rec)
    wa, wb = window if window is not None else (0, len(full_text))
    if parts is None and not 0 <= wa < wb <= len(full_text):  # media: text may be empty
        raise ValueError(f"window {window} outside the context (0, {len(full_text)})")
    cropped = (wa, wb) != (0, len(full_text))
    ctx_text = full_text[wa:wb]
    rend = model.renderer
    ectx = rend.context(parts if parts is not None else ctx_text)
    T = len(ectx.offsets)
    value = rec["value"]
    opts = FillOptions()
    # a cropped window belongs to a document longer than any window: never routed to index
    mode = model.array_mode if cropped else model.array_mode_for(T)

    def inside(sp: Any) -> bool:  # a whole-context example keeps every gold span as is
        return not cropped or (wa <= int(sp[0]) and int(sp[1]) <= wb)

    def shifted(sp: Any) -> tuple[int, int]:
        return int(sp[0]) - wa, int(sp[1]) - wa

    spans = {k: shifted(v) for k, v in (rec.get("spans") or {}).items() if inside(v)}
    all_spans = primary_spans(rec)  # whole-document chars: anchors are chosen there
    counts: dict[str, int] = {}
    anchors: dict[str, Any] = {}
    items: list[Item] = []
    n_anchored = n_skipped = dropped = 0
    while True:
        pending = [x for x in walk(plan.root, counts, anchors=anchors)
                   if isinstance(x, ArrayInstance) and x.ptr not in counts]
        if not pending:
            break
        for inst in pending:
            if _outside(inst.items):
                counts[inst.ptr] = 0  # inside an item outside the window
                continue
            ok, got = get_at(value, inst.parts)
            n = len(got) if ok and isinstance(got, list) else 0
            if is_anchored(mode, inst):
                lo, hi = anchor_bounds(inst, opts)
                r_lo, r_hi = inst.region(T)
                if not cropped:
                    toks = gold_item_anchors(inst.ptr, n, all_spans, ectx)
                    if (any(t is None or not r_lo <= t < r_hi for t in toks)
                            or len(set(toks)) != n
                            or not (lo <= n and (hi is None or n <= hi))):
                        counts[inst.ptr] = 0  # no complete gold anchors: subtree not trained
                        n_skipped += 1
                        continue
                    anchors[inst.ptr] = make_anchors(toks, r_hi, ctx_text, ectx)  # type: ignore[arg-type]
                    counts[inst.ptr] = n
                    n_anchored += 1
                    q = array_query(plan, inst, anchored=True)
                    items.append(_query_item(model, q, None, inst.ptr,
                                             anchors=sorted(toks)))  # type: ignore[arg-type]
                    if anchor_count:
                        aux = anchor_count_item(model, plan, inst, n, opts)
                        if aux is not None:
                            items.append(aux)
                    continue
                located = gold_item_anchor_spans(inst.ptr, n, all_spans)
                toks = gold_item_anchors(inst.ptr, n, all_spans, ectx, shift=wa,
                                         limit=wb - wa)
                seen_toks = [t for t in toks if t is not None]
                if (any(sp is None for sp in located)
                        or any(not r_lo <= t < r_hi for t in seen_toks)
                        or len(set(seen_toks)) != len(seen_toks)
                        or not (lo <= n and (hi is None or n <= hi))):
                    counts[inst.ptr] = 0  # no complete gold anchors: subtree not trained
                    n_skipped += 1
                    continue
                if not seen_toks:  # no item starts in the window: never "empty" by cropping
                    counts[inst.ptr] = 0
                    dropped += 1
                    continue
                shown = make_anchors(seen_toks, r_hi, ctx_text, ectx)
                it_shown = iter(shown)
                anchors[inst.ptr] = [OUTSIDE if t is None else next(it_shown) for t in toks]
                counts[inst.ptr] = n
                n_anchored += 1
                dropped += n - len(seen_toks)
                q = array_query(plan, inst, anchored=True)
                items.append(_query_item(model, q, None, inst.ptr, anchors=sorted(seen_toks)))
                continue
            lo, hi = count_bounds(inst, opts)
            counts[inst.ptr] = n
            if cropped:
                dropped += 1  # a count needs the whole array in view
                continue
            if lo <= n <= hi:
                items.append(_query_item(model, array_query(plan, inst),
                                         _onehot(MAX_COUNT + 1, n), inst.ptr, bounds=(lo, hi)))
    leaves = [x for x in walk(plan.root, counts, anchors=anchors)
              if isinstance(x, LeafInstance)]
    decide_probs = rec.get("decide_probs") or {}
    span_alts: dict[str, list[tuple[int, int]]] = {}
    span_probs = dict(rec.get("span_probs") or {})
    for k, alts in (rec.get("span_alternatives") or {}).items():
        keep = [i for i, sp in enumerate(alts) if inside(sp)]
        if len(keep) != len(alts):
            w = span_probs.pop(k, None)
            if w is not None and len(w) == len(alts):
                span_probs[k] = [w[i] for i in keep]
        span_alts[k] = [shifted(alts[i]) for i in keep]
    grounds = rec.get("grounds") or {}
    mspans = dground.media_spans(ectx, parts) if parts is not None else []
    known: list[tuple[str, Any]] = []
    for inst in leaves:
        fld = inst.field
        if fld.phase == "parallel":
            present, v = get_at(value, inst.parts)
            if present:
                known.append((inst.ptr, v))
    for inst in leaves:
        fld = inst.field
        present, v = get_at(value, inst.parts)
        if cropped:
            if _outside(inst.items) or not fld.mode.startswith("extract.") or not (
                    inst.ptr in spans or span_alts.get(inst.ptr)):
                dropped += 1  # gold not shown by this window (or no span to show)
                continue
        if fld.mode.startswith("decide."):
            q = field_query(fld, inst.qptr, list(inst.items))
            soft = decide_probs.get(inst.ptr)
            target = soft if soft is not None else _onehot(
                len(q.options) if q.options else 2, value_index(fld, v))
            q, target = _shuffled(q, list(target), rng, shuffle_p)
            items.append(_query_item(model, q, target, inst.ptr))
        elif fld.mode.startswith("extract."):
            q = field_query(fld, inst.qptr, list(inst.items))
            item = _query_item(model, q, None, inst.ptr)
            if (parts is not None and inst.ptr not in spans and not span_alts.get(inst.ptr)
                    and (inst.ptr in grounds or present)):
                tok = media_target(grounds, inst.ptr, mspans)
                if tok is None:
                    dropped += 1  # in the media but not located: no target
                    continue
                item.starts, item.ends, item.media = [tok[0]], [tok[1]], tok[0] >= 0
                items.append(item)
                continue
            if inst.ptr in spans or span_alts.get(inst.ptr):
                alts = span_alts.get(inst.ptr) or [spans[inst.ptr]]
                for a, b in alts:
                    i, j = token_span(ectx.offsets, a, b)
                    item.starts.append(i)
                    item.ends.append(j)
                w = span_probs.get(inst.ptr)
                item.weights = list(w) if w is not None and len(w) == len(alts) else None
            else:
                item.starts, item.ends = [-1], [-1]
            items.append(item)
        else:  # generate
            if not present or not isinstance(v, str):
                continue
            if judge_filter and "judge" in rec and not judged_ok(rec, inst.ptr):
                dropped += 1  # an open-judge record trains only the fields it kept
                continue
            # same layout as infer.generate_batch: decided/extracted values only (generate
            # fields are independent of each other), leaf order
            if base_layout:  # infer.fill_many with model.generator() base_layout
                prompt = rend.known_block(known) + rend.enc(
                    base_generate_header(fld, inst.ptr, list(inst.items), rend.gen_turn))
            else:
                prompt = rend.generate_prompt(known, generate_header(fld, inst.ptr,
                                                                     list(inst.items)))
            target = rend.enc(v) + [rend.eos_id]
            items.append(Item("generate", tuple(prompt + target), None, None,
                              prompt_len=len(prompt), ptr=inst.ptr))
    return Example("deg", rec["uid"], ectx, items, n_anchored, n_skipped,
                   (wa, wb) if cropped else None, dropped)


def media_target(grounds: dict[str, Any], ptr: str, spans: list[dground.MediaSpan]
                 ) -> tuple[int, int] | None:
    """Context-token (start, end) of a media extract leaf's gold ``grounds[ptr]``, (-1, -1)
    for a gold None (not in the media), None when there is no usable entry."""
    if ptr not in grounds:
        return None
    gold = grounds[ptr]
    if gold is None:
        return -1, -1
    by_id = {s.id: s for s in spans}
    span = by_id.get(str(gold.get("media"))) or (spans[0] if len(spans) == 1 else None)
    return dground.ground_target(span, gold) if span is not None else None


def anchor_count_item(model: DegModel, plan: Any, inst: ArrayInstance, n: int,
                      opts: FillOptions) -> Item | None:
    """Auxiliary count fork of an anchored array (the index-mode count query): target
    ``min(n, hi)`` over the bounds of :func:`deg.infer.anchor_count_bounds` (the count
    head's capacity :data:`MAX_COUNT` caps it: the top class means "at least hi");
    ``None`` when minItems exceeds the head."""
    lo, hi = anchor_count_bounds(inst, opts)
    if lo > hi:
        return None
    return _query_item(model, array_query(plan, inst), _onehot(MAX_COUNT + 1, min(n, hi)),
                       inst.ptr, bounds=(lo, hi), aux=True)


# --------------------------------------------------------------- filler augmentation


class FillerCorpus:
    """Unrelated long text for filler augmentation: the v3 document corpus
    (``<dir>/docs.jsonl`` + ``text/``; SEC EDGAR + CC-BY/CC0 arXiv, deduplicated against
    ExtractBench at fetch time). Only docs in the v3 **train** split
    (:func:`deg.datagen.v3.split_of`) are used, so v3 val documents never appear in
    training text. Never point this at ExtractBench."""

    def __init__(self, path: Path, cache: int = 64) -> None:
        from deg.datagen.v3 import split_of

        path = Path(path)
        docs = path / "docs.jsonl" if path.is_dir() else path
        if "extractbench" in str(docs).lower():
            raise ValueError(f"filler corpus must not be ExtractBench: {docs}")
        self.docs: list[tuple[str, Path]] = []
        for row in read_jsonl(docs):
            tp = Path(row.get("text_path") or docs.parent / "text" / f"{row['id']}.txt")
            if split_of(row["id"]) == "train" and tp.exists():
                self.docs.append((row["id"], tp))
        if not self.docs:
            raise ValueError(f"no train-split filler documents in {docs}")
        self._cache: dict[str, str] = {}
        self._order: deque[str] = deque()
        self._limit = cache

    def _read(self, doc_id: str, path: Path) -> str:
        if doc_id not in self._cache:
            self._cache[doc_id] = path.read_text(errors="replace")
            self._order.append(doc_id)
            while len(self._order) > self._limit:
                self._cache.pop(self._order.popleft(), None)
        return self._cache[doc_id]

    def text(self, rng: random.Random, chars: int, exclude: str | None = None) -> str:
        """About ``chars`` characters of contiguous text from a random document (more
        documents joined by a blank line when one is too short), never from ``exclude``."""
        parts: list[str] = []
        need = chars
        for _ in range(64):
            doc_id, path = rng.choice(self.docs)
            if doc_id == exclude:
                continue
            body = self._read(doc_id, path)
            if not body.strip():
                continue
            a = rng.randrange(max(1, len(body) - need)) if len(body) > need else 0
            piece = body[a:a + need]
            parts.append(piece)
            need -= len(piece) + 2
            if need <= 0:
                break
        return "\n\n".join(parts)


@dataclass(eq=False)
class FillerDraw:
    """A drawn short example to be rebuilt inside filler text of about ``tokens`` context
    tokens (``negative``: filler only, the record's top-level anchored arrays train empty);
    materialized when the batch packing reaches it, like a :class:`LongDoc` window."""

    ex: Example
    tokens: int
    negative: bool = False

    @property
    def est_len(self) -> int:
        return self.tokens


def _filler_cut(model: DegModel, seg: str, n_tokens: int) -> str:
    """The prefix of ``seg`` holding its first ``n_tokens`` tokens."""
    if n_tokens <= 0:
        return ""
    off = model.renderer.context(seg).offsets
    return seg if n_tokens >= len(off) else seg[:off[n_tokens][0]]


def filler_record(model: DegModel, rec: dict[str, Any], corpus: FillerCorpus,
                  rng: random.Random, filler_tokens: int) -> dict[str, Any]:
    """``rec`` with its document embedded at a random offset inside ``filler_tokens``
    tokens of unrelated text (``pre + "\n\n" + doc + "\n\n" + post``); every gold char
    span (``spans``, ``span_alternatives``) shifts by the inserted prefix. Values, decide
    labels and generate targets are unchanged: the filler is unrelated text."""
    text = record_text(rec)
    n_pre = rng.randint(0, filler_tokens)
    exclude = rec.get("doc_id")
    pre = _filler_cut(model, corpus.text(rng, 6 * n_pre + 64, exclude), n_pre)
    post = _filler_cut(model, corpus.text(rng, 6 * (filler_tokens - n_pre) + 64, exclude),
                       filler_tokens - n_pre)
    head = pre + "\n\n" if pre else ""
    shift = len(head)
    out = dict(rec)
    out["context"] = head + text + ("\n\n" + post if post else "")
    out["spans"] = {k: [int(v[0]) + shift, int(v[1]) + shift]
                    for k, v in (rec.get("spans") or {}).items()}
    out["span_alternatives"] = {k: [[int(a) + shift, int(b) + shift] for a, b in alts]
                                for k, alts in (rec.get("span_alternatives") or {}).items()}
    out["uid"] = f"{rec['uid']}#filler"
    return out


def negative_example(model: DegModel, rec: dict[str, Any], corpus: FillerCorpus,
                     rng: random.Random, tokens: int, anchor_count: bool = False
                     ) -> Example | None:
    """Filler text only (no record document) with ``rec``'s schema: every top-level
    anchored array whose minItems is 0 trains its anchor fork with no item start (all
    tokens negative) and, with ``anchor_count``, a count of 0. Nothing else trains (the
    filler is unrelated, but an extract field could still match it by chance). ``None``
    when the schema has no such array."""
    plan = compile_schema(rec["schema"])
    ctx_text = _filler_cut(model, corpus.text(rng, 6 * tokens + 64, rec.get("doc_id")),
                           tokens)
    if not ctx_text.strip():
        return None
    ectx = model.renderer.context(ctx_text)
    opts = FillOptions()
    items: list[Item] = []
    n_anchored = 0
    for x in walk(plan.root, {}, anchors={}):
        if not (isinstance(x, ArrayInstance) and not x.items
                and is_anchored(model.array_mode_for(len(ectx.offsets)), x)):
            continue
        if anchor_bounds(x, opts)[0] != 0:
            continue
        items.append(_query_item(model, array_query(plan, x, anchored=True), None, x.ptr))
        n_anchored += 1
        if anchor_count:
            aux = anchor_count_item(model, plan, x, 0, opts)
            if aux is not None:
                items.append(aux)
    if not items:
        return None
    return Example("deg", f"{rec['uid']}#negative", ectx, items, n_anchored, aug="negative")


def filler_example(model: DegModel, draw: FillerDraw, corpus: FillerCorpus,
                   rng: random.Random, max_context_tokens: int, anchor_count: bool = False
                   ) -> Example:
    """Materialize a :class:`FillerDraw`; the plain example when the rebuilt one would not
    fit ``max_context_tokens`` or trains nothing."""
    ex = draw.ex
    assert ex.rec is not None
    if draw.negative:
        neg = negative_example(model, ex.rec, corpus, rng,
                               min(draw.tokens, max_context_tokens - 32), anchor_count)
        if neg is not None and len(neg.ctx.ids) <= max_context_tokens:
            return neg
    filler = draw.tokens - len(ex.ctx.ids)
    for _ in range(3):
        if filler <= 0:
            break
        aug = deg_example(model, filler_record(model, ex.rec, corpus, rng, filler), rng,
                            anchor_count=anchor_count)
        if len(aug.ctx.ids) <= max_context_tokens:
            if aug.items:
                aug.aug = "filler"
                return aug
            break
        filler -= len(aug.ctx.ids) - max_context_tokens + 16
    return ex


# ------------------------------------------------------------------- long contexts


@dataclass(eq=False)
class LongDoc:
    """A train record whose context does not fit ``--max-context-tokens``. Each draw
    trains a fresh span-aware window (:func:`window_example`); inference is unaffected
    (it streams the whole context)."""

    uid: str
    rec: dict[str, Any]
    starts: list[int]  # whole-context token char offsets
    ends: list[int]
    text_len: int
    targets: list[tuple[int, int]]  # gold char spans a window is centred on
    window_tokens: int  # context tokens per window
    max_context_tokens: int

    @property
    def est_len(self) -> int:
        return self.max_context_tokens


def long_doc(model: DegModel, rec: dict[str, Any], max_context_tokens: int,
             margin: int = 16) -> LongDoc | None:
    """``None`` when no gold span exists (a cropped window could train nothing)."""
    text = record_text(rec)
    ectx = model.renderer.context(text)
    overhead = len(ectx.ids) - len(ectx.offsets)
    targets = sorted({(int(a), int(b)) for a, b in primary_spans(rec).values() if b > a}
                     | {(int(a), int(b)) for alts in (rec.get("span_alternatives") or {})
                        .values() for a, b in alts if b > a})
    if not targets:
        return None
    return LongDoc(rec["uid"], rec, [a for a, _ in ectx.offsets], [b for _, b in ectx.offsets],
                   len(text), targets, max(64, max_context_tokens - overhead - margin),
                   max_context_tokens)


def choose_window(doc: LongDoc, rng: random.Random, tokens: int | None = None
                  ) -> tuple[int, int]:
    """Context chars of a ``tokens``-token window holding a random gold span (uniform
    over the start positions that keep it inside; a span longer than the window starts
    the window)."""
    W = tokens or doc.window_tokens
    T = len(doc.starts)
    if T <= W:
        return 0, doc.text_len
    a, b = rng.choice(doc.targets)
    ta = min(bisect.bisect_right(doc.ends, a), T - 1)  # first token ending after a
    tb = max(bisect.bisect_left(doc.starts, b) - 1, ta)  # last token starting before b
    lo, hi = max(0, tb - W + 1), min(ta, T - W)
    s = rng.randint(lo, hi) if lo <= hi else min(ta, T - W)
    e = s + W
    # a gold start between two tokens (an untokenized gap) still opens the window
    wa = 0 if s == 0 else min(doc.starts[s], a if s == ta else doc.starts[s])
    return wa, (doc.text_len if e >= T else doc.starts[e])


def window_example(model: DegModel, doc: LongDoc, rng: random.Random, budget: int,
                   tries: int = 4, record_fork_tokens: int = 0) -> Example | None:
    """A cropped example of ``doc`` that fits ``max_context_tokens`` (re-tokenizing a window
    can shift its boundary tokens: shrink and retry) with its forks subsampled
    (:func:`cap_forks`) to fit ``budget`` tokens beside the trunk, and at least
    ``budget // 2`` fork tokens whatever
    the window length (a 16k window under a 12k budget still trains a field set like a
    6k one; the memory planner bounds the micro-batch); forks are independent loss
    terms. ``None`` when ``tries`` windows trained nothing."""
    W = doc.window_tokens
    for _ in range(tries):
        ex = deg_example(model, doc.rec, rng, window=choose_window(doc, rng, W))
        if len(ex.ctx.ids) > doc.max_context_tokens:
            W = max(32, W - (len(ex.ctx.ids) - doc.max_context_tokens) - 16)
            continue
        if not ex.items:
            continue
        return cap_forks(ex, rng, budget, record_fork_tokens)
    return None


WHOLE_ARRAY_KINDS = ("anchor", "count")  # one fork per array, never per item


def fork_room(ex: Example, budget: int, record_fork_tokens: int = 0) -> int:
    """Fork tokens one draw of ``ex`` may train: ``record_fork_tokens`` when set, else
    what the micro-batch ``budget`` leaves beside its trunk and at least ``budget // 2``."""
    if record_fork_tokens > 0:
        return record_fork_tokens
    return max(budget - len(ex.ctx.ids), budget // 2)


DECIDE_KINDS = ("choice", "score", "noul")


def field_of(ptr: str) -> str:
    """The schema field of a concrete pointer: array indices -> ``*``."""
    return "/".join("*" if part.isdigit() else part for part in ptr.split("/"))


def cap_decide_forks(ex: Example, rng: random.Random, per_field: int) -> Example:
    """At most ``per_field`` decide forks per schema field in one draw (a fresh random
    subset; ``0`` = no cap). Long filings repeat one in-array enum per table cell (v3-long:
    ~135 choice forks per record, 99.8% inside arrays, e.g. ``unit`` = usd / ``scale`` =
    thousands), which otherwise are most of the decide forks of a step that draws one
    (results/deg/regress-2b)."""
    if per_field <= 0:
        return ex
    groups: dict[str, list[int]] = defaultdict(list)
    for i, it in enumerate(ex.items):
        if it.kind in DECIDE_KINDS:
            groups[field_of(it.ptr)].append(i)
    drop: set[int] = set()
    for idx in groups.values():
        if len(idx) > per_field:
            drop.update(rng.sample(idx, len(idx) - per_field))
    if not drop:
        return ex
    return dataclasses.replace(ex, items=[it for i, it in enumerate(ex.items) if i not in drop],
                               subsampled=ex.subsampled + len(drop))


def cap_forks(ex: Example, rng: random.Random, budget: int, record_fork_tokens: int = 0,
              decide_per_field: int = 0) -> Example:
    """One draw's forks of ``ex`` (after :func:`cap_decide_forks` with
    ``decide_per_field``): all of them when they fit :func:`fork_room`, else a
    fresh random subset (a copy; the pooled example keeps every fork), so one record
    costs at most one trunk pass plus about ``room`` fork tokens per step whatever its
    field count (a long filing has ~1k field rows and each generate fork repeats the
    record's whole known-value block, ~16k tokens: millions of fork tokens unsubsampled).

    Whole-array targets (anchor forks and anchored arrays' count forks, one per array)
    are always kept; the remaining forks are taken in random order while they fit, and
    the first of them is taken whatever its length, so every fork (a generate fork longer
    than the room included) trains with positive probability. Each loss term stays a mean
    over the forks drawn (a subsample, as for windows; forks are independent loss terms),
    though forks longer than the room are drawn less often. Bound: ``room`` + the
    whole-array forks + one fork."""
    ex = cap_decide_forks(ex, rng, decide_per_field)
    room = fork_room(ex, budget, record_fork_tokens)
    if sum(len(it.seq) for it in ex.items) <= room:
        return ex
    whole = [i for i, it in enumerate(ex.items) if it.kind in WHOLE_ARRAY_KINDS]
    rest = [i for i, it in enumerate(ex.items) if it.kind not in WHOLE_ARRAY_KINDS]
    rng.shuffle(rest)
    keep = list(whole)
    used = sum(len(ex.items[i].seq) for i in whole)
    for n, i in enumerate(rest):
        size = len(ex.items[i].seq)
        if n and used + size > room:
            continue
        keep.append(i)
        used += size
    keep.sort()
    return dataclasses.replace(ex, items=[ex.items[i] for i in keep],
                               subsampled=ex.subsampled + len(ex.items) - len(keep))


def cap_record_parts(planner: MemPlanner, batch: list[Example], max_parts: int,
                     rng: random.Random) -> tuple[list[Example], int]:
    """At most ``max_parts`` micro-batches (parts) per record in one step: a record the
    planner would split by rows (:meth:`MemPlanner.split_rows`) into more chunks keeps a
    fresh random ``max_parts`` of them (their items; a subsample, like :func:`cap_forks`),
    so one pathological record (deg2b-1 step 2960: a long filing split into ~300 parts,
    ~30 min for one step while the other ranks waited) costs at most ``max_parts`` part
    passes. ``0`` = no cap. Returns the batch and the number of forks left out."""
    if max_parts <= 0:
        return batch, 0
    lim = planner.budget(True)
    out: list[Example] = []
    left = 0
    for ex in batch:
        if planner.estimate([ex]) <= lim:
            out.append(ex)
            continue
        chunks = planner.split_rows(ex)
        if len(chunks) <= max_parts:
            out.append(ex)
            continue
        keep = sorted(rng.sample(range(len(chunks)), max_parts))
        items = [it for i in keep for it in chunks[i].items]
        left += len(ex.items) - len(items)
        out.append(dataclasses.replace(ex, items=items,
                                       subsampled=ex.subsampled + len(ex.items) - len(items)))
    return out, left


def jev_example(model: DegModel, rec: dict[str, Any], rng: random.Random,
                shuffle_p: float = 0.5) -> Example | None:
    ctx_text = rec["state"] if isinstance(rec["state"], str) else render_context(rec["state"])
    ectx = model.renderer.context(ctx_text)
    items = []
    for n, raw in enumerate(rec["questions"]):
        label = raw.get("label") or {}
        kind = raw["type"]
        lines = ["FIELD /answer", f"DESCRIPTION: {raw['instructions']}"]
        options: tuple[str, ...] = ()
        if kind == "noul":
            crit = raw.get("criteria") or {}
            if crit.get("true"):
                lines.append(f"Yes means: {crit['true']}")
            if crit.get("false"):
                lines.append(f"No means: {crit['false']}")
            lines.append("KIND: yes or no (true or false)")
            size = 2
        elif kind == "choice":
            options = tuple(f"{k}: {d}" if d else k for k, d in raw["criteria"].items())
            lines.append("KIND: choose one option")
            size = len(options)
        else:
            options = tuple(f"{i}: {c}" for i, c in enumerate(raw["criteria"]))
            lines.append("KIND: choose one level of an integer scale")
            size = len(options)
        if label.get("probs") is not None:
            target = [float(p) for p in label["probs"]]
        elif label.get("hard") is not None:
            target = _onehot(size, int(label["hard"]))
        else:
            continue
        if len(target) != size:
            continue
        q, target = _shuffled(Query(kind, "\n".join(lines) + "\n", options), target, rng,
                              shuffle_p)
        items.append(_query_item(model, q, target, f"/q{n}"))
    return Example("jev", rec["uid"], ectx, items) if items else None


_SUITE_GUARD: Any = None
_SUITE_CHECKED: dict[str, bool] = {}  # uid -> contaminated (records are re-read each reload)


def suite_guard() -> Any:
    """deg-suite v1 items are eval only (:mod:`deg.bench.suite`): built once."""
    global _SUITE_GUARD
    if _SUITE_GUARD is None:
        from deg.bench.suite import SuiteGuard

        _SUITE_GUARD = SuiteGuard()
    return _SUITE_GUARD


def read_jsonl(path: Path, split: str | None = None, limit: int | None = None,
               pos: list[int] | None = None) -> list[dict[str, Any]]:
    """Records of ``split``. Reading the ``train`` split (training input) refuses suite
    files and drops records / jev examples containing a deg-suite context. ``pos``
    ([byte offset]): read only complete lines from that offset on and advance it (a
    file a datagen job is still appending to is re-read incrementally)."""
    guard = None
    if split == "train":
        from deg.bench.suite import guard_path

        guard_path(path)
        guard = suite_guard()
    out, dropped = [], 0

    def lines():
        if pos is None:
            with path.open() as fh:
                yield from fh
            return
        with path.open("rb") as fb:
            fb.seek(pos[0])
            for raw in fb:
                if not raw.endswith(b"\n"):  # a line still being written
                    break
                pos[0] += len(raw)
                yield raw.decode()

    for line in lines():
        if not line.strip():
            continue
        rec = json.loads(line)
        if split is None or rec.get("split") == split:
            if guard is not None:
                key = f"{path}:{rec.get('uid')}"
                if key not in _SUITE_CHECKED or rec.get("uid") is None:
                    _SUITE_CHECKED[key] = guard.match(
                        record_text(rec) if "context" in rec else rec.get("state", "")
                    ) is not None
                if _SUITE_CHECKED[key]:
                    dropped += 1
                    continue
            out.append(rec)
            if limit is not None and len(out) >= limit:
                break
    if guard is not None:
        print(json.dumps({"event": "suite_guard", "path": str(path), "texts": guard.loaded,
                          "dropped": dropped}), flush=True)
    return out


# ------------------------------------------------------------------------------- losses


def _group_forks(forks: list[tuple[int, Item]], fork_tokens: int, fork_rows: int
                 ) -> list[list[tuple[int, Item]]]:
    """See :func:`deg.memplan.group_forks` (the memory planner groups the same way)."""
    return group_forks(forks, fork_tokens, fork_rows)


def _kl(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """Row-wise KL(t || softmax(logits)); -inf logits must have t = 0."""
    logp = torch.log_softmax(logits.float(), -1)
    t = targets / targets.sum(-1, keepdim=True)
    tlogt = torch.where(t > 0, t * t.clamp_min(1e-30).log(), torch.zeros_like(t))
    return (tlogt - torch.where(t > 0, t * logp, torch.zeros_like(t))).sum(-1)


def forward_losses(model: DegModel, examples: list[Example], fork_tokens: int,
                   fork_rows: int, trunk: tuple[torch.Tensor, State] | None = None
                   ) -> tuple[dict[str, torch.Tensor], dict[str, float]]:
    """One micro-batch: trunks, then forks in length groups, then every head batched.
    ``trunk`` = the examples' already computed ``model.trunk`` output (row chunks of one
    record sharing one trunk pass, :func:`run_chunks`)."""
    device = model.device
    heads = model.heads
    trunk_hidden, state = trunk if trunk is not None else model.trunk(
        [ex.ctx for ex in examples])
    # LoRA: x-generate decodes adapter-off (= the frozen base) at inference, so generate
    # forks carry no trainable path and are not run
    forks = [(r, it) for r, ex in enumerate(examples) for it in ex.items
             if not (model.lora is not None and it.kind == "generate")]
    stats = {"trunk_tokens": float(sum(len(ex.ctx.ids) for ex in examples)),
             "trunk_padded": float(len(examples) * max(len(ex.ctx.ids) for ex in examples)),
             "fork_tokens": float(sum(len(it.seq) for _, it in forks)), "fork_padded": 0.0,
             "lm_targets": 0.0}
    answers: dict[str, list[torch.Tensor]] = defaultdict(list)
    options: dict[str, list[torch.Tensor]] = defaultdict(list)
    items: dict[str, list[tuple[int, Item]]] = defaultdict(list)
    lm_h: list[torch.Tensor] = []
    lm_y: list[int] = []
    lm_items: list[tuple[int, Item]] = []
    # a D/E/G task MoE runs one expert per fork call: group within each route
    routes = sorted({model.fork_route(it.kind) or "" for _, it in forks})
    groups = [g for rt in routes for g in _group_forks(
        [f for f in forks if (model.fork_route(f[1].kind) or "") == rt], fork_tokens, fork_rows)]
    for group in groups:
        rows = torch.tensor([r for r, _ in group], dtype=torch.long, device=device)
        hidden = model.run_forks(state, rows, [it.seq for _, it in group],
                                 kinds=[it.kind for _, it in group])
        stats["fork_padded"] += len(group) * max(len(it.seq) for _, it in group)
        for (r, it), h in zip(group, hidden):
            if it.kind == "generate":
                lm_h.append(h[it.prompt_len - 1: len(it.seq) - 1])
                lm_y += it.seq[it.prompt_len:]
                lm_items.append((r, it))
                continue
            assert it.query is not None
            answers[it.kind].append(h[it.query.answer_position])
            if it.kind in ("choice", "score"):
                options[it.kind].append(h[list(it.query.option_positions)])
            items[it.kind].append((r, it))
    terms: dict[str, torch.Tensor] = {}
    decide_terms: list[torch.Tensor] = []
    for kind in ("choice", "score"):
        if not items[kind]:
            continue
        a = torch.stack(answers[kind]).float()
        opts = torch.nn.utils.rnn.pad_sequence(options[kind], batch_first=True).float()
        n_opt = torch.tensor([len(o) for o in options[kind]], device=device)
        mask = torch.arange(opts.shape[1], device=device)[None] < n_opt[:, None]
        head = heads.choice if kind == "choice" else heads.score
        logits = head(a, opts).masked_fill(~mask, float("-inf"))
        tgt = torch.zeros_like(logits)
        for i, (_, it) in enumerate(items[kind]):
            tgt[i, :len(it.target)] = torch.tensor(it.target, device=device).clamp_min(0)
        kl = _kl(logits, tgt)
        decide_terms.append(kl)
        terms[f"decide_{kind}"] = kl.mean().detach()
    if items["noul"]:
        lg = heads.noul(torch.stack(answers["noul"]).float())
        logits = torch.cat([lg, torch.zeros_like(lg)], -1)
        tgt = torch.tensor([it.target for _, it in items["noul"]], device=device).clamp_min(0)
        kl = _kl(logits, tgt)
        decide_terms.append(kl)
        terms["decide_noul"] = kl.mean().detach()
    if decide_terms:
        terms["decide"] = torch.cat(decide_terms).mean()
    if items["count"]:
        logits = heads.count(torch.stack(answers["count"]).float())
        keep = torch.zeros_like(logits, dtype=torch.bool)
        y = []
        aux = torch.tensor([it.aux for _, it in items["count"]], device=device)
        for i, (_, it) in enumerate(items["count"]):
            assert it.bounds is not None and it.target is not None
            keep[i, it.bounds[0]: it.bounds[1] + 1] = True
            y.append(max(range(len(it.target)), key=it.target.__getitem__))
        ce = F.cross_entropy(logits.masked_fill(~keep, float("-inf")),
                             torch.tensor(y, device=device), reduction="none")
        # index-addressed arrays -> count; anchored arrays' auxiliary counts -> anchor_count
        for key, sel in (("count", ~aux), ("anchor_count", aux)):
            if bool(sel.any()):
                terms[key] = ce[sel].mean()
    if items["anchor"]:
        a_all = torch.stack(answers["anchor"]).float()
        anc_rows: dict[int, list[int]] = defaultdict(list)
        for i, (r, _) in enumerate(items["anchor"]):
            anc_rows[r].append(i)
        anc_terms = []
        for r, idx in anc_rows.items():
            ex = examples[r]
            ctx_h = trunk_hidden[r, ex.ctx.ctx_start: ex.ctx.ctx_start + len(ex.ctx.offsets)]
            logits = heads.anchor(a_all[idx], ctx_h.float())
            T = logits.shape[1]
            for k, i in enumerate(idx):
                it = items["anchor"][i][1]
                assert it.query is not None
                lo, hi = region_bounds(it.query.region, T)
                y = torch.zeros(hi - lo, device=device)
                if it.anchors:
                    y[torch.tensor(it.anchors, device=device) - lo] = 1.0
                bce = F.binary_cross_entropy_with_logits(logits[k, lo:hi], y, reduction="sum")
                anc_terms.append(bce / (1 + len(it.anchors)))
        terms["anchor"] = torch.stack(anc_terms).mean()
    if items["extract"]:
        a_all = torch.stack(answers["extract"]).float()
        by_row: dict[int, list[int]] = defaultdict(list)
        for i, (r, _) in enumerate(items["extract"]):
            by_row[r].append(i)
        ext_terms = []
        for r, idx in by_row.items():
            ex = examples[r]
            ctx_h = trunk_hidden[r, ex.ctx.ctx_start: ex.ctx.ctx_start + len(ex.ctx.offsets)]
            start, end = heads.span(a_all[idx], ctx_h.float())
            T = start.shape[1] - 1
            locs = [items["extract"][i][1].query for i in idx]
            if any(q is not None and q.anchor >= 0 for q in locs):
                bias = torch.zeros(2, len(idx), T + 1, device=device)
                for k, q in enumerate(locs):
                    if q is not None and q.anchor >= 0:
                        bs, be = span_rel_bias(heads.span_bias, q.anchor,
                                               region_bounds(q.region, T)[1], 0, T)
                        bias[0, k, :T] = bs
                        bias[1, k, :T] = be
                start, end = start + bias[0], end + bias[1]
            per_item = torch.zeros(len(idx), device=device)
            for logits, which in ((start, "starts"), (end, "ends")):
                logp = torch.log_softmax(logits, -1)
                w = torch.zeros_like(logp)
                pos = torch.zeros_like(logp, dtype=torch.bool)
                soft = torch.zeros(len(idx), dtype=torch.bool, device=device)
                for k, i in enumerate(idx):
                    it = items["extract"][i][1]
                    tok = [T if p < 0 else p for p in getattr(it, which)]
                    pos[k, tok] = True
                    if it.weights is not None:
                        soft[k] = True
                        ww = torch.tensor(it.weights, device=device, dtype=logp.dtype)
                        w[k].index_put_((torch.tensor(tok, device=device),), ww / ww.sum(),
                                        accumulate=True)
                marginal = -torch.logsumexp(logp.masked_fill(~pos, float("-inf")), -1)
                weighted = -(w * logp).sum(-1)
                per_item = per_item + torch.where(soft, weighted, marginal) / 2
            ext_terms.append(per_item)
        terms["extract"] = torch.cat(ext_terms).mean()
        in_media = torch.tensor([items["extract"][i][1].media for idx in by_row.values()
                                 for i in idx], device=device)
        if bool(in_media.any()):  # monitor only: media-region pointers of the extract term
            terms["extract_media"] = torch.cat(ext_terms)[in_media].mean().detach()
    if lm_h:
        logits = model.backbone.lm_logits(torch.cat(lm_h)).float()
        terms["lm"] = F.cross_entropy(logits, torch.tensor(lm_y, device=device))
        stats["lm_targets"] = float(len(lm_y))
        if model.base is not None:
            terms["lm_kl_base"] = base_kl(model, examples, lm_items, logits, fork_tokens,
                                          fork_rows)
    return terms, stats


KL_ROWS = 512  # target positions per KL block (bounds the base's fp32 [rows, vocab])


def base_kl(model: DegModel, examples: list[Example], lm_items: list[tuple[int, Item]],
            logits: torch.Tensor, fork_tokens: int, fork_rows: int) -> torch.Tensor:
    """Mean over the x-generate target positions of KL(p_base || softmax(logits)).
    ``logits`` [N, V] = deg's at those positions, in ``lm_items`` order. The frozen base
    (``model.base``) reads the same contexts and fork tokens under no_grad: one base
    trunk pass over the examples that have generate items, then their generate forks in
    the trainer's length groups."""
    view = model.base_view()
    used = sorted({r for r, _ in lm_items})
    local = {r: i for i, r in enumerate(used)}
    hid: dict[int, torch.Tensor] = {}
    with torch.no_grad():
        _, bstate = view.trunk([examples[r].ctx for r in used])
        for group in _group_forks([(local[r], it) for r, it in lm_items], fork_tokens,
                                  fork_rows):
            rows = torch.tensor([r for r, _ in group], dtype=torch.long, device=model.device)
            hidden = view.run_forks(bstate, rows, [it.seq for _, it in group])
            for (_, it), h in zip(group, hidden):
                hid[id(it)] = h[it.prompt_len - 1: len(it.seq) - 1]
        del bstate
        bh = torch.cat([hid[id(it)] for _, it in lm_items])
    total = logits.new_zeros(())
    for a in range(0, logits.shape[0], KL_ROWS):
        with torch.no_grad():
            blogp = torch.log_softmax(view.backbone.lm_logits(bh[a:a + KL_ROWS]).float(), -1)
        logq = torch.log_softmax(logits[a:a + KL_ROWS], -1)
        total = total + (blogp.exp() * (blogp - logq)).sum()
    return total / logits.shape[0]


# --------------------------------------------------------------------------- telemetry


class GpuMonitor:
    """Samples nvidia-smi utilization/memory for the visible device every second."""

    def __init__(self) -> None:
        visible = os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")
        rank = int(os.environ.get("LOCAL_RANK", "0"))
        self.index = visible[rank] if rank < len(visible) else str(rank)
        self.samples: list[tuple[float, float]] = []
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def _run(self) -> None:
        while not self._stop.is_set():
            try:
                out = subprocess.run(
                    ["nvidia-smi", "-i", self.index, "--query-gpu=utilization.gpu,memory.used",
                     "--format=csv,noheader,nounits"], capture_output=True, text=True,
                    timeout=5).stdout.strip().split(",")
                self.samples.append((float(out[0]), float(out[1])))
            except Exception:  # noqa: BLE001 - telemetry only
                pass
            self._stop.wait(1.0)

    def take(self) -> dict[str, float]:
        got, self.samples = self.samples, []
        if not got:
            return {}
        return {"gpu_util": round(sum(u for u, _ in got) / len(got), 1),
                "gpu_mem_mib": max(m for _, m in got)}

    def stop(self) -> None:
        self._stop.set()


# -------------------------------------------------------------------------------- train


LOSS_KEYS = ("decide", "extract", "count", "anchor", "anchor_count", "lm", "lm_kl_base")


def loss_grad_norm(loss: torch.Tensor, params: list[torch.nn.Parameter]) -> float:
    """L2 norm of d loss / d params (the graph is kept for the real backward)."""
    grads = torch.autograd.grad(loss, params, retain_graph=True, allow_unused=True)
    sq = [g.float().pow(2).sum() for g in grads if g is not None]
    return float(torch.stack(sq).sum().sqrt()) if sq else 0.0


@dataclass
class Weights:
    decide: float = 1.0
    extract: float = 1.0
    count: float = 0.5
    anchor: float = 1.0
    anchor_count: float = 0.0
    lm: float = 1.0
    lm_kl_base: float = 0.0


def make_weights(args: argparse.Namespace) -> Weights:
    """The loss weights the CLI flags select (decide/extract/count keep their defaults)."""
    if not args.lm_weight >= 0:
        raise ValueError(f"--lm-weight must be >= 0: {args.lm_weight}")
    kl = getattr(args, "lm_kl_base", 0.0)
    if not kl >= 0:
        raise ValueError(f"--lm-kl-base must be >= 0: {kl}")
    return Weights(anchor=args.anchor_weight, anchor_count=args.anchor_count_weight,
                   lm=args.lm_weight, lm_kl_base=kl)


@dataclass(eq=False)
class Source:
    """One ``--records`` file: its own train pool (this rank's shard), its share of the
    micro-batches, and the uids already read (re-reads add only new records)."""

    path: Path
    weight: float = 1.0
    name: str = ""
    pool: list[Example | LongDoc] = field(default_factory=list)
    seen: set[str] = field(default_factory=set)
    offset: list[int] = field(default_factory=lambda: [0])  # bytes of the file read so far

    def __post_init__(self) -> None:
        if not self.weight > 0:
            raise ValueError(f"--records weight must be > 0: {self.path}:{self.weight}")
        if not self.name:
            self.name = self.path.parent.name if self.path.name == "records.jsonl" \
                else self.path.stem


def parse_records_arg(spec: str) -> Source:
    """``path[:weight]``; a suffix that does not parse as a number belongs to the path."""
    head, sep, tail = spec.rpartition(":")
    if sep and head:
        try:
            weight = float(tail)
        except ValueError:
            return Source(Path(spec))
        return Source(Path(head), weight)
    return Source(Path(spec))


def make_sources(specs: list[str] | None) -> list[Source]:
    sources = [parse_records_arg(s) for s in (specs or [str(DEG_RECORDS)])]
    names: dict[str, int] = defaultdict(int)
    for s in sources:  # distinct names for the logs
        names[s.name] += 1
        if names[s.name] > 1:
            s.name = f"{s.name}#{names[s.name]}"
    return sources


def active_sources(sources: list[Source], min_train: int) -> list[Source]:
    """Sources with at least ``min_train`` train examples (this rank); while none has,
    every non-empty source (a run never stalls on the threshold)."""
    act = [s for s in sources if len(s.pool) >= min_train]
    return act or [s for s in sources if s.pool]


def pick_source(sources: list[Source], rng: random.Random, min_train: int) -> Source:
    """A micro-batch's deg source: by weight among the active sources, so a small source
    gets its fixed share as soon as it has ``min_train`` examples."""
    act = active_sources(sources, min_train)
    if not act:
        raise RuntimeError("no deg training records in any --records source")
    return rng.choices(act, weights=[s.weight for s in act])[0]


def _ctx_len(e: Example | LongDoc | FillerDraw) -> int:
    return e.est_len if isinstance(e, (LongDoc, FillerDraw)) else len(e.ctx.ids)


def _fits(batch: list[Example], ex: Example, budget: int) -> bool:
    width = max(len(e.ctx.ids) for e in (*batch, ex))  # padded trunk width
    forks = sum(e.tokens - len(e.ctx.ids) for e in (*batch, ex))
    return width * (len(batch) + 1) + forks <= budget


def batches(sources: list[Source], jev_pool: list[Example] | None, rng: random.Random,
            budget: int, jev_frac: float, max_records: int, draw: int = 96,
            min_train: int = 1, materialize: Any = None,
            fits: Callable[[list[Example]], bool] | None = None,
            augment: Callable[[Any, random.Random], Any] | None = None,
            cap: Callable[[Example, random.Random], Example] | None = None):
    """Infinite micro-batches. Each picks one deg source (:func:`pick_source`; pools are
    read live, so reloads count at once), draws ``draw`` random records (jev with prob
    jev_frac, else from that source), sorts them by context length and packs a contiguous
    run from a random start up to the token budget, so the padded trunk batch wastes
    little. A :class:`LongDoc` sorts at ``max_context_tokens`` and becomes a fresh window
    (``materialize(doc, rng)``, None = skip) only when the packing reaches it.
    ``augment(e, rng)`` may replace a drawn record by a :class:`FillerDraw` (sorted at its
    planned length, materialized the same way), so the packing and the memory planner
    only ever see real, fully built examples.

    ``fits(records)`` (the memory planner) is a second packing limit beside ``budget``;
    the first record of a batch is always taken (an oversized one is split by rows into
    several micro-batches of one step, :meth:`deg.memplan.MemPlanner.plan`).

    ``cap(ex, rng)`` (:func:`cap_forks`) draws the forks every packed example trains this
    time (whole-context, window and filler examples alike), before packing."""

    def ok(batch: list[Example], ex: Example) -> bool:
        return _fits(batch, ex, budget) and (fits is None or fits([*batch, ex]))

    def real(e: Example | LongDoc | FillerDraw) -> Example | None:
        ex = materialize(e, rng) if isinstance(e, (LongDoc, FillerDraw)) else e
        return cap(ex, rng) if cap is not None and ex is not None else ex

    while True:
        pool = pick_source(sources, rng, min_train).pool
        cand = [rng.choice(jev_pool if (jev_pool and rng.random() < jev_frac) else pool)
                for _ in range(draw)]
        if augment is not None:
            cand = [augment(e, rng) for e in cand]
        cand.sort(key=_ctx_len)
        i = rng.randrange(len(cand))
        batch: list[Example] = []
        while i < len(cand) and len(batch) < max_records:
            ex = real(cand[i])
            if ex is None:
                i += 1
                continue
            if batch and not ok(batch, ex):
                break
            batch.append(ex)
            i += 1
        if len(batch) < max_records // 4 and i >= len(cand):  # ran off the end: extend down
            j = len(cand) - len(batch) - 1
            while j >= 0 and len(batch) < max_records:
                ex = real(cand[j])
                j -= 1
                if ex is None:
                    continue
                if batch and not ok(batch, ex):
                    break
                batch.insert(0, ex)
        if batch:
            yield batch


_POOL_FN: Any = None


def _pool_call(item: Any) -> tuple[Any, str | None]:
    try:
        return _POOL_FN(item), None
    except Exception as exc:  # noqa: BLE001 - a bad record is skipped and counted
        return None, f"{type(exc).__name__}: {exc}"[:200]


def pool_map(fn: Callable[[Any], Any], items: list[Any], workers: int
             ) -> list[tuple[Any, str | None]]:
    """``[(fn(x), None) or (None, error)]`` in order; ``workers`` > 1: forked CPU
    processes (``fn`` and everything it reads are inherited, never pickled; the children
    only tokenize and build python objects, they never touch CUDA)."""
    global _POOL_FN
    _POOL_FN = fn
    try:
        if workers <= 1 or len(items) < 64:
            return [_pool_call(x) for x in items]
        import multiprocessing as mp

        os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
        with mp.get_context("fork").Pool(workers) as pool:
            return pool.map(_pool_call, items, chunksize=max(1, len(items) // (workers * 8)))
    finally:
        _POOL_FN = None


# bump when deg_example / long_doc / jev_example / Example / Item change what they build
POOL_FORMAT = 1


def pool_cache_key(args: argparse.Namespace, model: DegModel) -> str:
    """Everything an Example depends on: :data:`POOL_FORMAT` (the builders in this file),
    the query/schema rendering code, the tokenizer layout, the example-shaping arguments
    and the seed (per-record rngs derive from it)."""
    import hashlib

    here = Path(__file__).parent
    h = hashlib.sha256(f"pool-format-{POOL_FORMAT}".encode())
    for name in ("infer.py", "schema.py", "request.py"):
        h.update((here / name).read_bytes())
    h.update(json.dumps({
        "backbone": model.meta.get("backbone"), "prefix": model.renderer.context("").ids,
        "eos": model.renderer.eos_id, "gen_layout": model.gen_layout,
        "array_mode": model.array_mode, "anchor_min_tokens": model.meta.get("anchor_min_tokens"),
        **{k: str(getattr(args, k, None)) for k in (
            "max_context_tokens", "anchor_count_weight", "filler_frac", "seed", "jev_limit")},
        # media runs only (a text run's key, and so its pool cache, is unchanged)
        **({"media_cache": str(args.media_cache)} if getattr(args, "media_cache", None) else {}),
    }, sort_keys=True).encode())
    return h.hexdigest()[:16]


def load_sources(sources: list[Source], build: Any, rank: int, world: int,
                 log: Any = lambda row: None, workers: int = 0,
                 cache: Path | None = None) -> dict[str, int]:
    """Read every source's ``split == "train"`` records not seen yet (this rank's uid
    shard) into its pool; ``build(rec)`` returns an Example to keep or None. Each source
    keeps its own split labels (val/test never train). Returns added per source.

    ``workers`` > 1 builds in forked CPU processes (``build`` must not depend on shared
    mutable state: per-record rngs). ``cache`` (a directory keyed by
    :func:`pool_cache_key`): a source's pool, seen uids and byte offset are restored from
    ``<cache>/<source>.r<rank>of<world>.pkl`` when present (restarts and the other runs of
    the same recipe skip the build) and written after every read; later reads (and
    ``--reload-every``) only parse the lines appended since."""
    import pickle

    added: dict[str, int] = {}
    for src in sources:
        n = 0
        if not src.path.exists():
            added[src.name] = 0
            continue
        cfile = (cache / f"{src.name}.r{rank}of{world}.pkl") if cache is not None else None
        if cfile is not None and not src.seen and cfile.exists():
            with cfile.open("rb") as fh:
                saved = pickle.load(fh)
            if saved["path"] == str(src.path) and saved["offset"] <= src.path.stat().st_size:
                src.pool, src.seen, src.offset = saved["pool"], saved["seen"], [saved["offset"]]
                n += len(src.pool)
                log({"event": "pool_cache", "source": src.name, "loaded": len(src.pool),
                     "offset": saved["offset"]})
        recs = [rec for rec in read_jsonl(src.path, "train", pos=src.offset)
                if rec["uid"] not in src.seen and zlib_shard(rec["uid"], world) == rank]
        for rec in recs:
            src.seen.add(rec["uid"])
        for rec, (ex, err) in zip(recs, pool_map(build, recs, workers)):
            if err is not None:
                log({"skip": rec["uid"], "source": src.name, "error": err})
                continue
            if ex is not None:
                src.pool.append(ex)
                n += 1
        if cfile is not None and recs:
            cfile.parent.mkdir(parents=True, exist_ok=True)
            tmp = cfile.with_suffix(f".tmp{os.getpid()}")
            with tmp.open("wb") as fh:
                pickle.dump({"path": str(src.path), "offset": src.offset[0], "seen": src.seen,
                             "pool": src.pool}, fh, protocol=pickle.HIGHEST_PROTOCOL)
            tmp.replace(cfile)
        added[src.name] = n
    return added


def lr_scale(step: int, steps: int, warmup: int, start_step: int = 0) -> float:
    """Linear warmup over ``warmup`` steps after ``start_step`` (a restart's optimizer
    moments start fresh), times the cosine to 10% over the whole run (absolute step)."""
    warm = min(1.0, (step - start_step) / max(1, warmup))
    progress = min(1.0, max(0.0, (step - warmup) / max(1, steps - warmup)))
    return warm * (0.5 * (1 + math.cos(math.pi * progress)) * 0.9 + 0.1)


def backbone_live(step: int, start_step: int, heads_only_steps: int) -> bool:
    """Whether ``step`` updates the backbone: not during the first ``heads_only_steps``
    steps after ``start_step``."""
    return step > start_step + heads_only_steps


# ------------------------------------------------------------- micro-batches + OOM recovery


def loss_counts(examples: list[Example]) -> dict[str, int]:
    """Items behind each loss term of :func:`forward_losses` (each term is a mean over
    them; ``lm`` counts target tokens)."""
    n: dict[str, int] = defaultdict(int)
    for ex in examples:
        for it in ex.items:
            if it.kind in ("choice", "score", "noul"):
                n["decide"] += 1
                n[f"decide_{it.kind}"] += 1
            elif it.kind == "generate":
                n["lm"] += len(it.seq) - it.prompt_len
                n["lm_kl_base"] += len(it.seq) - it.prompt_len
            elif it.kind == "count" and it.aux:
                n["anchor_count"] += 1
            else:  # extract | count | anchor
                n[it.kind] += 1
                if it.media:
                    n["extract_media"] += 1
    return n


# every key :func:`loss_counts` can produce (data-parallel count all-reduce)
COUNT_KEYS = ("decide", "decide_choice", "decide_score", "decide_noul", "lm", "lm_kl_base",
              "anchor_count", "extract", "extract_media", "count", "anchor")


def is_oom(exc: BaseException) -> bool:
    if isinstance(exc, torch.OutOfMemoryError):
        return True
    msg = str(exc)
    return isinstance(exc, RuntimeError) and (
        "out of memory" in msg or "CUBLAS_STATUS_ALLOC_FAILED" in msg
        or "CUDNN_STATUS_ALLOC_FAILED" in msg)


def split_part(part: list[Example], min_rows: int = 1) -> list[list[Example]] | None:
    """Halve a micro-batch that ran out of memory: its records in two halves; a single
    record's items (field rows) in two halves, each keeping the whole context. ``None``
    when ``part`` is one record with at most ``min_rows`` items (the minimal unit)."""
    if len(part) > 1:
        mid = len(part) // 2
        return [part[:mid], part[mid:]]
    ex = part[0]
    if len(ex.items) <= max(1, min_rows):
        return None
    mid = len(ex.items) // 2
    return [[dataclasses.replace(ex, items=ex.items[:mid])],
            [dataclasses.replace(ex, items=ex.items[mid:])]]


def free_device_memory() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


@dataclass
class StepResult:
    """One optimizer step's forward/backward over its micro-batches (parts).

    ``ok``: every item of the step contributed its gradient; the backbone's fp32 master
    grads and the heads' grads then hold exactly the gradient of the full step loss
    (each part's loss terms weighted by its share of the step's items). Not ``ok``: a
    minimal unit ran out of memory, every gradient of the step was discarded (``.grad``
    is None everywhere) and the caller must not step the optimizer."""

    ok: bool
    losses: dict[str, float] = field(default_factory=dict)  # step means, as one batch
    stats: dict[str, float] = field(default_factory=lambda: defaultdict(float))
    parts: int = 0  # micro-batches run successfully
    oom_retry: int = 0  # micro-batches that ran out of memory and were split
    oom_skip: list[dict[str, Any]] = field(default_factory=list)  # minimal units that OOMed
    peaks: list[tuple[list[Example], float]] = field(default_factory=list)  # activation bytes
    max_allocated: float = 0.0  # device peak (bytes allocated) over the step's parts
    trunk_shared: int = 0  # records whose row chunks ran on one trunk pass


def _discard_grads(params_bb: list[torch.nn.Parameter], master: Any,
                   params_hd: list[torch.nn.Parameter]) -> None:
    if isinstance(master, ZeroMaster):
        master.discard()
        master = []
    for q in (*params_bb, *master, *params_hd):
        q.grad = None


def detach_trunk(hidden: torch.Tensor, state: State
                 ) -> tuple[torch.Tensor, State, list[tuple[torch.Tensor, torch.Tensor]]]:
    """Leaf copies of a trunk's outputs (final hidden, recurrent state, conv tails, K/V)
    for forks that backpropagate into them, plus the ``(output, leaf)`` pairs: the trunk
    then backpropagates once from the leaves' accumulated ``.grad``."""
    pairs: list[tuple[torch.Tensor, torch.Tensor]] = []

    def leaf(t: torch.Tensor) -> torch.Tensor:
        if t is None or not t.requires_grad:
            return t
        x = t.detach().requires_grad_(True)
        pairs.append((t, x))
        return x

    st = State([(leaf(r), leaf(t)) for r, t in state.lin],
               [(leaf(k), leaf(v)) for k, v in state.kv], state.kv_mask, state.lengths)
    return leaf(hidden), st, pairs


@dataclass
class _Chunk:
    part: list[Example]
    losses: dict[str, torch.Tensor]
    stats: dict[str, float]
    n: dict[str, int]
    peak: float = 0.0  # activation bytes (device peak minus allocated before the part)
    top: float = 0.0  # device peak (bytes allocated)


def run_chunks(model: DegModel, chunks: list[list[Example]], weights: Weights,
               total_n: dict[str, int], fork_tokens: int, fork_rows: int,
               fault: Callable[[list[Example], str], None] | None = None,
               pre_backward: Callable[[dict[str, torch.Tensor]], None] | None = None
               ) -> list[_Chunk]:
    """Row chunks of ONE record (every chunk = ``[ex]`` with the same ``ex.ctx``) on one
    trunk pass: the trunk runs forward once, its outputs are detached into leaves
    (:func:`detach_trunk`), each chunk's forks + heads run forward and backward into those
    leaves (and the fork/head parameters), then the trunk runs backward once from the
    leaves' accumulated gradients. The gradient is the unsplit record's (each chunk's loss
    terms weighted by its share of the step's items, as in :func:`run_step`); the trunk
    activations stay alive across the chunks, which the memory planner already counts in
    every chunk. Parameter gradients accumulate in ``.grad`` over the chunks (the caller
    moves them to its accumulators once the whole record succeeded). ``peak`` per chunk:
    activation peak from before the trunk (device only; the trunk backward's peak is
    folded into the last chunk). Out-of-memory errors propagate (``.grad`` partial)."""
    cuda = model.device.type == "cuda"
    dev = model.device
    if cuda:
        torch.cuda.reset_peak_memory_stats(dev)
        before = torch.cuda.memory_allocated(dev)
    ex0 = chunks[0][0]
    assert all(len(c) == 1 and c[0].ctx is ex0.ctx for c in chunks)
    hidden, state = model.trunk([ex0.ctx])
    h_leaf, st_leaf, pairs = detach_trunk(hidden, state)
    out: list[_Chunk] = []
    for k, part in enumerate(chunks):
        if fault is not None:
            fault(part, "forward")
        if cuda and k:
            torch.cuda.reset_peak_memory_stats(dev)
        losses, stats = forward_losses(model, part, fork_tokens, fork_rows,
                                       trunk=(h_leaf, st_leaf))
        if k:  # one trunk pass for the whole record
            stats["trunk_tokens"] = stats["trunk_padded"] = 0.0
        n = loss_counts(part)
        share = {key: n[key] / total_n[key] for key in n if total_n[key]}
        total = sum(getattr(weights, key) * v * share.get(key, 0.0)
                    for key, v in losses.items() if key in LOSS_KEYS)
        if pre_backward is not None:
            pre_backward(losses)
        if isinstance(total, torch.Tensor) and total.requires_grad:
            total.backward()
        del total
        if fault is not None:
            fault(part, "backward")
        c = _Chunk(part, {key: v.detach() for key, v in losses.items()}, stats, n)
        if cuda:
            c.top = float(torch.cuda.max_memory_allocated(dev))
            c.peak = c.top - before
        out.append(c)
        del losses
    if cuda:
        torch.cuda.reset_peak_memory_stats(dev)
    outs = [(t, x.grad) for t, x in pairs if x.grad is not None]
    del h_leaf, st_leaf, pairs
    if outs:
        torch.autograd.backward([t for t, _ in outs], [g for _, g in outs])
    del outs, hidden, state
    if fault is not None:
        fault(chunks[-1], "trunk")
    if cuda:
        top = float(torch.cuda.max_memory_allocated(dev))
        out[-1].top = max(out[-1].top, top)
        out[-1].peak = max(out[-1].peak, top - before)
    return out


def _chunk_run(queue: deque, first: tuple[list[Example], bool]) -> list[list[Example]]:
    """``first`` and the queued parts right after it that are row chunks of the same
    record (single-record parts sharing its context object, :meth:`MemPlanner.split_rows`)
    and may share a trunk pass."""
    part, shared = first
    run = [part]
    if not shared or len(part) != 1:
        return run
    while queue and queue[0][1] and len(queue[0][0]) == 1 and queue[0][0][0].ctx is part[0].ctx:
        run.append(queue.popleft()[0])
    return run


def run_step(model: DegModel, parts: list[list[Example]], weights: Weights,
             fork_tokens: int, fork_rows: int, params_bb: list[torch.nn.Parameter],
             master: list[torch.Tensor] | ZeroMaster, params_hd: list[torch.nn.Parameter],
             live: bool, min_rows: int = 1,
             fault: Callable[[list[Example], str], None] | None = None,
             pre_backward: Callable[[dict[str, torch.Tensor]], None] | None = None,
             grad_bytes: float = 0.0, share_trunk: bool = True,
             total_n: dict[str, int] | None = None) -> StepResult:
    """Forward + backward of every part, accumulating into the fp32 master grads
    (backbone) and the heads' grads; every loss term is the mean over the whole step's
    items whatever the split.

    Consecutive parts that are row chunks of one record (the planner's split of a record
    too big for one micro-batch) run on ONE trunk pass (:func:`run_chunks`; ``share_trunk``
    off: one trunk pass per chunk). If such a run hits out-of-memory, its gradients are
    dropped and its chunks run again one by one with their own trunk passes.

    A part that raises an out-of-memory error (forward or backward) is dropped together
    with any gradient it had produced (a part's grads join the accumulators only after its
    backward completed), memory is freed, and the part runs again as two halves
    (:func:`split_part`). A minimal unit that still runs out of memory is skipped
    (``oom_skip``) and the whole step is abandoned: no partial-gradient step, so the
    optimizer only ever sees complete steps. ``fault(part, phase)`` (tests) may raise
    before the forward (``"forward"``), after the backward (``"backward"``) or, on a
    shared trunk, after the trunk's backward (``"trunk"``).

    ``master`` may be a :class:`deg.zero.ZeroMaster` (ZeRO-2 sharded masters): the
    step's backbone gradient then stays in it (bf16 on the device for a one-part step,
    fp32 on the host otherwise) until the caller's ``reduce_scatter``.

    ``total_n`` (data parallel): the item counts of the GLOBAL step (every rank's batch,
    all-reduced); each loss term is then this rank's share of the global mean and the
    ranks' gradients SUM to the gradient of the global step loss."""
    res = StepResult(ok=True)
    total_n = total_n or loss_counts([ex for part in parts for ex in part])
    zero = master if isinstance(master, ZeroMaster) else None
    acc_hd: list[torch.Tensor | None] = [None] * len(params_hd)
    _discard_grads(params_bb, master, params_hd)
    cuda = model.device.type == "cuda"
    # (part, may share a trunk pass with the next parts of the same record)
    queue: deque[tuple[list[Example], bool]] = deque((p, share_trunk) for p in parts)

    def drop_grads() -> None:
        for q in (*params_bb, *params_hd):  # this part's (partial) grads
            q.grad = None
        free_device_memory()

    def merge(done: list[_Chunk]) -> None:
        if zero is not None:
            zero.collect(more=bool(queue))
        for q, mq in zip(params_bb, [] if zero is not None else master):
            if q.grad is not None:
                if mq.grad is None:
                    mq.grad = q.grad.float()
                else:
                    mq.grad.add_(q.grad)
                q.grad = None
        for i, q in enumerate(params_hd):
            if q.grad is not None:
                if acc_hd[i] is None:
                    acc_hd[i] = q.grad
                else:
                    acc_hd[i].add_(q.grad)  # type: ignore[union-attr]
                q.grad = None
        for c in done:
            if cuda:
                res.peaks.append((c.part, max(0.0, c.peak - (grad_bytes if live else 0.0))))
            for k, v in c.losses.items():
                if total_n.get(k):
                    res.losses[k] = res.losses.get(k, 0.0) + float(v) * c.n[k] / total_n[k]
            for k, v in c.stats.items():
                res.stats[k] += v
            res.max_allocated = max(res.max_allocated, c.top)
            res.parts += 1

    while queue:
        run = _chunk_run(queue, queue.popleft())
        if len(run) > 1:
            try:
                done = run_chunks(model, run, weights, total_n, fork_tokens, fork_rows,
                                  fault, pre_backward)
            except Exception as exc:  # noqa: BLE001 - only out-of-memory is recovered
                if not is_oom(exc):
                    raise
                drop_grads()
                res.oom_retry += 1
                queue.extendleft((c, False) for c in reversed(run))
                continue
            merge(done)
            res.trunk_shared += 1
            continue
        part = run[0]
        oom = False
        chunk: _Chunk | None = None
        before = 0
        try:
            if fault is not None:
                fault(part, "forward")
            if cuda:
                torch.cuda.reset_peak_memory_stats(model.device)
                before = torch.cuda.memory_allocated(model.device)
            losses, stats = forward_losses(model, part, fork_tokens, fork_rows)
            n = loss_counts(part)
            share = {k: n[k] / total_n[k] for k in n if total_n[k]}
            total = sum(getattr(weights, k) * v * share.get(k, 0.0)
                        for k, v in losses.items() if k in LOSS_KEYS)
            if pre_backward is not None:
                pre_backward(losses)
            if isinstance(total, torch.Tensor) and total.requires_grad:
                total.backward()
            del total
            if fault is not None:
                fault(part, "backward")
            chunk = _Chunk(part, {k: v.detach() for k, v in losses.items()}, stats, n)
            del losses
            if cuda:
                chunk.top = float(torch.cuda.max_memory_allocated(model.device))
                chunk.peak = chunk.top - before
        except Exception as exc:  # noqa: BLE001 - only out-of-memory is recovered
            if not is_oom(exc):
                raise
            oom = True
        if oom:
            drop_grads()
            halves = split_part(part, min_rows)
            if halves is None:
                ex = part[0]
                res.oom_skip.append({
                    "event": "oom_skip", "uid": ex.uid, "source": ex.source,
                    "context_tokens": len(ex.ctx.ids), "items": len(ex.items),
                    "fork_tokens": sum(len(it.seq) for it in ex.items),
                    "window": list(ex.window) if ex.window else None})
                res.ok = False
                break
            res.oom_retry += 1
            queue.extendleft((h, False) for h in reversed(halves))
            continue
        assert chunk is not None
        merge([chunk])
    if not res.ok:
        _discard_grads(params_bb, master, params_hd)
        acc_hd.clear()
        free_device_memory()
        return res
    for q, a in zip(params_hd, acc_hd):
        q.grad = a
    for q, mq in zip(params_bb, [] if zero is not None else master):
        # heads-only steps: no grad -> AdamW skips the backbone (no decay, no moments)
        mq.grad = None if not live else (mq.grad if mq.grad is not None
                                         else torch.zeros_like(mq))
    res.losses["total"] = sum(getattr(weights, k) * v for k, v in res.losses.items()
                              if k in LOSS_KEYS)
    return res


def planner_features(planner: MemPlanner, part: list[Example]) -> dict[str, Any]:
    from deg.memplan import features

    return features(part, planner.fork_tokens, planner.fork_rows).as_dict()


def all_ranks_ok(ok: bool, world: int, device: torch.device) -> bool:
    if world == 1:
        return ok
    import torch.distributed as dist

    flag = torch.tensor([1.0 if ok else 0.0], device=device)
    dist.all_reduce(flag, op=dist.ReduceOp.MIN)
    return bool(flag.item() > 0)


def _dist() -> tuple[int, int]:
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        import torch.distributed as dist

        if not dist.is_initialized():
            import datetime

            # rank 0 alone calibrates/evaluates/saves at the end while the others wait
            dist.init_process_group(os.environ.get("DEG_DIST_BACKEND", "nccl"),
                                    timeout=datetime.timedelta(hours=4))
        return dist.get_rank(), dist.get_world_size()
    return 0, 1


def global_counts(batch: list[Example], world: int, device: torch.device
                  ) -> dict[str, int] | None:
    """Item counts of every loss term over ALL ranks' batches of this step (None with one
    rank): :func:`run_step` then makes each rank's loss its share of the global mean."""
    if world == 1:
        return None
    import torch.distributed as dist

    local = loss_counts(batch)
    keys = sorted(COUNT_KEYS)
    extra = sorted(set(local) - set(keys))
    if extra:
        raise RuntimeError(f"loss counts without a global key: {extra}")
    t = torch.tensor([float(local.get(k, 0)) for k in keys], device=device,
                     dtype=torch.float64)
    dist.all_reduce(t)
    return {k: int(v) for k, v in zip(keys, t.tolist())}


def sum_over_ranks(values: dict[str, float], world: int, device: torch.device
                   ) -> dict[str, float]:
    """Sum of every rank's ``values`` (global-mean shares -> the global means); keys are
    the union over ranks."""
    if world == 1:
        return values
    import torch.distributed as dist

    keys_all: list[Any] = [None] * world
    dist.all_gather_object(keys_all, sorted(values))
    keys = sorted(set().union(*keys_all))
    t = torch.tensor([float(values.get(k, 0.0)) for k in keys], device=device,
                     dtype=torch.float64)
    dist.all_reduce(t)
    return dict(zip(keys, t.tolist()))


def replicas_equal(model: DegModel, world: int) -> bool:
    """Whether every rank holds bitwise the same backbone and heads (per-tensor float64
    sums and sums of squares of the raw bits, compared across ranks)."""
    import torch.distributed as dist

    vals = []
    for p in [*model.backbone.parameters(), *model.heads.parameters()]:
        x = p.detach().reshape(-1)
        x = x.view(torch.int16 if x.element_size() == 2 else torch.int32)
        for c in x.split(1 << 24):  # bounded float64 temporaries
            bits = c.double()
            vals += [bits.sum(), (bits * bits).sum()]
    mine = torch.stack(vals)
    lo, hi = mine.clone(), mine.clone()
    dist.all_reduce(lo, op=dist.ReduceOp.MIN)
    dist.all_reduce(hi, op=dist.ReduceOp.MAX)
    return bool(torch.equal(lo, hi))


def clip_global_(zero: ZeroMaster | None, params_hd: list[torch.nn.Parameter],
                 max_norm: float) -> torch.Tensor:
    """``clip_grad_norm_`` over the sharded backbone gradient (``zero``; None = not
    stepped) and the replicated heads' gradients: the same total norm and clip factor on
    every rank. Returns the total norm before clipping."""
    hd = [p.grad for p in params_hd if p.grad is not None]
    sq = (torch.stack([g.float().pow(2).sum() for g in hd]).sum() if hd
          else None)
    if zero is not None:
        zsq = zero.grad_sq_norm()
        sq = zsq if sq is None else sq + zsq.to(sq.device)
    if sq is None:  # nothing stepped (frozen heads and no backbone step)
        sq = torch.zeros((), device=zero.grad_buf.device if zero is not None else "cpu")
    norm = sq.sqrt()
    coef = torch.clamp(max_norm / (norm + 1e-6), max=1.0)
    for g in hd:
        g.mul_(coef.to(g.dtype))
    if zero is not None:
        zero.grad_buf.mul_(coef)
    return norm


def _allreduce_grads(params: list[torch.nn.Parameter], world: int,
                     bucket: int = 1 << 25, average: bool = True) -> None:
    """Average (``average``) or sum gradients across ranks in ~128 MB fp32 buckets
    (missing grads count as 0)."""
    if world == 1:
        return
    import torch.distributed as dist

    for p in params:
        if p.grad is None:
            p.grad = torch.zeros_like(p)
    group: list[torch.nn.Parameter] = []
    size = 0
    for i, p in enumerate(params):
        group.append(p)
        size += p.numel()
        if size >= bucket or i == len(params) - 1:
            flat = torch.cat([q.grad.reshape(-1) for q in group])
            dist.all_reduce(flat)
            if average:
                flat /= world
            o = 0
            for q in group:
                q.grad.copy_(flat[o:o + q.numel()].view_as(q))
                o += q.numel()
            group, size = [], 0


def manifest_records(path: Path, skip: list[str], scale: dict[str, float] | None = None
                     ) -> list[str]:
    """``--records`` specs (``PATH:WEIGHT``) of a datagen manifest's sets (``sets[]`` with
    ``path``, ``group``, ``proposed_weight``), leaving out the groups in ``skip``; a set
    under one of the manifest's ``eval_only_never_train`` patterns is refused. ``scale``
    multiplies the weights of a group's sets (``{"mm": 4}``)."""
    import fnmatch

    m = json.loads(Path(path).read_text())
    never = [str(p) for p in m.get("eval_only_never_train", [])]
    out = []
    for st in m["sets"]:
        if st.get("group") in skip:
            continue
        p = str(st["path"])
        if any(p.startswith(n.rstrip("/") + "/") or fnmatch.fnmatch(p, n)
               or fnmatch.fnmatch(p, "*/" + n) for n in never):
            raise SystemExit(f"manifest set {st.get('name')} is eval-only: {p}")
        w = float(st["proposed_weight"]) * (scale or {}).get(st.get("group"), 1.0)
        if w > 0:
            out.append(f"{p}:{w}")
    return out


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--records", action="append", default=None, metavar="PATH[:WEIGHT]",
                    help="datagen records file, repeatable; WEIGHT (default 1) is the source's "
                         "share of the deg micro-batches once it has --source-min-train "
                         f"train records (default: {DEG_RECORDS})")
    ap.add_argument("--manifest", type=Path, default=None,
                    help="a datagen ALL-*.json manifest: every listed set (minus "
                         "--manifest-skip groups) becomes --records PATH:proposed_weight; "
                         "its eval_only_never_train paths are refused")
    ap.add_argument("--manifest-skip", default="mm",
                    help="comma list of manifest groups to leave out (default mm: media "
                         "records need --media-cache; pass '' to train every group)")
    ap.add_argument("--manifest-group-scale", default="",
                    help="comma list GROUP=FACTOR multiplying a manifest group's set weights "
                         "(e.g. mm=4: the mm sets are ~5%% of ALL-v5's weight)")
    ap.add_argument("--media-cache", type=Path,
                    default=Path(os.environ["DEG_MEDIA_CACHE"])
                    if os.environ.get("DEG_MEDIA_CACHE") else None,
                    help="tower-embedding cache of the media records (python -m deg.media "
                         "cache; default $DEG_MEDIA_CACHE): media records train from it, a "
                         "record whose media is not cached is skipped (logged)")
    ap.add_argument("--source-min-train", type=int, default=50,
                    help="train records (all ranks) before a source gets its share")
    ap.add_argument("--jev", type=Path, default=JEV_TRAIN)
    ap.add_argument("--jev-limit", type=int, default=None)
    ap.add_argument("--jev-frac", type=float, default=0.3)
    ap.add_argument("--init", type=Path, default=None, help="deg checkpoint to start from")
    ap.add_argument("--backbone", "--base-model", dest="backbone", default=None,
                    help=f"base model for fresh heads: {sorted(BACKBONES)}, any "
                         "qwen3.5-<size>, or its HF id (Qwen/Qwen3.5-2B); default "
                         "qwen3.5-0.8b; with --init the checkpoint's backbone (a different "
                         "value is refused)")
    ap.add_argument("--steps", type=int, default=300)
    ap.add_argument("--start-step", type=int, default=0,
                    help="restart: with --init <run>/step-N, continue the schedule and the "
                         "step numbering at N (optimizer moments restart with a warmup)")
    ap.add_argument("--lr", type=float, default=2e-5)
    ap.add_argument("--head-lr", type=float, default=1e-3)
    ap.add_argument("--warmup", type=int, default=30)
    ap.add_argument("--budget", type=int, default=12000, help="tokens per micro-batch")
    ap.add_argument("--max-records", type=int, default=16)
    ap.add_argument("--fork-tokens", type=int, default=8192)
    ap.add_argument("--fork-rows", type=int, default=64)
    ap.add_argument("--record-fork-tokens", type=int, default=0,
                    help="fork tokens one record may train per step (a fresh random subset "
                         "of its field rows each draw, anchor/count forks always kept; "
                         "cap_forks); 0 = what --budget leaves beside its context, at least "
                         "--budget/2")
    ap.add_argument("--decide-forks-per-field", type=int, default=0,
                    help="at most this many decide forks per schema field per draw (a fresh "
                         "random subset; cap_decide_forks): repeated in-array enums of long "
                         "filings stop dominating the decide loss; 0 = no cap")
    ap.add_argument("--max-context-tokens", type=int, default=6000,
                    help="longer contexts train span-aware windows of this many tokens "
                         "(fields whose gold the window does not show are dropped)")
    ap.add_argument("--log-every", type=int, default=10)
    ap.add_argument("--save-every", type=int, default=0)
    ap.add_argument("--pool-workers", type=int, default=0,
                    help="build each rank's example pools in this many forked CPU processes "
                         "(0 = in the trainer process)")
    ap.add_argument("--pool-cache", type=Path, default=None,
                    help="directory of built example pools, keyed by code + tokenizer + "
                         "example arguments + seed: restarts and runs of the same recipe "
                         "load them instead of rebuilding; new records are appended")
    ap.add_argument("--pool-only", default=None, metavar="RANK/WORLD",
                    help="build (into --pool-cache) the example pools of rank RANK of a "
                         "WORLD-rank run on the CPU, then exit: pools are ready before the "
                         "GPUs are")
    ap.add_argument("--reload-every", type=int, default=0,
                    help="re-read --records every N steps (the datagen file grows)")
    ap.add_argument("--eval-records", type=int, default=60)
    ap.add_argument("--eval-jev", type=int, default=300)
    ap.add_argument("--no-grad-ckpt", action="store_true")
    ap.add_argument("--no-final-eval", action="store_true",
                    help="save final/ without calibration and eval")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--array-mode", choices=ARRAY_MODES, default="anchor",
                    help="anchor: position-anchored array items, no count cap (SPEC §12.6); "
                         "index: 'ITEM i of' + the capped count head (run-1)")
    ap.add_argument("--anchor-weight", type=float, default=1.0)
    ap.add_argument("--anchor-min-tokens", type=int, default=0,
                    help="anchor mode: documents shorter than this many context tokens "
                         "address array items by index + count (run-1's path, better on "
                         "short records) in training and inference; saved in deg.json")
    ap.add_argument("--lm-weight", type=float, default=1.0,
                    help="weight of the x-generate next-token loss (lm); the LM trains only "
                         "through the backbone, so anchor/count terms that dominate the "
                         "clipped gradient dilute it (results/deg/gen-regress)")
    ap.add_argument("--gen-layout", choices=GEN_LAYOUTS, default=None,
                    help="how the model writes x-generate text with its own weights: deg "
                         "(WRITE/VALUE layout) or chat (the base model's chat turn, "
                         "infer.base_generate_header); saved in deg.json; default: the "
                         "--init checkpoint's, else chat for Gemma 4 (an instruction-tuned "
                         "base writes well there from step 0), deg otherwise")
    ap.add_argument("--moe-experts", default=None,
                    help="'d,g': make a fresh model a D/E/G task MoE (model.add_task_moe): "
                         "decide forks train FFN expert D, extract forks + the context the "
                         "layer's own FFN (E), generate G; all start as the base FFN")
    ap.add_argument("--lm-kl-base", type=float, default=0.0,
                    help="weight of KL(frozen base || deg) on x-generate target tokens "
                         "(anti-forgetting; loads the base body beside the model, +~1 GB)")
    ap.add_argument("--lora-rank", type=int, default=0,
                    help="LoRA adapters of this rank on a frozen base backbone (0 = full "
                         "fine-tune); x-generate then decodes with the base")
    ap.add_argument("--lora-alpha", type=float, default=None, help="default 2 * rank")
    ap.add_argument("--adapters", default="shared",
                    help="'shared' or a comma list of decide,extract (one adapter each)")
    ap.add_argument("--lora-context", default="base",
                    help="'base' (context read once by the base weights) or an adapter "
                         "name that also reads the context")
    ap.add_argument("--anchor-count-weight", type=float, default=0.0,
                    help="weight of the count head's CE on anchored arrays (anchor_count; "
                         "0 = not trained, no count forks for anchored arrays)")
    ap.add_argument("--filler-frac", type=float, default=0.0,
                    help="probability that a drawn short datagen record is embedded in "
                         "unrelated filler text (gold spans shifted)")
    ap.add_argument("--filler-corpus", type=Path, default=FILLER_CORPUS,
                    help="v3 docs.jsonl (or its directory); train-split docs only")
    ap.add_argument("--filler-max-tokens", type=int, default=16384,
                    help="filler examples' context length is log-uniform from the record's "
                         "own length to this (capped at --max-context-tokens)")
    ap.add_argument("--filler-negative-frac", type=float, default=0.25,
                    help="share of filler draws that are filler only: the record's top-level "
                         "anchored arrays train with no item start (and count 0)")
    ap.add_argument("--heads-only-steps", type=int, default=0,
                    help="the first N steps after --start-step train the heads only (fresh "
                         "anchor head settles before its gradients reach the backbone); the "
                         "backbone LR then warms up from there")
    ap.add_argument("--loss-gnorm-every", type=int, default=0,
                    help="every N steps log each weighted loss's backbone gradient norm "
                         "(gnorm_<loss>; one extra backward pass per loss)")
    ap.add_argument("--mem-target-gb", type=float, default=18.0,
                    help="device memory (GiB) a micro-batch may plan to reach: batches are "
                         "packed to it and an oversized record is split into field-row "
                         "chunks (several micro-batches, one optimizer step); 0 = off")
    ap.add_argument("--mem-coef", type=Path, default=None,
                    help="JSON of deg.memplan.MemCoef overrides (bytes per feature)")
    ap.add_argument("--oom-min-rows", type=int, default=1,
                    help="an out-of-memory micro-batch is halved down to one record with "
                         "this many field rows; a minimal unit that still OOMs is skipped "
                         "(oom_skip) and the step redrawn")
    ap.add_argument("--oom-max-redraws", type=int, default=4,
                    help="fresh batches per step after an oom_skip before the step is "
                         "skipped without an optimizer update (oom_step_skipped)")
    ap.add_argument("--max-record-parts", type=int, default=0,
                    help="at most this many micro-batches per record in one step (a record "
                         "split into more row chunks keeps a random subset of them; "
                         "cap_record_parts); 0 = no cap")
    ap.add_argument("--init-heads", type=Path, default=None,
                    help="load only the heads of this deg checkpoint (same backbone "
                         "size) onto a fresh/base backbone, e.g. LoRA on the base with the "
                         "heads of a full fine-tune")
    ap.add_argument("--no-share-trunk", action="store_true",
                    help="a record split into row chunks re-runs its trunk per chunk "
                         "(default: one trunk pass, its gradient accumulated over the chunks)")
    ap.add_argument("--freeze-experts", default="",
                    help="comma list of task-MoE experts (d,g) whose FFNs stay frozen, e.g. g: "
                         "the writer keeps the base model's FFN while D/E and the shared trunk "
                         "train (gemma E2B: training G caused pointer-echo loops)")
    ap.add_argument("--train-only-experts", default="",
                    help="comma list of task-MoE experts (e.g. g): ONLY their FFNs train; the "
                         "shared trunk, the other experts and the heads (except --train-heads) stay frozen")
    ap.add_argument("--train-heads", default="",
                    help="with --train-only-experts: comma list of heads that also train "
                         "(e.g. choice,score,noul for a D + decide-heads fine-tune)")
    ap.add_argument("--freeze-backbone", action="store_true",
                    help="train the heads only (smoke tests on a GPU shared with serving)")
    ap.add_argument("--no-zero", action="store_true",
                    help="data parallel without ZeRO-2 (every rank keeps whole fp32 masters, "
                         "moments and grads); default with several ranks: sharded")
    ap.add_argument("--zero-sim-world", type=int, default=0,
                    help="smoke test on ONE GPU: shard as rank 0 of this many ranks (real "
                         "per-rank memory and step time, no collectives; other shards are "
                         "not updated - never for a real run)")
    args = ap.parse_args(argv)
    if args.manifest is not None:
        scale = {k: float(v) for k, v in (x.split("=") for x in
                                          args.manifest_group_scale.split(",") if x)}
        args.records = (args.records or []) + manifest_records(
            args.manifest, args.manifest_skip.split(","), scale)
    if args.backbone is not None:
        try:
            args.backbone = backbone_spec(args.backbone).name
        except ValueError as exc:
            ap.error(str(exc))

    if args.loss_gnorm_every:  # per-loss backward passes retain the graph
        import torch._functorch.config as functorch_config

        functorch_config.donated_buffer = False
    if args.pool_only:
        if args.pool_cache is None:
            raise SystemExit("--pool-only needs --pool-cache")
        rank, world = (int(x) for x in args.pool_only.split("/"))
        local = -1
    else:
        local = int(os.environ.get("LOCAL_RANK", "0")) % torch.cuda.device_count()
        torch.cuda.set_device(local)
        rank, world = _dist()
    torch.manual_seed(args.seed + rank)
    # a restart draws a fresh data order (the same seed would replay the first steps)
    rng = random.Random(args.seed * 1000 + rank + 7919 * args.start_step)
    args.out.mkdir(parents=True, exist_ok=True)
    log_path = args.out / "train.log.jsonl"

    def log(row: dict[str, Any]) -> None:
        if rank == 0:
            print(json.dumps(row), flush=True)
            with log_path.open("a") as fh:
                fh.write(json.dumps(row) + "\n")

    # bf16 weights for compute, fp32 master copies in the optimizer (no autocast casts)
    model = load_model(args.init, device="cpu" if args.pool_only else f"cuda:{local}",
                       dtype=torch.bfloat16, backbone=args.backbone)
    if args.init_heads is not None:
        if args.init is not None:
            raise SystemExit("--init-heads is for a fresh backbone (no --init)")
        init_heads(model, args.init_heads)
        log({"event": "init_heads", "path": str(args.init_heads)})
    log({"event": "backbone", "backbone": model.meta["backbone"], "base": model.meta["base"],
         "params": sum(p.numel() for p in model.backbone.parameters())})
    model.backbone.grad_ckpt = not args.no_grad_ckpt
    model.meta["array_mode"] = args.array_mode
    model.meta["anchor_min_tokens"] = args.anchor_min_tokens
    if args.gen_layout is not None:
        model.meta["gen_layout"] = args.gen_layout
    elif "gen_layout" not in model.meta and isinstance(model.backbone, GemmaBackbone):
        model.meta["gen_layout"] = "chat"
    log({"event": "gen_layout", "gen_layout": model.gen_layout})
    if args.moe_experts and model.moe is None:
        from deg.model import add_task_moe, task_moe_params

        add_task_moe(model, args.moe_experts.split(","))
        log({"event": "task_moe", **model.moe, "params": task_moe_params(model)})
    emb = model.backbone.text.embed_tokens.weight
    emb.requires_grad_(False)  # tied LM head stays the base model's (no optimizer state)
    if args.lm_kl_base > 0 and not args.pool_only:
        if args.lora_rank or model.lora is not None:
            raise SystemExit("--lm-kl-base is for full fine-tunes (LoRA generates with the "
                             "base, adapter off)")
        from deg.model import attach_base

        attach_base(model)  # frozen; shares the frozen embeddings when equal
    if args.lora_rank and model.lora is None:
        from deg.model import add_lora

        model.backbone.requires_grad_(False)
        add_lora(model, args.lora_rank, args.lora_alpha, args.adapters.split(","),
                 args.lora_context, seed=args.seed)
    elif model.lora is not None:  # resuming a LoRA checkpoint: only the adapters train
        for name, p in model.backbone.named_parameters():
            p.requires_grad_(".lora_" in name)
    if model.lora is not None:
        log({"event": "lora", **model.lora,
             "params": sum(p.numel() for p in model.backbone.parameters()
                           if p.requires_grad)})
    if args.freeze_backbone:
        model.backbone.requires_grad_(False)
    if args.freeze_experts:
        keys = tuple(f".mlp_{x}." for x in args.freeze_experts.split(",") if x)
        frozen = 0
        for name, p in model.backbone.named_parameters():
            if any(k in name for k in keys):
                p.requires_grad_(False)
                frozen += p.numel()
        log({"event": "freeze_experts", "experts": args.freeze_experts, "params": frozen})
    if args.train_only_experts:
        # "e" = the reader/extract FFN, which keeps the dense name `.mlp.` (D/G are `.mlp_d.`/`.mlp_g.`)
        keys = tuple(".mlp." if x == "e" else f".mlp_{x}." for x in args.train_only_experts.split(",") if x)
        live = 0
        for name, p in model.backbone.named_parameters():
            p.requires_grad_(any(k in name for k in keys))
            live += p.numel() * p.requires_grad
        model.heads.requires_grad_(False)
        for h in (x for x in args.train_heads.split(",") if x):
            sub = getattr(model.heads, h)  # AttributeError on an unknown head name
            sub.requires_grad_(True)
            live += sum(p.numel() for p in sub.parameters())
        log({"event": "train_only_experts", "experts": args.train_only_experts,
             "heads": args.train_heads, "params": live})
    if world > 1 and not args.pool_only:  # fresh heads from per-rank seeds: equal replicas
        broadcast_module(model.heads)
        broadcast_module(model.backbone.lm)
    model.backbone.train()
    model.heads.train()

    sources = make_sources(args.records)
    # JevBench is eval only (deg.bench.jevbench): refuse its files and drop any record or
    # jev example whose text contains one of its public states
    from deg.bench.jevbench import ContaminationGuard, guard_path

    for x in sources:
        guard_path(x.path)
    guard_path(args.jev)
    guard_path(args.filler_corpus)
    jevbench_guard = ContaminationGuard()
    jevbench_dropped = {"deg": 0, "jev": 0}
    min_train = max(1, math.ceil(args.source_min_train / world))
    arrays = {"anchored": 0, "skipped": 0}

    count_aux = args.anchor_count_weight > 0
    filler = FillerCorpus(args.filler_corpus) if args.filler_frac > 0 else None
    filler_cap = min(args.filler_max_tokens, args.max_context_tokens)

    media_enc = None
    if args.media_cache is not None:
        from deg.media import EmbedCache, MediaEncoder

        mcache = EmbedCache(args.media_cache)  # read only: built ahead by deg.media cache
        media_enc = MediaEncoder(model.meta.get("base"), cache=mcache, encode=False)
        log({"event": "media_cache", "path": str(args.media_cache), "items": len(mcache.index),
             **mcache.config})

    def build(rec: dict[str, Any]) -> Example | LongDoc | None:
        # no shared mutable state (pool_map may run it in forked workers): the record's own
        # rng; guard drops and array counts are tallied from the results (count_built)
        if jevbench_guard.match(record_text(rec)):
            return None
        if is_media_record(rec):  # whole context only: no windows, no filler augmentation
            ex = deg_example(model, rec, random.Random(f"{args.seed}:{rec['uid']}"),
                             anchor_count=count_aux, media=media_enc)
            if not ex.items or len(ex.ctx.ids) > args.max_context_tokens:
                return None
            return ex
        if len(model.renderer.context(record_text(rec)).ids) > args.max_context_tokens:
            return long_doc(model, rec, args.max_context_tokens)  # span-aware windows
        r = random.Random(f"{args.seed}:{rec['uid']}")
        ex = deg_example(model, rec, r, anchor_count=count_aux)
        if not ex.items:
            return None
        if filler is not None:
            ex.rec = rec
        return ex

    def count_built() -> None:
        arrays["anchored"] = sum(e.anchored_arrays for x in sources for e in x.pool
                                 if isinstance(e, Example))
        arrays["skipped"] = sum(e.skipped_arrays for x in sources for e in x.pool
                                if isinstance(e, Example))

    def augment(e: Any, r: random.Random) -> Any:
        if (filler is None or not isinstance(e, Example) or e.rec is None
                or r.random() >= args.filler_frac):
            return e
        lo = len(e.ctx.ids) + 64
        if filler_cap <= lo:
            return e
        tokens = int(math.exp(r.uniform(math.log(lo), math.log(filler_cap))))
        return FillerDraw(e, tokens, r.random() < args.filler_negative_frac)

    def materialize(doc: LongDoc | FillerDraw, r: random.Random) -> Example | None:
        if isinstance(doc, FillerDraw):
            assert filler is not None
            return filler_example(model, doc, filler, r, args.max_context_tokens, count_aux)
        return window_example(model, doc, r, args.budget,
                              record_fork_tokens=args.record_fork_tokens)

    cache = (args.pool_cache / pool_cache_key(args, model)) if args.pool_cache else None

    def load_deg() -> dict[str, int]:
        added = load_sources(sources, build, rank, world, log, args.pool_workers, cache)
        count_built()
        return added

    load_deg()
    jev_pool: list[Example] = []
    if args.jev_frac > 0 and args.jev.exists():
        import pickle

        jfile = cache / f"jev.r{rank}of{world}.pkl" if cache is not None else None
        if jfile is not None and jfile.exists():
            with jfile.open("rb") as fh:
                jev_pool = pickle.load(fh)
            log({"event": "pool_cache", "source": "jev", "loaded": len(jev_pool)})
        else:
            mine = [rec for i, rec in enumerate(read_jsonl(args.jev, "train", args.jev_limit))
                    if i % world == rank]

            def build_jev(rec: dict[str, Any]) -> Example | None:
                if jevbench_guard.match(rec["state"]):
                    return None
                ex = jev_example(model, rec, random.Random(f"{args.seed}:jev:{rec.get('uid')}"))
                if ex is not None and len(ex.ctx.ids) <= args.max_context_tokens:
                    return ex
                return None

            jev_pool = [ex for ex, _ in pool_map(build_jev, mine, args.pool_workers)
                        if ex is not None]
            # jev states not kept: JevBench guard matches, empty or over-long examples
            jevbench_dropped["jev"] = len(mine) - len(jev_pool)
            if jfile is not None:
                jfile.parent.mkdir(parents=True, exist_ok=True)
                with jfile.open("wb") as fh:
                    pickle.dump(jev_pool, fh, protocol=pickle.HIGHEST_PROTOCOL)
    log({"event": "data", "deg_train": sum(len(x.pool) for x in sources),
         "sources": {x.name: {"path": str(x.path), "weight": x.weight, "train": len(x.pool),
                              "long": sum(isinstance(e, LongDoc) for e in x.pool),
                              "active": any(x is a for a in active_sources(sources, min_train))}
                     for x in sources},
         "jev_train": len(jev_pool),
         "jevbench_guard": {"states": jevbench_guard.loaded, "dropped": jevbench_dropped},
         "filler_docs": len(filler.docs) if filler is not None else 0,
         "anchor_count": count_aux,
         "world": world, "array_mode": args.array_mode, "anchored_arrays": arrays["anchored"],
         "anchor_skipped_arrays": arrays["skipped"]})
    if args.pool_only:
        return

    params_bb = [p for p in model.backbone.parameters() if p.requires_grad]
    params_hd = [p for p in model.heads.parameters() if p.requires_grad]  # frozen heads: none
    use_zero = bool(params_bb) and not args.no_zero and (world > 1 or args.zero_sim_world > 1)
    zero: ZeroMaster | None = None
    master: list[torch.Tensor] | ZeroMaster
    if use_zero:
        if world > 1 and args.zero_sim_world:
            raise SystemExit("--zero-sim-world is for one process")
        zero = ZeroMaster(params_bb, rank, world, sim_world=args.zero_sim_world)
        master = zero
        master_params = [zero.master]
        log({"event": "zero", "world": zero.world, "simulated": zero.sim,
             "backbone_params": zero.N, "shard": zero.shard_len,
             "buckets": len(zero.buckets)})
    else:
        master = [p.detach().float().clone().requires_grad_(True) for p in params_bb]
        master_params = master
    groups = ([{"params": master_params, "lr": args.lr}] if master_params else []) + (
        [{"params": params_hd, "lr": args.head_lr}] if params_hd else [])
    opt = torch.optim.AdamW(groups, betas=(0.9, 0.98), weight_decay=0.01, fused=True)
    base_lrs = ([args.lr] if master_params else []) + ([args.head_lr] if params_hd else [])
    weights = make_weights(args)
    planner: MemPlanner | None = None
    n_bb = sum(q.numel() for q in params_bb)
    n_hd = sum(q.numel() for q in params_hd)
    grad_bytes = 2.0 * n_bb  # bf16 backbone grads while a backward runs
    if args.mem_target_gb > 0:
        w_bytes = sum(q.numel() * q.element_size() for m in model.modules()
                      for q in m.parameters())
        planner = MemPlanner(
            target=args.mem_target_gb * GIB,
            static=(sharded_static_bytes(n_bb, n_hd, w_bytes, zero.world) if zero is not None
                    else static_bytes(n_bb, n_hd, w_bytes, accumulate=False)),
            # ZeRO: a several-part step sums its gradient on the host
            accum_extra=4.0 * n_hd if zero is not None else 4.0 * (n_bb + n_hd),
            fork_tokens=args.fork_tokens,
            fork_rows=args.fork_rows,
            coef=MemCoef.load(args.mem_coef, model.meta["backbone"]))
        log({"event": "mem_plan", "target_gib": args.mem_target_gb,
             "static_gib": round(planner.static / GIB, 2),
             "activation_budget_gib": round(planner.budget() / GIB, 2),
             "coef": dataclasses.asdict(planner.coef)})
    static_analytic = planner.static if planner else 0.0
    gen = batches(sources, jev_pool, rng, args.budget, args.jev_frac, args.max_records,
                  min_train=min_train, materialize=materialize,
                  fits=planner.fits if planner else None,
                  augment=augment if filler is not None else None,
                  cap=lambda ex, r: cap_forks(ex, r, args.budget, args.record_fork_tokens,
                                              args.decide_forks_per_field))
    mem_log = args.out / "mem.jsonl"
    mon = GpuMonitor()
    window: dict[str, list[float]] = defaultdict(list)
    tok_count = {"real": 0.0, "padded": 0.0}
    win = {"windows": 0, "window_targets": 0, "window_dropped": 0, "filler": 0,
           "filler_negative": 0, "filler_tokens": 0, "subsampled_records": 0,
           "subsampled_forks": 0}
    oom = {"oom_retry": 0, "oom_skip": 0, "oom_redraw": 0, "oom_step_skipped": 0,
           "mem_parts": 0, "mem_split_steps": 0, "trunk_shared": 0, "parts_capped_forks": 0}
    peak_win = 0.0
    t0 = time.time()
    for step in range(args.start_step + 1, args.steps + 1):
        live = backbone_live(step, args.start_step, args.heads_only_steps)
        if params_bb and live != params_bb[0].requires_grad:
            for p in params_bb:
                p.requires_grad_(live)
        scale = lr_scale(step, args.steps, args.warmup, args.start_step)
        bb_scale = lr_scale(step, args.steps, args.warmup,
                            args.start_step + args.heads_only_steps) if live else 0.0
        for i, (g, lr) in enumerate(zip(opt.param_groups, base_lrs)):
            g["lr"] = lr * (bb_scale if master_params and i == 0 else scale)  # 0 = backbone
        opt.zero_grad(set_to_none=True)
        if planner is not None:  # resident model + optimizer state, as measured now
            resident = torch.cuda.memory_allocated() + grad_bytes
            planner.static = max(static_analytic, float(resident))
        pre_backward = None
        if args.loss_gnorm_every and params_bb and live and step % args.loss_gnorm_every == 0:
            def pre_backward(losses: dict[str, torch.Tensor]) -> None:
                for k, v in losses.items():
                    w = getattr(weights, k, 0.0) if k in LOSS_KEYS else 0.0
                    if w and v.requires_grad:
                        window[f"gnorm_{k}"].append(loss_grad_norm(w * v, params_bb))
        result: StepResult | None = None
        agreed = False
        for attempt in range(args.oom_max_redraws + 1):
            batch = next(gen)
            if planner is not None and args.max_record_parts > 0:
                batch, left = cap_record_parts(planner, batch, args.max_record_parts, rng)
                oom["parts_capped_forks"] += left
            for ex in batch:
                if ex.aug:
                    win["filler" if ex.aug == "filler" else "filler_negative"] += 1
                    win["filler_tokens"] += len(ex.ctx.ids)
                if ex.subsampled:
                    win["subsampled_records"] += 1
                    win["subsampled_forks"] += ex.subsampled
                if ex.window is not None:
                    win["windows"] += 1
                    win["window_targets"] += len(ex.items)
                    win["window_dropped"] += ex.dropped
            parts = planner.plan(batch) if planner is not None else [batch]
            result = run_step(model, parts, weights, args.fork_tokens, args.fork_rows,
                              params_bb, master, params_hd, live, args.oom_min_rows,
                              pre_backward=pre_backward, grad_bytes=grad_bytes,
                              share_trunk=not args.no_share_trunk,
                              total_n=global_counts(batch, world, model.device))
            oom["oom_retry"] += result.oom_retry
            oom["oom_skip"] += len(result.oom_skip)
            oom["mem_parts"] += result.parts
            oom["trunk_shared"] += result.trunk_shared
            oom["mem_split_steps"] += len(parts) > 1
            peak_win = max(peak_win, result.max_allocated)
            for ev in result.oom_skip:
                log({**ev, "step": step, "attempt": attempt})
            if planner is not None:
                for part, peak in result.peaks:
                    planner.observe(part, peak)
                    if rank == 0:
                        with mem_log.open("a") as fh:
                            fh.write(json.dumps({
                                "step": step, "peak": peak,
                                "est": planner.estimate(part),
                                "scale": planner.scale, "static": planner.static,
                                **planner_features(planner, part)}) + "\n")
            agreed = all_ranks_ok(result.ok, world, model.device)
            if agreed:
                break
            _discard_grads(params_bb, master, params_hd)  # another rank skipped: all redraw
            oom["oom_redraw"] += 1
        assert result is not None
        if not agreed:
            _discard_grads(params_bb, master, params_hd)
            oom["oom_step_skipped"] += 1
            log({"event": "oom_step_skipped", "step": step,
                 "redraws": args.oom_max_redraws})
        elif zero is not None:
            if live:
                zero.reduce_scatter()  # sum over ranks -> this rank's fp32 shard
            _allreduce_grads(params_hd, world, average=False)
            gnorm = clip_global_(zero if live else None, params_hd, 1.0)
            opt.step()
            if live:
                zero.gather_params()
        else:
            stepped = (master if live else []) + params_hd
            _allreduce_grads(stepped, world, average=False)
            gnorm = torch.nn.utils.clip_grad_norm_(stepped, 1.0)
            opt.step()
            with torch.no_grad():
                for p, mp in zip(params_bb, master):
                    p.copy_(mp)
        if agreed:
            for k, v in sum_over_ranks(result.losses, world, model.device).items():
                window[k].append(float(v))
            window["gnorm"].append(float(gnorm))
            stats = result.stats
            tok_count["real"] += stats["trunk_tokens"] + stats["fork_tokens"]
            tok_count["padded"] += stats["trunk_padded"] + stats["fork_padded"]
        if step % args.log_every == 0 or step == args.steps:
            torch.cuda.synchronize()
            dt = time.time() - t0
            row = {"step": step,
                   "lr": round(base_lrs[0] * (bb_scale if master_params else scale), 8),
                   **{k: round(sum(v) / len(v), 4) for k, v in window.items()},
                   "tok_per_s": round(tok_count["real"] * world / dt, 1),
                   "padded_tok_per_s": round(tok_count["padded"] * world / dt, 1),
                   "s_per_step": round(dt / args.log_every, 3),
                   # device peak over this log window (peak stats reset per micro-batch)
                   "max_mem_gib": round(peak_win / GIB, 2),
                   **oom, **win, **mon.take()}
            if planner is not None:
                row["mem_scale"] = round(planner.scale, 3)
            log(row)
            window.clear()
            tok_count = {"real": 0.0, "padded": 0.0}
            win = dict.fromkeys(win, 0)
            oom = dict.fromkeys(oom, 0)
            peak_win = 0.0
            t0 = time.time()
        if args.reload_every and step % args.reload_every == 0:
            added = load_deg()
            if any(added.values()):
                log({"event": "reload", "added": added,
                     "train": {x.name: len(x.pool) for x in sources},
                     "active": [x.name for x in active_sources(sources, min_train)]})
        if args.save_every and step % args.save_every == 0:
            if world > 1:  # every rank must hold the same weights (rank 0 saves them)
                same = replicas_equal(model, world)
                log({"event": "replicas", "step": step, "equal": same})
                if not same:
                    raise RuntimeError(f"step {step}: model replicas diverged across ranks")
            if rank == 0:
                save(model, args.out / f"step-{step:07d}", args)
    mon.stop()
    if rank == 0 and args.no_final_eval:
        save(model, args.out / "final", args)
    elif rank == 0:
        model.backbone.eval()
        model.heads.eval()
        from deg.evaluate import calibrate, evaluate

        temps = calibrate(model, [x.path for x in sources if x.path.exists()], JEV_DEV,
                          args.eval_jev, source="datagen-val",
                          max_context_tokens=args.max_context_tokens)
        log({"event": "calibrated", "temperature": temps})
        save(model, args.out / "final", args)
        report = evaluate(model, sources[0].path, JEV_DEV, args.eval_records, args.eval_jev)
        (args.out / "eval.json").write_text(json.dumps(report, indent=2))
        log({"event": "eval", **report})
    if world > 1:
        import torch.distributed as dist

        dist.barrier()
        dist.destroy_process_group()


def init_heads(model: DegModel, ckpt: Path) -> None:
    """Load the heads of deg checkpoint ``ckpt`` (its backbone must be ``model``'s)."""
    from safetensors.torch import load_file

    meta = read_meta(ckpt)
    want = meta.get("backbone", "qwen3.5-0.8b")
    if want != model.meta["backbone"]:
        raise SystemExit(f"--init-heads {ckpt} is a {want} checkpoint, not "
                         f"{model.meta['backbone']}")
    state = load_file(str(Path(ckpt) / "heads.safetensors"))
    missing, unexpected = model.heads.load_state_dict(state, strict=False)
    if missing or unexpected:
        raise SystemExit(f"--init-heads {ckpt}: heads mismatch {(missing + unexpected)[:5]}")


def zlib_shard(uid: str, world: int) -> int:
    import zlib

    return zlib.crc32(uid.encode()) % world


def save(model: DegModel, path: Path, args: argparse.Namespace) -> None:
    model.save(path, {"train_args": {k: str(v) for k, v in vars(args).items()},
                      "temperature": dict(zip(KINDS, model.heads.temperature.tolist()))})


if __name__ == "__main__":
    main()