review round 3 (probe/model-card/eval): rank-visible eval failure + cross-rank peak memory + complete run_summary; card derived-and-asserted from config; harness install --no-deps and task ids confirmed; torchrun redirects 3 + tee 3

train/train_ounce100m.py

@@ -1,703 +1,760 @@
-# The training run in one script: frozen architecture, map-style packed data, fp16 on T4, trapezoid LR,
-# and a checkpoint cycle that satisfies §3.13 by construction.
-
-# Deliberately thin glue over `transformers.Trainer` (docs/01-plan.md §5): data loading, DDP, gradient
-# accumulation, scaler, RNG capture and resumption are the framework's job. What is ours is the shard
-# reader, the LR shape, and the upload/verify/prune discipline -- the three things a generic stack cannot
-# know about this project.
-
-# Two resume facts that are easy to get wrong and are asserted here rather than assumed:
-#   1. `ignore_data_skip=True`. Trainer's own resume would replay discarded batches; we advance the data
-#      position by slicing the permutation in shard_dataset instead. Leaving the flag at False would
-#      double-skip and silently LOSE training data -- the exact §3.1 violation.
-#   2. The cursor is checked against the checkpoint, not trusted. A resumed run whose cursor, seed,
-#      sequence length or corpus fingerprint disagrees with what is on disk stops before step 1.
-
-import argparse
-import json
-import math
-import os
-import time
-
-import torch
-from torch.optim.lr_scheduler import LambdaLR
-from transformers import (LlamaConfig, LlamaForCausalLM, Trainer, TrainerCallback,
-                          TrainingArguments, default_data_collator, set_seed)
-
-import hubckpt
-import shard_dataset as SD
-
-# ---- frozen at Gate 1 (docs/01-plan.md §2.3, §4). Only --layers/--hidden change, for test T1.
-BASE = dict(hidden=576, layers=22, heads=9, kv_heads=3, ffn=1536, vocab=49152,
-            rope_theta=10000.0, rms_eps=1e-5, tie=True)
-
-
-def _kv_dividing(heads, hidden):
-    """GQA requires num_key_value_heads to divide num_attention_heads; SDPA raises otherwise
-    ("Number of heads in key and value must divide the number of heads in Q"). Variant shapes are built
-    from hidden/64, which need not divide 9 -- take the largest divisor of `heads` that fits."""
-    per_head = max(1, hidden // 64)
-    return max(d for d in range(1, min(per_head, heads) + 1) if heads % d == 0)
-
-
-def build_config(seq_len, layers=None, hidden=None, attn="eager"):
-    """The frozen shape, and the two variants test T1 needs (20L at the same width, and a wider/shallower
-    alternative). SwiGLU width is int(2*4h/3) rounded to a 64-multiple -- 1536 at h=576, which is the
-    value the frozen 106,194,240 count was derived from. Getting this factor wrong silently changes the
-    parameter count and makes T1 compare two different models.
-
-    `attn` is explicit because it is a measured choice, not a default: probe p0c found eager attention
-    beating memory-efficient SDPA at seq 1024 on Turing (docs/00-platform-notes.md §9), and Gate 3 P2
-    then clocked 4,071 tok/s with the transformers default -- half the rate p0c saw in a raw loop. Leaving
-    the attention kernel to whatever the library defaults to is how a 68-hour run becomes a 27-hour run.
-    """
-    h = hidden or BASE["hidden"]
-    L = layers or BASE["layers"]
-    # "variant" means a different WIDTH, not merely a different layer count: the frozen GQA ratio (9 Q /
-    # 3 KV) belongs to the frozen width, and passing --hidden 576 explicitly must not quietly turn the
-    # model into MHA-9 and make test T1 measure something else.
-    frozen_width = (h == BASE["hidden"])
-    inter = BASE["ffn"] if frozen_width else (int(round(8 * h / 3)) // 64) * 64
-    kv = BASE["kv_heads"] if frozen_width else _kv_dividing(BASE["heads"], h)
-    cfg = LlamaConfig(vocab_size=BASE["vocab"], hidden_size=h, intermediate_size=inter,
-                      num_hidden_layers=L, num_attention_heads=BASE["heads"],
-                      num_key_value_heads=kv,
-                      max_position_embeddings=seq_len, rope_theta=BASE["rope_theta"],
-                      rms_norm_eps=BASE["rms_eps"], tie_word_embeddings=BASE["tie"],
-                      attention_dropout=0.0, hidden_act="silu", mlp_bias=False,
-                      attn_implementation=attn,
-                      pad_token_id=BASE["vocab"] - 1, eos_token_id=0)
-    return cfg
-
-
-def count_params(model):
-    """Total, embeddings included -- the number §2 budgets against. Reported both ways: numel sum and
-    the unique-parameter count, which differ when embeddings are tied."""
-    total = sum(p.numel() for p in model.parameters())
-    uniq = sum(p.numel() for p in set(model.parameters()))
-    return {"sum_numel": total, "unique": uniq,
-            "trainable": sum(p.numel() for p in model.parameters() if p.requires_grad)}
-
-
-# The cursor arithmetic in one place, because three separate bugs live here:
-#   - a step consumes micro_batch * world * accum sequences, not micro_batch (DDP shards the batch);
-#   - Trainer restores `global_step` from the checkpoint, so on resume it is ALREADY absolute. Adding the
-#     resume point a second time inflates the cursor, and an inflated cursor makes the next run skip data
-#     it never read -- silent data loss with no error anywhere, which is §3.1's worst case. Writing this
-#     as a pure function of the absolute step removes the whole class of bug.
-#   - max_steps is the TOTAL horizon (it is what the LR schedule decays against), so resuming mid-run
-#     must keep the same steps_planned or the trapezoid changes shape.
-def samples_at(absolute_step, seqs_per_step):
-    return int(absolute_step) * seqs_per_step
-
-
-class TrapezoidTrainer(Trainer):
-    """Trainer with (a) a hand-rolled trapezoid schedule and (b) a data position we own."""
-
-    def __init__(self, *a, lr_shape=None, **kw):
-        self.lr_shape = lr_shape
-        super().__init__(*a, **kw)
-
-    def create_scheduler(self, num_training_steps, optimizer=None):
-        warm = self.lr_shape["warmup_steps"]
-        decay_at = self.lr_shape["decay_start"]
-        total = self.lr_shape["total_steps"]
-
-        def fn(cur):
-            if cur < warm:
-                return (cur + 1) / max(1, warm)
-            if cur < decay_at:
-                return 1.0
-            # linear decay to exactly zero at the end: the trapezoid in docs/01-plan.md §4.3
-            return max(0.0, 1.0 - (cur - decay_at) / max(1, total - decay_at))
-
-        if self.lr_scheduler is None:
-            self.lr_scheduler = LambdaLR(self.optimizer if optimizer is None else optimizer, fn)
-        return self.lr_scheduler
-
-
-class HubPush(TrainerCallback):
-    """§3.13 after every checkpoint: write the cursor in, upload, verify from the Hub, prune locally.
-
-    Only rank 0 acts -- the ranks share the checkpoint directory and would race. A failed verification
-    raises inside the callback, so the run stops AT that checkpoint rather than continuing with its only
-    copy on a disk that can be taken away at any moment.
-    """
-
-    def __init__(self, repo, api, token, seq_len, seed, fingerprint, seqs_per_step,
-                 prune=True, min_free_gb=6.0, repo_type="dataset", stop_after_steps=0, order_sha=None):
-        self.repo, self.api, self.token = repo, api, token
-        self.seq_len, self.seed, self.fingerprint = seq_len, seed, fingerprint
-        self.seqs_per_step, self.repo_type = seqs_per_step, repo_type
-        self.order_sha = order_sha
-        self.prune, self.min_free_gb = prune, min_free_gb
-        self.stop_after = int(stop_after_steps or 0)
-        self.stopped_at = None
-        self.n = 0
-        self.steps = []
-        self.results = []
-
-    def on_step_end(self, args, state, control, **kw):
-        # A session that ends ON a verified checkpoint loses nothing when the platform takes the disk.
-        # That is how the 10 % cadence is made compatible with an unknown session cap: instead of pushing
-        # extra intermediate copies (which multiplies Hub LFS storage by the number of sessions), the run
-        # is cut into segments whose length is a whole number of checkpoint intervals.
-        if self.stop_after and state.global_step >= self.stop_after:
-            self.stopped_at = state.global_step
-            control.should_training_stop = True
-        return control
-
-    def on_save(self, args, state, control, **kw):
-        if not state.is_world_process_zero:
-            return control
-        step = state.global_step
-        d = os.path.join(args.output_dir, f"checkpoint-{step}")
-        if not os.path.isdir(d):
-            raise RuntimeError(f"Trainer reported a save but {d} does not exist")
-        samples = samples_at(step, self.seqs_per_step)
-        SD.save_cursor(d, SD.cursor_dict(self.seq_len, samples, self.seed, self.fingerprint,
-                                         step=step, order_sha=self.order_sha))
-        t0 = time.time()
-        # Upload and verify *without* pruning, roll the pointer and read it back, and only then delete the
-        # local copy. The old order (prune inside push_and_prune, pointer written afterwards) meant one
-        # failed pointer commit left a Hub-verified checkpoint that nothing pointed at and no local copy to
-        # recover it with -- and every guard a resume runs would still pass, because weights, cursor,
-        # fingerprint and seed are all self-consistent. Only the pointer would be wrong, and the next
-        # session would quietly re-read 100 M tokens.
-        res = hubckpt.push_and_prune(self.repo, d, f"ckpt/checkpoint-{step}", self.api,
-                                     repo_type=self.repo_type, token=self.token, prune=False)
-        ptr_body = {"path_in_repo": f"ckpt/checkpoint-{step}", "step": step,
-                    "samples_consumed": samples, "tokens_consumed": samples * self.seq_len,
-                    "fingerprint": self.fingerprint, "seed": self.seed}
-        hubckpt.write_latest_pointer(self.repo, ptr_body, self.api, repo_type=self.repo_type,
-                                     token=self.token)
-        if self.prune:
-            hubckpt.prune_local(d, res, self.min_free_gb)
-        loss = next((h["loss"] for h in reversed(state.log_history) if "loss" in h), None)
-        print(f"CKPT {step} samples={samples:,} push+verify={round(time.time()-t0,1)}s "
-              f"ok={res['verify']['ok']} pruned={res.get('pruned')} free={res['free_after_gb']}GB "
-              f"loss={loss}", flush=True)
-        self.n += 1
-        self.steps.append(step)
-        self.results.append(res)
-        return control
-
-
-def run_benchmark(args, model, dataset, steps, world, seqs_per_step):
-    """The same data path, optimizer and fp16 recipe as the real run, with no Trainer in the loop.
-
-    Why this exists: P2 measured 5,105 tok/s for 22L/seq-1024 where probe p0c measured 11,062 tok/s for a
-    near-identical shape with a hand-rolled loop. The gap is either (a) the attention kernel, (b) per-step
-    framework overhead, or (c) a measurement artefact, and which one it is decides whether the main run
-    needs 27 h or 68 h. Guessing wrong in either direction is expensive, so measure both paths.
-    """
-    import torch.distributed as dist
-    from torch.nn.parallel import DistributedDataParallel as DDP
-    from torch.utils.data import DataLoader
-
-    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
-    device = torch.device("cuda", local_rank)
-    torch.cuda.set_device(device)
-    if world > 1 and not dist.is_initialized():
-        dist.init_process_group("nccl")
-    model = model.to(device)
-    if args.grad_ckpt:
-        model.gradient_checkpointing_enable()
-        model.config.use_cache = False
-    net = DDP(model, device_ids=[local_rank]) if world > 1 else model
-
-    opt = torch.optim.AdamW(net.parameters(), lr=args.lr, betas=(0.9, 0.95),
-                            weight_decay=0.1, eps=1e-8)
-    try:
-        scaler = torch.amp.GradScaler("cuda", enabled=True)     # torch >= 2.4
-    except TypeError:
-        scaler = torch.cuda.amp.GradScaler(enabled=True)        # older signature
-    # Same reader, same shuffle, no accumulation: the comparison against the Trainer cells is then
-    # per-step cost with the framework removed, and nothing else differs.
-    dl = DataLoader(dataset, batch_size=args.micro_batch, shuffle=False, num_workers=2,
-                    pin_memory=False, drop_last=True)
-    t0 = time.time()
-    toks = 0
-    for i, batch in enumerate(dl):
-        if i >= steps:
-            break
-        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
-        opt.zero_grad(set_to_none=True)
-        with torch.autocast("cuda", dtype=torch.float16):
-            out = net(**batch)
-        scaler.scale(out.loss).backward()
-        scaler.unscale_(opt)
-        torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
-        scaler.step(opt)
-        scaler.update()
-        toks += batch["input_ids"].numel()
-        if i == 0:
-            torch.cuda.synchronize()
-            t0 = time.time()   # exclude CUDA context + first-trace cost from the rate
-            toks = 0
-    sec = time.time() - t0
-    res = {"benchmark_steps": i + 1 if steps else 0, "tok_per_s": round(toks / max(1e-6, sec)),
-           "seconds": round(sec, 1), "tokens": toks,
-           "peak_gpu_gb": round(torch.cuda.max_memory_allocated() / 1e9, 2),
-           "world": world, "seq_len": args.seq_len, "micro_batch": args.micro_batch,
-           "attn": args.attn, "grad_ckpt": bool(args.grad_ckpt),
-           "params": sum(p.numel() for p in model.parameters())}
-    print("BENCH_JSON " + json.dumps(res, sort_keys=True), flush=True)
-    if world > 1:
-        dist.barrier()
-        dist.destroy_process_group()
-    return res
-
-
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--root", default="/kaggle/working/mixroot")
-    ap.add_argument("--out", default="/kaggle/working/run")
-    ap.add_argument("--hub-repo", default="", help="dataset repo for checkpoints; empty = no upload")
-    ap.add_argument("--hub-repo-type", default="dataset", choices=["dataset", "model"])
-    # D-011: these ARE the frozen main-run values (measured 7,732-7,828 tok/s at 1024/eager on 2xT4), so
-    # that launching the run does not depend on remembering a flag string.
-    ap.add_argument("--seq-len", type=int, default=1024)
-    ap.add_argument("--attn", default="eager", choices=["eager", "sdpa", "flash_attention_2"],
-                    help="p0c measured eager beating SDPA on Turing; Gate 3 P2/A-B measured +52 %% for "
-                         "eager at seq 1024 and an OOM at 2048. flash_attention_2 needs sm80+")
-    ap.add_argument("--grad-ckpt", default=True, action=argparse.BooleanOptionalAction,
-                    help="on by default: it is what keeps peak memory at ~7 GB instead of 12 GB, and the "
-                         "measured rate cost at 1024 was nil (7,828 with GC at m8 vs 7,732 without at m4)")
-    ap.add_argument("--layers", type=int, default=0)
-    ap.add_argument("--hidden", type=int, default=0)
-    ap.add_argument("--micro-batch", type=int, default=4,
-                    help="4 x 1024 x 32 accum x 2 cards = 262,144 tokens/step, the frozen batch size")
-    ap.add_argument("--accum", type=int, default=32)
-    ap.add_argument("--lr", type=float, default=6e-4)
-    ap.add_argument("--optim", default="adamw_torch",
-                    help="only mathematically-equivalent implementations are fair game at freeze time "
-                         "(adamw_torch / adamw_torch_fused / adamw_apex_fused); 8-bit variants change the "
-                         "update, so they are excluded by rule, not by preference")
-    ap.add_argument("--gc-nonreentrant", default=False, action=argparse.BooleanOptionalAction,
-                    help="pass use_reentrant=False to gradient checkpointing: same recomputation, "
-                         "different autograd path, and reports on whether it is faster under DDP conflict")
-    ap.add_argument("--torch-compile", default=False, action=argparse.BooleanOptionalAction,
-                    help="Trainer's torch_compile hook; fused kernels, no change to the update rule")
-    ap.add_argument("--tokens", type=int, default=1_000_000_000)
-    ap.add_argument("--max-steps", type=int, default=0, help="0 = derive from --tokens; >0 overrides")
-    ap.add_argument("--warmup-frac", type=float, default=0.02)
-    ap.add_argument("--decay-frac", type=float, default=0.80)
-    ap.add_argument("--seed", type=int, default=20260919)
-    ap.add_argument("--data-seed", type=int, default=20260919)
-    ap.add_argument("--resume", default="", help="checkpoint dir; empty = auto from Hub latest")
-    ap.add_argument("--val-tokens", type=int, default=2_000_000)
-    ap.add_argument("--log-every", type=int, default=10)
-    ap.add_argument("--push-every-steps", type=int, default=0, help="0 = only at 10%% boundaries")
-    ap.add_argument("--stop-after-steps", type=int, default=0,
-                    help="absolute global step at which THIS session stops training, so it ends on a "
-                         "verified checkpoint instead of being killed between two of them. Must be a "
-                         "whole multiple of the checkpoint interval; 0 = run to the horizon.")
-    ap.add_argument("--prune", action="store_true", help="delete local checkpoint after hub verify")
-    ap.add_argument("--dry-run", action="store_true", help="build everything, train 0 steps")
-    ap.add_argument("--benchmark", type=int, default=0,
-                    help="N steps with no Trainer in the loop (isolates framework per-step cost)")
-    args = ap.parse_args()
-
-    rank = int(os.environ.get("RANK", "0"))
-    world = int(os.environ.get("WORLD_SIZE", "1"))
-    say = (lambda *m: print(*m, flush=True)) if rank == 0 else (lambda *m: None)
-
-    man_path = os.path.join(args.root, "manifest.json")
-    if not os.path.exists(man_path):
-        raise SystemExit(f"no manifest at {man_path}: build the mix first (Phase 2 is not optional)")
-    man = json.load(open(man_path))
-    if man["total_tokens"] < args.tokens:
-        say(f"WARNING mix has {man['total_tokens']:,} tokens < target {args.tokens:,}")
-
-    store = SD.PackedTokenStore(args.root, man)
-    fingerprint = SD.mix_fingerprint(args.root, man)
-    # One definition of the step, used by the cursor, the schedule and the token count: a step consumes
-    # micro_batch * world sequences (DDP shards the batch), so tokens/step follows from it.
-    seqs_per_step = args.micro_batch * world * args.accum
-    tokens_per_step = seqs_per_step * args.seq_len
-    steps_planned = args.max_steps or int(args.tokens / tokens_per_step)
-    say(f"mix: {store.total_tokens:,} tokens in {len(store.shards)} shards fp={fingerprint} "
-        f"| {tokens_per_step:,} tok/step x {steps_planned} steps = "
-        f"{tokens_per_step * steps_planned:,} tokens")
-
-    # ---- resume point: the cursor is the run's identity as much as the weights are.
-    # `--resume auto` is what a fresh instance with an empty disk passes: it asks the Hub where the run
-    # stopped and pulls that checkpoint down. §3.1 forbids recovering from anything left on an instance.
-    if args.resume == "auto":
-        if not args.hub_repo:
-            raise SystemExit("--resume auto needs --hub-repo")
-        import ounce100m_credentials
-        ounce100m_credentials.install()
-        from huggingface_hub import HfApi
-        tk = os.environ.get("HF_TOKEN")
-        ptr = hubckpt.latest_pointer(args.hub_repo, repo_type=args.hub_repo_type, token=tk)
-        if ptr.get("not_found"):
-            # A 404 is normally a fact -- "this run has never checkpointed" -- and step 0 is then right.
-            # But a missing pointer beside existing checkpoints means the pointer upload failed after a
-            # good checkpoint, and starting from zero would silently retrain the corpus from the top,
-            # which §3.1 forbids and E-018 is about. So list the repo before believing the 404.
-            try:
-                listed = sorted(k for k in hubckpt.hub_listing(args.hub_repo, args.hub_repo_type,
-                                                               token=tk)
-                                if k.startswith("ckpt/checkpoint-"))
-            except hubckpt.RepoMissing:
-                # The checkpoint repo has never been created: nothing has ever been pushed, so step 0 is
-                # the only correct answer. The first session of the run always takes this branch, and
-                # treating a 404 as "listing failed" would refuse it (a generic `except Exception` below
-                # makes the list non-empty and the refusal fires).
-                say(f"{args.hub_repo} does not exist yet -- nothing has ever been pushed")
-                listed = []
-            except Exception as e:
-                listed = [f"<listing failed: {type(e).__name__}>"]
-            if listed:
-                raise SystemExit(f"--resume auto found no latest.json on {args.hub_repo}, but the repo "
-                                 f"holds {len(listed)} checkpoint path(s) ({listed[:3]}). The pointer is "
-                                 "missing, not the run -- refusing to start from step 0.")
-            say(f"no latest.json and no checkpoints on {args.hub_repo} -> starting from step 0")
-            args.resume = ""
-        elif "error" in ptr:
-            # Anything else (401, timeout, DNS) means the lookup failed, not that there is nothing there.
-            # Falling through to step 0 here would silently retrain the first stretch of the run, which is
-            # exactly the §3.1 violation this whole branch exists to prevent (E-018's shape).
-            raise SystemExit(f"--resume auto could not read the hub pointer: {ptr['error']}")
-        else:
-            # A pointer is a resume point only if nothing newer was ever pushed. This scan catches the one
-            # case write_latest_pointer's own read-back cannot see: a pointer left behind by a rollback, a
-            # deletion, or a session that died between its checkpoint push and its pointer roll. Resuming
-            # from the stale pointer would re-train an interval with weights and a cursor that agree with
-            # each other and disagree with the Hub.
-            import re as _re
-            steps_on_hub = set()
-            for attempt in (1, 2):
-                try:
-                    steps_on_hub = {int(m.group(1)) for k in
-                                    hubckpt.hub_listing(args.hub_repo, args.hub_repo_type, token=tk)
-                                    if (m := _re.match(r"ckpt/checkpoint-(\d+)/", k))}
-                    break
-                except Exception as e:
-                    say(f"   checkpoint scan attempt {attempt}: {type(e).__name__}: {str(e)[:140]}")
-                    time.sleep(15)
-            else:
-                raise SystemExit("--resume auto could not enumerate the checkpoint repo, so it cannot "
-                                 f"prove latest.json (step {ptr.get('step')}) is the newest state")
-            if steps_on_hub and max(steps_on_hub) > int(ptr.get("step") or 0):
-                raise SystemExit(f"latest.json points at step {ptr.get('step')} but {args.hub_repo} also "
-                                 f"holds checkpoint-{max(steps_on_hub)}: the pointer is behind the run. "
-                                 "Resuming here would re-train that interval, so refusing.")
-            # Both ranks need the checkpoint on disk, and this block runs before TrainingArguments has
-            # created a process group -- so there is no barrier to synchronise on. Downloading into a
-            # per-rank directory instead removes the race entirely (two writers, one path), at the cost of
-            # a second 1.7 GB pull that happens in parallel. Each rank then reads its own verified copy and
-            # derives the same start_sample/start_step from the same cursor.json bytes.
-            dest = os.path.join(args.out, f"resume_r{rank}",
-                                os.path.basename(ptr["path_in_repo"]))
-            say(f"auto-resume: hub says step {ptr.get('step')} at {ptr['path_in_repo']}")
-            got = hubckpt.download_checkpoint(args.hub_repo, ptr["path_in_repo"], dest, HfApi(),
-                                              repo_type=args.hub_repo_type, token=tk)
-            say(f"downloaded {got['files']} files, {got['bytes']/1e6:.0f} MB in {got['seconds']}s")
-            if got["bytes"] < 1_000_000:
-                raise SystemExit("auto-resume downloaded a checkpoint that is not real weights")
-            args.resume = got["dir"]
-            say(f"resuming from {args.resume}")
-    cur = SD.load_cursor(args.resume) if args.resume else None
-    if args.resume and cur is None:
-        # Weights at step N with the data stream back at sample 0 would retrain the opening of the corpus
-        # and report a continuous loss curve. Refuse instead.
-        raise SystemExit(f"resume refused: {args.resume} has no cursor.json -- a checkpoint without a "
-                         "recorded data position cannot be resumed, only restarted (§3.1)")
-    if args.resume:
-        say(f"resume from {args.resume} cursor={cur}")
-    start_sample = (cur or {}).get("samples_consumed", 0)
-    start_step = (cur or {}).get("step") or 0
-    # Identity of the visit order, derived from the corpus size and the seed rather than trusted to the
-    # seed: numpy's permutation is only stable within a version, and Kaggle updates its image underneath
-    # long runs. ~50 ms once per session, and it is what makes "resumed onto a different order" loud.
-    order_sha = SD.perm_sha(store.total_tokens // args.seq_len, args.data_seed)
-    if cur:
-        for k, v in (("seq_len", args.seq_len), ("shuffle_seed", args.data_seed),
-                     ("shuffle_perm_sha", order_sha), ("dataset_files_sha", fingerprint)):
-            if cur.get(k) != v:
-                raise SystemExit(f"resume refused: cursor {k}={cur.get(k)} != current {v}. "
-                                 "A resumed run must read the same data in the same order (§3.1).")
-        if start_sample % seqs_per_step:
-            raise SystemExit(f"resume refused: cursor samples_consumed={start_sample} is not a whole "
-                             f"number of steps (seqs/step={seqs_per_step}) -- batch geometry changed")
-        if cur.get("step") is not None and cur["step"] * seqs_per_step != start_sample:
-            raise SystemExit(f"resume refused: cursor step={cur['step']} and samples_consumed="
-                             f"{start_sample} disagree at seqs/step={seqs_per_step}")
-        if start_step * seqs_per_step != start_sample:
-            start_step = start_sample // seqs_per_step
-
-    train_ds = SD.make_dataset(store, args.seq_len, args.data_seed, start_sample=start_sample)
-    assert getattr(train_ds, "order_sha", None) == order_sha, "two derivations of the order disagree"
-    # Validation comes from the mix's held-out val/ shards, never from the training stream: scoring on
-    # tokens the model already consumed would be a number nobody could reproduce (and §3.12).
-    vs = [s["file"] for s in man.get("val_shards", [])]
-    if not vs:
-        raise SystemExit("manifest lists no val_shards -- Gate 2 requires a held-out split")
-    vstore = SD.PackedTokenStore(args.root, man, files=vs)
-    n_val = min(max(1, args.val_tokens // args.seq_len), vstore.total_tokens // args.seq_len)
-    val_ds = SD.make_dataset(vstore, args.seq_len, args.data_seed, shuffle=False, count=n_val)
-    need = (steps_planned - start_step) * seqs_per_step
-    if len(train_ds) < need:
-        # Reading past the end of the stream makes the Trainer begin epoch 2 at the top of the same
-        # permutation: the run would then report the full token target while quietly re-training on the
-        # first windows. The mix is built with ~12 % headroom over the target for exactly this reason, so
-        # arriving here means that assumption stopped being true (review finding), and the honest response
-        # is to stop before step 1 rather than to discover it in the final report.
-        raise SystemExit(
-            f"dataset too small for the plan: {len(train_ds):,} windows available from sample "
-            f"{start_sample:,}, {need:,} needed for {steps_planned} steps x {seqs_per_step} sequences "
-            f"({need * args.seq_len:,} tokens). Lower --tokens to the 0.9 B fallback (D-012) or rebuild "
-            "the mix larger -- never let the epoch wrap.")
-    say(f"data: {len(train_ds):,} training windows available, {need:,} needed "
-        f"({100.0 * len(train_ds) / need - 100:.1f} % margin), {n_val:,} held-out val windows")
-
-    cfg = build_config(args.seq_len, args.layers or None, args.hidden or None, attn=args.attn)
-    set_seed(args.seed)  # before init: both ranks must start from identical weights, not two draws
-    model = LlamaForCausalLM(cfg).to(dtype=torch.float32)
-    pc = count_params(model)
-    say("params:", json.dumps(pc))
-    if not (90_000_000 <= pc["sum_numel"] <= 110_000_000) and not args.layers and not args.hidden:
-        raise SystemExit(f"parameter count {pc['sum_numel']:,} is outside §2's 90-110M budget")
-
-    if args.benchmark:
-        # Both ranks iterate the same samples here (no DistributedSampler): this path measures step cost,
-        # not a training run, and it must never be used to produce a checkpoint.
-        say(f"benchmark mode: {args.benchmark} steps, no Trainer, micro={args.micro_batch} "
-            f"seq={args.seq_len} attn={args.attn} world={world}")
-        run_benchmark(args, model, train_ds, args.benchmark, world, seqs_per_step)
-        return
-
-    # 10 checkpoints per run is §2's cadence; --push-every-steps overrides for the preflight.
-    ckpt_every = args.push_every_steps or max(1, steps_planned // 10)
-    if args.stop_after_steps:
-        if args.stop_after_steps % ckpt_every:
-            raise SystemExit(f"--stop-after-steps {args.stop_after_steps} is not a whole multiple of the "
-                             f"checkpoint interval {ckpt_every}: the session would end between two "
-                             f"checkpoints, which is exactly what this flag exists to prevent")
-        if args.stop_after_steps > steps_planned:
-            raise SystemExit(f"--stop-after-steps {args.stop_after_steps} exceeds the horizon "
-                             f"{steps_planned}; use 0 to run to the end")
-
-    api = None
-    if args.hub_repo:
-        import ounce100m_credentials
-        ounce100m_credentials.install()
-        from huggingface_hub import HfApi
-        api = HfApi()
-        # rank 0 only: two ranks committing the same .gitattributes to a brand-new repo is a write race
-        # on the one path that must not fail (only rank 0 ever uploads checkpoints, so rank 1 needs none)
-        if rank == 0:
-            say("ckpt repo:", json.dumps(hubckpt.ensure_repo(args.hub_repo, api,
-                                                             os.environ.get("HF_TOKEN"),
-                                                             args.hub_repo_type)))
-        say(f"checkpoint hub target {args.hub_repo} (prune={args.prune}, every {ckpt_every} steps)")
-
-    tacc = TrainingArguments(
-        output_dir=args.out,
-        per_device_train_batch_size=args.micro_batch,
-        # Same geometry for the validation pass: with the library default (8) the end-of-run eval measures
-        # a peak that has nothing to do with the training configuration, which is how an OOM in the no-
-        # checkpointing cell could be mistaken for a passing one (review finding).
-        per_device_eval_batch_size=args.micro_batch,
-        gradient_accumulation_steps=args.accum,
-        learning_rate=args.lr,
-        weight_decay=0.1,
-        adam_beta1=0.9, adam_beta2=0.95, adam_epsilon=1e-8,
-        max_grad_norm=1.0,
-        lr_scheduler_type="constant",   # replaced by TrapezoidTrainer.create_scheduler
-        warmup_ratio=0.0,               # the trapezoid owns warmup
-        num_train_epochs=1,
-        max_steps=0 if args.dry_run else steps_planned,
-        fp16=True, bf16=False,
-        # measured decision, not a default: A-B found eager attention 52 % faster at seq 1024 and OOM at
-        # 2048; checkpointing is what makes eager at 2048 possible at all, at the cost of a recompute pass
-        gradient_checkpointing=args.grad_ckpt,
-        # Same recomputation, different autograd path -- offered as a cells-vs-baseline question, not a
-        # belief. Reports on use_reentrant=False under DDP contradict each other across versions.
-        gradient_checkpointing_kwargs={"use_reentrant": not args.gc_nonreentrant}
-        if args.grad_ckpt else None,
-        optim=args.optim,
-        torch_compile=args.torch_compile,
-        ddp_find_unused_parameters=False,
-        # Rank 0 spends minutes inside on_save uploading ~1.7 GB, listing the repo and downloading one file
-        # back to hash it, while rank 1 waits at the next collective. The library's 1,800 s default would
-        # abort the run on a slow-egress Kaggle box -- precisely during the checkpoint that the other rank
-        # was protecting. Four hours makes the collective patient without hiding a real deadlock.
-        ddp_timeout=14400,
-        dataloader_num_workers=2,
-        dataloader_pin_memory=False,
-        remove_unused_columns=False,    # our dict keys must reach the model untouched
-        ignore_data_skip=True,          # see header note 1 -- we advance position ourselves
-        save_strategy="no" if args.dry_run else "steps",
-        save_steps=ckpt_every,
-        save_total_limit=1,             # §3.13: the Hub is the store, the instance is scratch
-        logging_steps=args.log_every, report_to=[],
-        seed=args.seed, data_seed=args.seed,
-        accelerator_config={"dispatch_batches": False},
-        average_tokens_across_devices=False,
-    )
-
-    tr = TrapezoidTrainer(
-        model=model, args=tacc, train_dataset=train_ds, eval_dataset=val_ds,
-        data_collator=default_data_collator, lr_shape={
-            "warmup_steps": max(100, int(steps_planned * args.warmup_frac)),
-            "decay_start": int(steps_planned * args.decay_frac),
-            "total_steps": steps_planned},
-    )
-    hub_cb = None
-    if api is not None:
-        hub_cb = HubPush(args.hub_repo, api, os.environ.get("HF_TOKEN"), args.seq_len,
-                         args.data_seed, fingerprint, seqs_per_step,
-                         prune=args.prune, repo_type=args.hub_repo_type,
-                         stop_after_steps=args.stop_after_steps, order_sha=order_sha)
-        tr.add_callback(hub_cb)
-
-    # Precision must be ASSERTED, not configured: bf16 "works" silently on Turing and fp32 master
-    # weights can be lost in a refactor (docs/00-platform-notes.md §8, memory/ERRORS.md E-006). The
-    # scaler's home moved between transformers versions, so look in both places instead of trusting one
-    # attribute path (E-010: prefer version-free surfaces).
-    scaler = getattr(tr, "scaler", None) or getattr(getattr(tr, "accelerator", None), "scaler", None)
-    eff_attn = (getattr(model.config, "_attn_implementation", None)
-                or getattr(model.config, "attn_implementation", None))
-    say(f"precision: fp16={tacc.fp16} bf16={tacc.bf16} scaler={type(scaler).__name__} "
-        f"master_dtype={next(model.parameters()).dtype} requested_attn={args.attn} "
-        f"effective_attn={eff_attn}")
-    assert tacc.fp16 and not tacc.bf16, "not the frozen T4 fp16 recipe"
-    assert next(model.parameters()).dtype == torch.float32, "master weights are not fp32"
-
-    prior = None
-    if args.resume:
-        tsp = os.path.join(args.resume, "trainer_state.json")
-        if os.path.exists(tsp):
-            h0 = [h for h in json.load(open(tsp)).get("log_history", []) if "loss" in h]
-            prior = h0[-1]["loss"] if h0 else None
-        say(f"prior loss at resume: {prior} -- loss continuity is judged against this number")
-
-    if args.dry_run:
-        say(json.dumps({"dry_run": True, "train_samples": len(train_ds), "val_samples": len(val_ds),
-                        "steps_planned": steps_planned, "ckpt_every": ckpt_every,
-                        "params": pc, "fingerprint": fingerprint, "scaler": type(scaler).__name__}))
-        return
-
-    t0 = time.time()
-    tr.train(resume_from_checkpoint=args.resume or None)
-
-    elapsed = time.time() - t0
-    step = tr.state.global_step
-    # Did this segment really end on a pushed checkpoint? `stopped_at` is set when the step counter reaches
-    # the target, but performing the save at that step is the library's decision. If it ever ordered the
-    # stop ahead of the save, the session would exit 0 having pushed nothing, and every relaunch would
-    # redo the same interval while `latest` stayed put -- which looks like a stuck run, not a broken one.
-    if hub_cb is not None and args.stop_after_steps:
-        if args.stop_after_steps <= start_step:
-            raise SystemExit(f"--stop-after-steps {args.stop_after_steps} is not beyond the resume point "
-                             f"step {start_step}: this session would train nothing")
-        if step < args.stop_after_steps:
-            raise SystemExit(f"asked to stop at step {args.stop_after_steps}, stopped at {step}")
-        if args.stop_after_steps < steps_planned and args.stop_after_steps not in hub_cb.steps:
-            raise SystemExit(f"segment ended at step {step} but no checkpoint was pushed and verified "
-                             f"for it (pushed: {sorted(hub_cb.steps)}). The next session would resume "
-                             "from the previous pointer and re-read this interval.")
-    samples = samples_at(step, seqs_per_step)
-    consumed = samples * args.seq_len
-    consumed_this_run = max(0, step - start_step) * tokens_per_step
-    hist = [h for h in tr.state.log_history if "loss" in h]
-    final_loss = hist[-1]["loss"] if hist else None
-    say(f"TRAIN DONE step={step}/{steps_planned} elapsed={elapsed/60:.1f} min "
-        f"tokens_consumed={consumed:,} last_loss={final_loss}")
-    # E-006, again: a run that produced NaN still reports a confident tokens/sec. Never let that pass.
-    if final_loss is None or not math.isfinite(final_loss):
-        raise SystemExit(f"final loss is not finite ({final_loss}) -- not exporting anything")
-
-    # Validation first, and on every rank: `evaluate()` is collective, so gating it would deadlock.
-    ppl = None
-    try:
-        ev = tr.evaluate()
-        ppl = math.exp(min(20.0, ev["eval_loss"]))
-        say(f"validation loss={ev['eval_loss']:.4f} ppl={ppl:.2f} over {n_val} held-out windows")
-    except Exception as e:
-        say(f"validation eval failed: {type(e).__name__}: {str(e)[:200]}")
-
-    # Everything below writes files, so it is rank 0's alone (E-029). Two ranks building the same export
-    # directory is a write race: `shutil.copy(src, out)` with `out` not yet created writes a regular FILE
-    # called `final` where rank 0 is making a directory of the same name, and the observed result was
-    # `training_args.bin` vanishing mid-hash -- which killed a fully trained, fully checkpointed leg on
-    # its last line. The variable is also no longer called `out`, which until now meant both the run
-    # directory and the export inside it.
-    exp = os.path.join(args.out, "final")
-    if rank == 0:
-        tr.save_model(exp)
-        # The tokenizer travels with the run so `final/` is loadable by AutoTokenizer on its own (§8). If
-        # it is not in the mix root, say so in the summary instead of skipping quietly and being surprised.
-        shipped = []
-        import shutil as _sh
-        os.makedirs(exp, exist_ok=True)
-        for fn in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json",
-                   "merges.txt"):
-            src = os.path.join(args.root, fn)
-            if os.path.exists(src):
-                _sh.copy(src, exp)
-                shipped.append(fn)
-        say(f"tokenizer files shipped with the export: {shipped or 'NONE -- Phase 5 must add them'}")
-        SD.save_cursor(exp, SD.cursor_dict(args.seq_len, samples, args.data_seed, fingerprint,
-                                           step=step, order_sha=order_sha))
-        json.dump({"step": step, "start_step": start_step, "tokens_consumed": consumed,
-                   "tokens_this_run": consumed_this_run, "seconds": round(elapsed, 1),
-                   "params": pc, "fingerprint": fingerprint, "shuffle_perm_sha": order_sha,
-                   "final_loss": final_loss,
-                   "prior_loss": prior, "samples_consumed": samples,
-                   "steps_planned": steps_planned},
-                  open(os.path.join(exp, "run_summary.json"), "w"), indent=1, sort_keys=True)
-    segment = bool(hub_cb and hub_cb.stopped_at and step < steps_planned)
-    if api is not None and not segment and rank == 0:
-        r = hubckpt.push_and_prune(args.hub_repo, exp, "final", api, repo_type=args.hub_repo_type,
-                                   token=os.environ.get("HF_TOKEN"), prune=args.prune)
-        # .get, not [k]: `pruned` and `free_after_gb` are only present on some paths through
-        # push_and_prune, and a KeyError here would throw away a completed run (E-023).
-        say("final push:", json.dumps({k: r.get(k) for k in
-                                      ("verify", "pruned", "free_before_gb", "free_after_gb",
-                                       "prune_blocked_reason")}, default=str)[:900])
-        # The run is over and `latest.json` has to say so. `on_save` only ever rolls it at a 381-step
-        # boundary, so without this the pointer stops at 3,810, the launcher's `start_step >= horizon`
-        # check never fires, and every relaunch after completion trains four more steps and pushes again.
-        if step >= steps_planned:
-            hubckpt.write_latest_pointer(args.hub_repo, {
-                "path_in_repo": "final", "step": step, "samples_consumed": samples,
-                "tokens_consumed": consumed, "fingerprint": fingerprint, "seed": args.data_seed},
-                api, repo_type=args.hub_repo_type, token=os.environ.get("HF_TOKEN"))
-            say(f"latest.json rolled to the completed run: final at step {step}")
-    elif segment:
-        # The session stopped on a checkpoint interval, and HubPush already pushed AND verified that
-        # checkpoint and rolled latest.json to it. Re-pushing it as "final" would add a 1.7 GB LFS
-        # revision per session for a copy the Hub already holds.
-        say(f"segment boundary reached at step {step} of {steps_planned}: the checkpoint HubPush just "
-            "verified is the resume point, so no separate final push")
-    say("RUN_JSON " + json.dumps({"final_step": step, "val_ppl": ppl, "tokens_consumed": consumed,
-                                 "segment_stop": segment, "steps_planned": steps_planned,
-                                 "tok_per_s": round(consumed_this_run / max(1, elapsed)),
-                                 "final_loss": final_loss, "prior_loss": prior,
-                                 "micro_batch": args.micro_batch, "seq_len": args.seq_len,
-                                 "accum": args.accum, "world": world,
-                                 "params": pc["sum_numel"], "attn": args.attn,
-                                 "grad_ckpt": bool(args.grad_ckpt),
-                                 "optim": args.optim, "gc_nonreentrant": bool(args.gc_nonreentrant),
-                                 "torch_compile": bool(args.torch_compile),
-                                 "effective_attn": eff_attn,
-                                 "peak_gpu_gb": peak_gb()}, sort_keys=True))
-
-
-def peak_gb():
-    """Peak CUDA allocation, in GB. The batch size at seq 2048 is the one thing p0c could not tell us
-    (it measured seq 1024), and an OOM mid-main-run is a lost session, so this number is part of every
-    preflight cell's evidence rather than something to discover at step 300 of the real run."""
-    try:
-        return round(torch.cuda.max_memory_allocated() / 1e9, 2) if torch.cuda.is_available() else None
-    except Exception:
-        return None
-
-
-if __name__ == "__main__":
-    main()
+# The training run in one script: frozen architecture, map-style packed data, fp16 on T4, trapezoid LR,
+# and a checkpoint cycle that satisfies §3.13 by construction.
+
+# Deliberately thin glue over `transformers.Trainer` (docs/01-plan.md §5): data loading, DDP, gradient
+# accumulation, scaler, RNG capture and resumption are the framework's job. What is ours is the shard
+# reader, the LR shape, and the upload/verify/prune discipline -- the three things a generic stack cannot
+# know about this project.
+
+# Two resume facts that are easy to get wrong and are asserted here rather than assumed:
+#   1. `ignore_data_skip=True`. Trainer's own resume would replay discarded batches; we advance the data
+#      position by slicing the permutation in shard_dataset instead. Leaving the flag at False would
+#      double-skip and silently LOSE training data -- the exact §3.1 violation.
+#   2. The cursor is checked against the checkpoint, not trusted. A resumed run whose cursor, seed,
+#      sequence length or corpus fingerprint disagrees with what is on disk stops before step 1.
+
+import argparse
+import json
+import math
+import os
+import time
+
+import torch
+from torch.optim.lr_scheduler import LambdaLR
+from transformers import (LlamaConfig, LlamaForCausalLM, Trainer, TrainerCallback,
+                          TrainingArguments, default_data_collator, set_seed)
+
+import hubckpt
+import shard_dataset as SD
+
+# ---- frozen at Gate 1 (docs/01-plan.md §2.3, §4). Only --layers/--hidden change, for test T1.
+BASE = dict(hidden=576, layers=22, heads=9, kv_heads=3, ffn=1536, vocab=49152,
+            rope_theta=10000.0, rms_eps=1e-5, tie=True)
+
+
+def _kv_dividing(heads, hidden):
+    """GQA requires num_key_value_heads to divide num_attention_heads; SDPA raises otherwise
+    ("Number of heads in key and value must divide the number of heads in Q"). Variant shapes are built
+    from hidden/64, which need not divide 9 -- take the largest divisor of `heads` that fits."""
+    per_head = max(1, hidden // 64)
+    return max(d for d in range(1, min(per_head, heads) + 1) if heads % d == 0)
+
+
+def build_config(seq_len, layers=None, hidden=None, attn="eager"):
+    """The frozen shape, and the two variants test T1 needs (20L at the same width, and a wider/shallower
+    alternative). SwiGLU width is int(2*4h/3) rounded to a 64-multiple -- 1536 at h=576, which is the
+    value the frozen 106,194,240 count was derived from. Getting this factor wrong silently changes the
+    parameter count and makes T1 compare two different models.
+
+    `attn` is explicit because it is a measured choice, not a default: probe p0c found eager attention
+    beating memory-efficient SDPA at seq 1024 on Turing (docs/00-platform-notes.md §9), and Gate 3 P2
+    then clocked 4,071 tok/s with the transformers default -- half the rate p0c saw in a raw loop. Leaving
+    the attention kernel to whatever the library defaults to is how a 68-hour run becomes a 27-hour run.
+    """
+    h = hidden or BASE["hidden"]
+    L = layers or BASE["layers"]
+    # "variant" means a different WIDTH, not merely a different layer count: the frozen GQA ratio (9 Q /
+    # 3 KV) belongs to the frozen width, and passing --hidden 576 explicitly must not quietly turn the
+    # model into MHA-9 and make test T1 measure something else.
+    frozen_width = (h == BASE["hidden"])
+    inter = BASE["ffn"] if frozen_width else (int(round(8 * h / 3)) // 64) * 64
+    kv = BASE["kv_heads"] if frozen_width else _kv_dividing(BASE["heads"], h)
+    cfg = LlamaConfig(vocab_size=BASE["vocab"], hidden_size=h, intermediate_size=inter,
+                      num_hidden_layers=L, num_attention_heads=BASE["heads"],
+                      num_key_value_heads=kv,
+                      max_position_embeddings=seq_len, rope_theta=BASE["rope_theta"],
+                      rms_norm_eps=BASE["rms_eps"], tie_word_embeddings=BASE["tie"],
+                      attention_dropout=0.0, hidden_act="silu", mlp_bias=False,
+                      attn_implementation=attn,
+                      pad_token_id=BASE["vocab"] - 1, eos_token_id=0)
+    return cfg
+
+
+def count_params(model):
+    """Total, embeddings included -- the number §2 budgets against. Reported both ways: numel sum and
+    the unique-parameter count, which differ when embeddings are tied."""
+    total = sum(p.numel() for p in model.parameters())
+    uniq = sum(p.numel() for p in set(model.parameters()))
+    return {"sum_numel": total, "unique": uniq,
+            "trainable": sum(p.numel() for p in model.parameters() if p.requires_grad)}
+
+
+# The cursor arithmetic in one place, because three separate bugs live here:
+#   - a step consumes micro_batch * world * accum sequences, not micro_batch (DDP shards the batch);
+#   - Trainer restores `global_step` from the checkpoint, so on resume it is ALREADY absolute. Adding the
+#     resume point a second time inflates the cursor, and an inflated cursor makes the next run skip data
+#     it never read -- silent data loss with no error anywhere, which is §3.1's worst case. Writing this
+#     as a pure function of the absolute step removes the whole class of bug.
+#   - max_steps is the TOTAL horizon (it is what the LR schedule decays against), so resuming mid-run
+#     must keep the same steps_planned or the trapezoid changes shape.
+def samples_at(absolute_step, seqs_per_step):
+    return int(absolute_step) * seqs_per_step
+
+
+class TrapezoidTrainer(Trainer):
+    """Trainer with (a) a hand-rolled trapezoid schedule and (b) a data position we own."""
+
+    def __init__(self, *a, lr_shape=None, **kw):
+        self.lr_shape = lr_shape
+        super().__init__(*a, **kw)
+
+    def create_scheduler(self, num_training_steps, optimizer=None):
+        warm = self.lr_shape["warmup_steps"]
+        decay_at = self.lr_shape["decay_start"]
+        total = self.lr_shape["total_steps"]
+
+        def fn(cur):
+            if cur < warm:
+                return (cur + 1) / max(1, warm)
+            if cur < decay_at:
+                return 1.0
+            # linear decay to exactly zero at the end: the trapezoid in docs/01-plan.md §4.3
+            return max(0.0, 1.0 - (cur - decay_at) / max(1, total - decay_at))
+
+        if self.lr_scheduler is None:
+            self.lr_scheduler = LambdaLR(self.optimizer if optimizer is None else optimizer, fn)
+        return self.lr_scheduler
+
+
+class HubPush(TrainerCallback):
+    """§3.13 after every checkpoint: write the cursor in, upload, verify from the Hub, prune locally.
+
+    Only rank 0 acts -- the ranks share the checkpoint directory and would race. A failed verification
+    raises inside the callback, so the run stops AT that checkpoint rather than continuing with its only
+    copy on a disk that can be taken away at any moment.
+    """
+
+    def __init__(self, repo, api, token, seq_len, seed, fingerprint, seqs_per_step,
+                 prune=True, min_free_gb=6.0, repo_type="dataset", stop_after_steps=0, order_sha=None):
+        self.repo, self.api, self.token = repo, api, token
+        self.seq_len, self.seed, self.fingerprint = seq_len, seed, fingerprint
+        self.seqs_per_step, self.repo_type = seqs_per_step, repo_type
+        self.order_sha = order_sha
+        self.prune, self.min_free_gb = prune, min_free_gb
+        self.stop_after = int(stop_after_steps or 0)
+        self.stopped_at = None
+        self.n = 0
+        self.steps = []
+        self.results = []
+
+    def on_step_end(self, args, state, control, **kw):
+        # A session that ends ON a verified checkpoint loses nothing when the platform takes the disk.
+        # That is how the 10 % cadence is made compatible with an unknown session cap: instead of pushing
+        # extra intermediate copies (which multiplies Hub LFS storage by the number of sessions), the run
+        # is cut into segments whose length is a whole number of checkpoint intervals.
+        if self.stop_after and state.global_step >= self.stop_after:
+            self.stopped_at = state.global_step
+            control.should_training_stop = True
+        return control
+
+    def on_save(self, args, state, control, **kw):
+        if not state.is_world_process_zero:
+            return control
+        step = state.global_step
+        d = os.path.join(args.output_dir, f"checkpoint-{step}")
+        if not os.path.isdir(d):
+            raise RuntimeError(f"Trainer reported a save but {d} does not exist")
+        samples = samples_at(step, self.seqs_per_step)
+        SD.save_cursor(d, SD.cursor_dict(self.seq_len, samples, self.seed, self.fingerprint,
+                                         step=step, order_sha=self.order_sha))
+        t0 = time.time()
+        # Upload and verify *without* pruning, roll the pointer and read it back, and only then delete the
+        # local copy. The old order (prune inside push_and_prune, pointer written afterwards) meant one
+        # failed pointer commit left a Hub-verified checkpoint that nothing pointed at and no local copy to
+        # recover it with -- and every guard a resume runs would still pass, because weights, cursor,
+        # fingerprint and seed are all self-consistent. Only the pointer would be wrong, and the next
+        # session would quietly re-read 100 M tokens.
+        res = hubckpt.push_and_prune(self.repo, d, f"ckpt/checkpoint-{step}", self.api,
+                                     repo_type=self.repo_type, token=self.token, prune=False)
+        ptr_body = {"path_in_repo": f"ckpt/checkpoint-{step}", "step": step,
+                    "samples_consumed": samples, "tokens_consumed": samples * self.seq_len,
+                    "fingerprint": self.fingerprint, "seed": self.seed}
+        hubckpt.write_latest_pointer(self.repo, ptr_body, self.api, repo_type=self.repo_type,
+                                     token=self.token)
+        if self.prune:
+            hubckpt.prune_local(d, res, self.min_free_gb)
+        loss = next((h["loss"] for h in reversed(state.log_history) if "loss" in h), None)
+        print(f"CKPT {step} samples={samples:,} push+verify={round(time.time()-t0,1)}s "
+              f"ok={res['verify']['ok']} pruned={res.get('pruned')} free={res['free_after_gb']}GB "
+              f"loss={loss}", flush=True)
+        self.n += 1
+        self.steps.append(step)
+        self.results.append(res)
+        return control
+
+
+def run_benchmark(args, model, dataset, steps, world, seqs_per_step):
+    """The same data path, optimizer and fp16 recipe as the real run, with no Trainer in the loop.
+
+    Why this exists: P2 measured 5,105 tok/s for 22L/seq-1024 where probe p0c measured 11,062 tok/s for a
+    near-identical shape with a hand-rolled loop. The gap is either (a) the attention kernel, (b) per-step
+    framework overhead, or (c) a measurement artefact, and which one it is decides whether the main run
+    needs 27 h or 68 h. Guessing wrong in either direction is expensive, so measure both paths.
+    """
+    import torch.distributed as dist
+    from torch.nn.parallel import DistributedDataParallel as DDP
+    from torch.utils.data import DataLoader
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    device = torch.device("cuda", local_rank)
+    torch.cuda.set_device(device)
+    if world > 1 and not dist.is_initialized():
+        dist.init_process_group("nccl")
+    model = model.to(device)
+    if args.grad_ckpt:
+        model.gradient_checkpointing_enable()
+        model.config.use_cache = False
+    net = DDP(model, device_ids=[local_rank]) if world > 1 else model
+
+    opt = torch.optim.AdamW(net.parameters(), lr=args.lr, betas=(0.9, 0.95),
+                            weight_decay=0.1, eps=1e-8)
+    try:
+        scaler = torch.amp.GradScaler("cuda", enabled=True)     # torch >= 2.4
+    except TypeError:
+        scaler = torch.cuda.amp.GradScaler(enabled=True)        # older signature
+    # Same reader, same shuffle, no accumulation: the comparison against the Trainer cells is then
+    # per-step cost with the framework removed, and nothing else differs.
+    dl = DataLoader(dataset, batch_size=args.micro_batch, shuffle=False, num_workers=2,
+                    pin_memory=False, drop_last=True)
+    t0 = time.time()
+    toks = 0
+    for i, batch in enumerate(dl):
+        if i >= steps:
+            break
+        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+        opt.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.float16):
+            out = net(**batch)
+        scaler.scale(out.loss).backward()
+        scaler.unscale_(opt)
+        torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
+        scaler.step(opt)
+        scaler.update()
+        toks += batch["input_ids"].numel()
+        if i == 0:
+            torch.cuda.synchronize()
+            t0 = time.time()   # exclude CUDA context + first-trace cost from the rate
+            toks = 0
+    sec = time.time() - t0
+    res = {"benchmark_steps": i + 1 if steps else 0, "tok_per_s": round(toks / max(1e-6, sec)),
+           "seconds": round(sec, 1), "tokens": toks,
+           "peak_gpu_gb": round(torch.cuda.max_memory_allocated() / 1e9, 2),
+           "world": world, "seq_len": args.seq_len, "micro_batch": args.micro_batch,
+           "attn": args.attn, "grad_ckpt": bool(args.grad_ckpt),
+           "params": sum(p.numel() for p in model.parameters())}
+    print("BENCH_JSON " + json.dumps(res, sort_keys=True), flush=True)
+    if world > 1:
+        dist.barrier()
+        dist.destroy_process_group()
+    return res
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--root", default="/kaggle/working/mixroot")
+    ap.add_argument("--out", default="/kaggle/working/run")
+    ap.add_argument("--hub-repo", default="", help="dataset repo for checkpoints; empty = no upload")
+    ap.add_argument("--hub-repo-type", default="dataset", choices=["dataset", "model"])
+    # D-011: these ARE the frozen main-run values (measured 7,732-7,828 tok/s at 1024/eager on 2xT4), so
+    # that launching the run does not depend on remembering a flag string.
+    ap.add_argument("--seq-len", type=int, default=1024)
+    ap.add_argument("--attn", default="eager", choices=["eager", "sdpa", "flash_attention_2"],
+                    help="p0c measured eager beating SDPA on Turing; Gate 3 P2/A-B measured +52 %% for "
+                         "eager at seq 1024 and an OOM at 2048. flash_attention_2 needs sm80+")
+    ap.add_argument("--grad-ckpt", default=True, action=argparse.BooleanOptionalAction,
+                    help="on by default: it is what keeps peak memory at ~7 GB instead of 12 GB, and the "
+                         "measured rate cost at 1024 was nil (7,828 with GC at m8 vs 7,732 without at m4)")
+    ap.add_argument("--layers", type=int, default=0)
+    ap.add_argument("--hidden", type=int, default=0)
+    ap.add_argument("--micro-batch", type=int, default=4,
+                    help="4 x 1024 x 32 accum x 2 cards = 262,144 tokens/step, the frozen batch size")
+    ap.add_argument("--accum", type=int, default=32)
+    ap.add_argument("--lr", type=float, default=6e-4)
+    ap.add_argument("--optim", default="adamw_torch",
+                    help="only mathematically-equivalent implementations are fair game at freeze time "
+                         "(adamw_torch / adamw_torch_fused / adamw_apex_fused); 8-bit variants change the "
+                         "update, so they are excluded by rule, not by preference")
+    ap.add_argument("--gc-nonreentrant", default=False, action=argparse.BooleanOptionalAction,
+                    help="pass use_reentrant=False to gradient checkpointing: same recomputation, "
+                         "different autograd path, and reports on whether it is faster under DDP conflict")
+    ap.add_argument("--torch-compile", default=False, action=argparse.BooleanOptionalAction,
+                    help="Trainer's torch_compile hook; fused kernels, no change to the update rule")
+    ap.add_argument("--tokens", type=int, default=1_000_000_000)
+    ap.add_argument("--max-steps", type=int, default=0, help="0 = derive from --tokens; >0 overrides")
+    ap.add_argument("--warmup-frac", type=float, default=0.02)
+    ap.add_argument("--decay-frac", type=float, default=0.80)
+    ap.add_argument("--seed", type=int, default=20260919)
+    ap.add_argument("--data-seed", type=int, default=20260919)
+    ap.add_argument("--resume", default="", help="checkpoint dir; empty = auto from Hub latest")
+    ap.add_argument("--val-tokens", type=int, default=2_000_000)
+    ap.add_argument("--log-every", type=int, default=10)
+    ap.add_argument("--push-every-steps", type=int, default=0, help="0 = only at 10%% boundaries")
+    ap.add_argument("--stop-after-steps", type=int, default=0,
+                    help="absolute global step at which THIS session stops training, so it ends on a "
+                         "verified checkpoint instead of being killed between two of them. Must be a "
+                         "whole multiple of the checkpoint interval; 0 = run to the horizon.")
+    ap.add_argument("--prune", action="store_true", help="delete local checkpoint after hub verify")
+    ap.add_argument("--dry-run", action="store_true", help="build everything, train 0 steps")
+    ap.add_argument("--benchmark", type=int, default=0,
+                    help="N steps with no Trainer in the loop (isolates framework per-step cost)")
+    args = ap.parse_args()
+
+    rank = int(os.environ.get("RANK", "0"))
+    world = int(os.environ.get("WORLD_SIZE", "1"))
+    say = (lambda *m: print(*m, flush=True)) if rank == 0 else (lambda *m: None)
+
+    man_path = os.path.join(args.root, "manifest.json")
+    if not os.path.exists(man_path):
+        raise SystemExit(f"no manifest at {man_path}: build the mix first (Phase 2 is not optional)")
+    man = json.load(open(man_path))
+    if man["total_tokens"] < args.tokens:
+        say(f"WARNING mix has {man['total_tokens']:,} tokens < target {args.tokens:,}")
+
+    store = SD.PackedTokenStore(args.root, man)
+    fingerprint = SD.mix_fingerprint(args.root, man)
+    # One definition of the step, used by the cursor, the schedule and the token count: a step consumes
+    # micro_batch * world sequences (DDP shards the batch), so tokens/step follows from it.
+    seqs_per_step = args.micro_batch * world * args.accum
+    tokens_per_step = seqs_per_step * args.seq_len
+    steps_planned = args.max_steps or int(args.tokens / tokens_per_step)
+    say(f"mix: {store.total_tokens:,} tokens in {len(store.shards)} shards fp={fingerprint} "
+        f"| {tokens_per_step:,} tok/step x {steps_planned} steps = "
+        f"{tokens_per_step * steps_planned:,} tokens")
+
+    # ---- resume point: the cursor is the run's identity as much as the weights are.
+    # `--resume auto` is what a fresh instance with an empty disk passes: it asks the Hub where the run
+    # stopped and pulls that checkpoint down. §3.1 forbids recovering from anything left on an instance.
+    if args.resume == "auto":
+        if not args.hub_repo:
+            raise SystemExit("--resume auto needs --hub-repo")
+        import ounce100m_credentials
+        ounce100m_credentials.install()
+        from huggingface_hub import HfApi
+        tk = os.environ.get("HF_TOKEN")
+        ptr = hubckpt.latest_pointer(args.hub_repo, repo_type=args.hub_repo_type, token=tk)
+        if ptr.get("not_found"):
+            # A 404 is normally a fact -- "this run has never checkpointed" -- and step 0 is then right.
+            # But a missing pointer beside existing checkpoints means the pointer upload failed after a
+            # good checkpoint, and starting from zero would silently retrain the corpus from the top,
+            # which §3.1 forbids and E-018 is about. So list the repo before believing the 404.
+            try:
+                listed = sorted(k for k in hubckpt.hub_listing(args.hub_repo, args.hub_repo_type,
+                                                               token=tk)
+                                if k.startswith("ckpt/checkpoint-"))
+            except hubckpt.RepoMissing:
+                # The checkpoint repo has never been created: nothing has ever been pushed, so step 0 is
+                # the only correct answer. The first session of the run always takes this branch, and
+                # treating a 404 as "listing failed" would refuse it (a generic `except Exception` below
+                # makes the list non-empty and the refusal fires).
+                say(f"{args.hub_repo} does not exist yet -- nothing has ever been pushed")
+                listed = []
+            except Exception as e:
+                listed = [f"<listing failed: {type(e).__name__}>"]
+            if listed:
+                raise SystemExit(f"--resume auto found no latest.json on {args.hub_repo}, but the repo "
+                                 f"holds {len(listed)} checkpoint path(s) ({listed[:3]}). The pointer is "
+                                 "missing, not the run -- refusing to start from step 0.")
+            say(f"no latest.json and no checkpoints on {args.hub_repo} -> starting from step 0")
+            args.resume = ""
+        elif "error" in ptr:
+            # Anything else (401, timeout, DNS) means the lookup failed, not that there is nothing there.
+            # Falling through to step 0 here would silently retrain the first stretch of the run, which is
+            # exactly the §3.1 violation this whole branch exists to prevent (E-018's shape).
+            raise SystemExit(f"--resume auto could not read the hub pointer: {ptr['error']}")
+        else:
+            # A pointer is a resume point only if nothing newer was ever pushed. This scan catches the one
+            # case write_latest_pointer's own read-back cannot see: a pointer left behind by a rollback, a
+            # deletion, or a session that died between its checkpoint push and its pointer roll. Resuming
+            # from the stale pointer would re-train an interval with weights and a cursor that agree with
+            # each other and disagree with the Hub.
+            import re as _re
+            steps_on_hub = set()
+            for attempt in (1, 2):
+                try:
+                    steps_on_hub = {int(m.group(1)) for k in
+                                    hubckpt.hub_listing(args.hub_repo, args.hub_repo_type, token=tk)
+                                    if (m := _re.match(r"ckpt/checkpoint-(\d+)/", k))}
+                    break
+                except Exception as e:
+                    say(f"   checkpoint scan attempt {attempt}: {type(e).__name__}: {str(e)[:140]}")
+                    time.sleep(15)
+            else:
+                raise SystemExit("--resume auto could not enumerate the checkpoint repo, so it cannot "
+                                 f"prove latest.json (step {ptr.get('step')}) is the newest state")
+            if steps_on_hub and max(steps_on_hub) > int(ptr.get("step") or 0):
+                raise SystemExit(f"latest.json points at step {ptr.get('step')} but {args.hub_repo} also "
+                                 f"holds checkpoint-{max(steps_on_hub)}: the pointer is behind the run. "
+                                 "Resuming here would re-train that interval, so refusing.")
+            # Both ranks need the checkpoint on disk, and this block runs before TrainingArguments has
+            # created a process group -- so there is no barrier to synchronise on. Downloading into a
+            # per-rank directory instead removes the race entirely (two writers, one path), at the cost of
+            # a second 1.7 GB pull that happens in parallel. Each rank then reads its own verified copy and
+            # derives the same start_sample/start_step from the same cursor.json bytes.
+            dest = os.path.join(args.out, f"resume_r{rank}",
+                                os.path.basename(ptr["path_in_repo"]))
+            say(f"auto-resume: hub says step {ptr.get('step')} at {ptr['path_in_repo']}")
+            got = hubckpt.download_checkpoint(args.hub_repo, ptr["path_in_repo"], dest, HfApi(),
+                                              repo_type=args.hub_repo_type, token=tk)
+            say(f"downloaded {got['files']} files, {got['bytes']/1e6:.0f} MB in {got['seconds']}s")
+            if got["bytes"] < 1_000_000:
+                raise SystemExit("auto-resume downloaded a checkpoint that is not real weights")
+            args.resume = got["dir"]
+            say(f"resuming from {args.resume}")
+    cur = SD.load_cursor(args.resume) if args.resume else None
+    if args.resume and cur is None:
+        # Weights at step N with the data stream back at sample 0 would retrain the opening of the corpus
+        # and report a continuous loss curve. Refuse instead.
+        raise SystemExit(f"resume refused: {args.resume} has no cursor.json -- a checkpoint without a "
+                         "recorded data position cannot be resumed, only restarted (§3.1)")
+    if args.resume:
+        say(f"resume from {args.resume} cursor={cur}")
+    start_sample = (cur or {}).get("samples_consumed", 0)
+    start_step = (cur or {}).get("step") or 0
+    # Identity of the visit order, derived from the corpus size and the seed rather than trusted to the
+    # seed: numpy's permutation is only stable within a version, and Kaggle updates its image underneath
+    # long runs. ~50 ms once per session, and it is what makes "resumed onto a different order" loud.
+    order_sha = SD.perm_sha(store.total_tokens // args.seq_len, args.data_seed)
+    if cur:
+        for k, v in (("seq_len", args.seq_len), ("shuffle_seed", args.data_seed),
+                     ("shuffle_perm_sha", order_sha), ("dataset_files_sha", fingerprint)):
+            if cur.get(k) != v:
+                raise SystemExit(f"resume refused: cursor {k}={cur.get(k)} != current {v}. "
+                                 "A resumed run must read the same data in the same order (§3.1).")
+        if start_sample % seqs_per_step:
+            raise SystemExit(f"resume refused: cursor samples_consumed={start_sample} is not a whole "
+                             f"number of steps (seqs/step={seqs_per_step}) -- batch geometry changed")
+        if cur.get("step") is not None and cur["step"] * seqs_per_step != start_sample:
+            raise SystemExit(f"resume refused: cursor step={cur['step']} and samples_consumed="
+                             f"{start_sample} disagree at seqs/step={seqs_per_step}")
+        if start_step * seqs_per_step != start_sample:
+            start_step = start_sample // seqs_per_step
+
+    train_ds = SD.make_dataset(store, args.seq_len, args.data_seed, start_sample=start_sample)
+    assert getattr(train_ds, "order_sha", None) == order_sha, "two derivations of the order disagree"
+    # Validation comes from the mix's held-out val/ shards, never from the training stream: scoring on
+    # tokens the model already consumed would be a number nobody could reproduce (and §3.12).
+    vs = [s["file"] for s in man.get("val_shards", [])]
+    if not vs:
+        raise SystemExit("manifest lists no val_shards -- Gate 2 requires a held-out split")
+    vstore = SD.PackedTokenStore(args.root, man, files=vs)
+    n_val = min(max(1, args.val_tokens // args.seq_len), vstore.total_tokens // args.seq_len)
+    val_ds = SD.make_dataset(vstore, args.seq_len, args.data_seed, shuffle=False, count=n_val)
+    need = (steps_planned - start_step) * seqs_per_step
+    if len(train_ds) < need:
+        # Reading past the end of the stream makes the Trainer begin epoch 2 at the top of the same
+        # permutation: the run would then report the full token target while quietly re-training on the
+        # first windows. The mix is built with ~12 % headroom over the target for exactly this reason, so
+        # arriving here means that assumption stopped being true (review finding), and the honest response
+        # is to stop before step 1 rather than to discover it in the final report.
+        raise SystemExit(
+            f"dataset too small for the plan: {len(train_ds):,} windows available from sample "
+            f"{start_sample:,}, {need:,} needed for {steps_planned} steps x {seqs_per_step} sequences "
+            f"({need * args.seq_len:,} tokens). Lower --tokens to the 0.9 B fallback (D-012) or rebuild "
+            "the mix larger -- never let the epoch wrap.")
+    say(f"data: {len(train_ds):,} training windows available, {need:,} needed "
+        f"({100.0 * len(train_ds) / need - 100:.1f} % margin), {n_val:,} held-out val windows")
+
+    cfg = build_config(args.seq_len, args.layers or None, args.hidden or None, attn=args.attn)
+    set_seed(args.seed)  # before init: both ranks must start from identical weights, not two draws
+    model = LlamaForCausalLM(cfg).to(dtype=torch.float32)
+    pc = count_params(model)
+    say("params:", json.dumps(pc))
+    if not (90_000_000 <= pc["sum_numel"] <= 110_000_000) and not args.layers and not args.hidden:
+        raise SystemExit(f"parameter count {pc['sum_numel']:,} is outside §2's 90-110M budget")
+
+    if args.benchmark:
+        # Both ranks iterate the same samples here (no DistributedSampler): this path measures step cost,
+        # not a training run, and it must never be used to produce a checkpoint.
+        say(f"benchmark mode: {args.benchmark} steps, no Trainer, micro={args.micro_batch} "
+            f"seq={args.seq_len} attn={args.attn} world={world}")
+        run_benchmark(args, model, train_ds, args.benchmark, world, seqs_per_step)
+        return
+
+    # 10 checkpoints per run is §2's cadence; --push-every-steps overrides for the preflight.
+    ckpt_every = args.push_every_steps or max(1, steps_planned // 10)
+    if args.stop_after_steps:
+        if args.stop_after_steps % ckpt_every:
+            raise SystemExit(f"--stop-after-steps {args.stop_after_steps} is not a whole multiple of the "
+                             f"checkpoint interval {ckpt_every}: the session would end between two "
+                             f"checkpoints, which is exactly what this flag exists to prevent")
+        if args.stop_after_steps > steps_planned:
+            raise SystemExit(f"--stop-after-steps {args.stop_after_steps} exceeds the horizon "
+                             f"{steps_planned}; use 0 to run to the end")
+
+    api = None
+    if args.hub_repo:
+        import ounce100m_credentials
+        ounce100m_credentials.install()
+        from huggingface_hub import HfApi
+        api = HfApi()
+        # rank 0 only: two ranks committing the same .gitattributes to a brand-new repo is a write race
+        # on the one path that must not fail (only rank 0 ever uploads checkpoints, so rank 1 needs none)
+        if rank == 0:
+            say("ckpt repo:", json.dumps(hubckpt.ensure_repo(args.hub_repo, api,
+                                                             os.environ.get("HF_TOKEN"),
+                                                             args.hub_repo_type)))
+        say(f"checkpoint hub target {args.hub_repo} (prune={args.prune}, every {ckpt_every} steps)")
+
+    tacc = TrainingArguments(
+        output_dir=args.out,
+        per_device_train_batch_size=args.micro_batch,
+        # Same geometry for the validation pass: with the library default (8) the end-of-run eval measures
+        # a peak that has nothing to do with the training configuration, which is how an OOM in the no-
+        # checkpointing cell could be mistaken for a passing one (review finding).
+        per_device_eval_batch_size=args.micro_batch,
+        gradient_accumulation_steps=args.accum,
+        learning_rate=args.lr,
+        weight_decay=0.1,
+        adam_beta1=0.9, adam_beta2=0.95, adam_epsilon=1e-8,
+        max_grad_norm=1.0,
+        lr_scheduler_type="constant",   # replaced by TrapezoidTrainer.create_scheduler
+        warmup_ratio=0.0,               # the trapezoid owns warmup
+        num_train_epochs=1,
+        max_steps=0 if args.dry_run else steps_planned,
+        fp16=True, bf16=False,
+        # measured decision, not a default: A-B found eager attention 52 % faster at seq 1024 and OOM at
+        # 2048; checkpointing is what makes eager at 2048 possible at all, at the cost of a recompute pass
+        gradient_checkpointing=args.grad_ckpt,
+        # Same recomputation, different autograd path -- offered as a cells-vs-baseline question, not a
+        # belief. Reports on use_reentrant=False under DDP contradict each other across versions.
+        gradient_checkpointing_kwargs={"use_reentrant": not args.gc_nonreentrant}
+        if args.grad_ckpt else None,
+        optim=args.optim,
+        torch_compile=args.torch_compile,
+        ddp_find_unused_parameters=False,
+        # Rank 0 spends minutes inside on_save uploading ~1.7 GB, listing the repo and downloading one file
+        # back to hash it, while rank 1 waits at the next collective. The library's 1,800 s default would
+        # abort the run on a slow-egress Kaggle box -- precisely during the checkpoint that the other rank
+        # was protecting. Four hours makes the collective patient without hiding a real deadlock.
+        ddp_timeout=14400,
+        dataloader_num_workers=2,
+        dataloader_pin_memory=False,
+        remove_unused_columns=False,    # our dict keys must reach the model untouched
+        ignore_data_skip=True,          # see header note 1 -- we advance position ourselves
+        save_strategy="no" if args.dry_run else "steps",
+        save_steps=ckpt_every,
+        save_total_limit=1,             # §3.13: the Hub is the store, the instance is scratch
+        logging_steps=args.log_every, report_to=[],
+        seed=args.seed, data_seed=args.seed,
+        accelerator_config={"dispatch_batches": False},
+        average_tokens_across_devices=False,
+    )
+
+    tr = TrapezoidTrainer(
+        model=model, args=tacc, train_dataset=train_ds, eval_dataset=val_ds,
+        data_collator=default_data_collator, lr_shape={
+            "warmup_steps": max(100, int(steps_planned * args.warmup_frac)),
+            "decay_start": int(steps_planned * args.decay_frac),
+            "total_steps": steps_planned},
+    )
+    hub_cb = None
+    if api is not None:
+        hub_cb = HubPush(args.hub_repo, api, os.environ.get("HF_TOKEN"), args.seq_len,
+                         args.data_seed, fingerprint, seqs_per_step,
+                         prune=args.prune, repo_type=args.hub_repo_type,
+                         stop_after_steps=args.stop_after_steps, order_sha=order_sha)
+        tr.add_callback(hub_cb)
+
+    # Precision must be ASSERTED, not configured: bf16 "works" silently on Turing and fp32 master
+    # weights can be lost in a refactor (docs/00-platform-notes.md §8, memory/ERRORS.md E-006). The
+    # scaler's home moved between transformers versions, so look in both places instead of trusting one
+    # attribute path (E-010: prefer version-free surfaces).
+    scaler = getattr(tr, "scaler", None) or getattr(getattr(tr, "accelerator", None), "scaler", None)
+    eff_attn = (getattr(model.config, "_attn_implementation", None)
+                or getattr(model.config, "attn_implementation", None))
+    say(f"precision: fp16={tacc.fp16} bf16={tacc.bf16} scaler={type(scaler).__name__} "
+        f"master_dtype={next(model.parameters()).dtype} requested_attn={args.attn} "
+        f"effective_attn={eff_attn}")
+    assert tacc.fp16 and not tacc.bf16, "not the frozen T4 fp16 recipe"
+    assert next(model.parameters()).dtype == torch.float32, "master weights are not fp32"
+
+    prior = None
+    if args.resume:
+        tsp = os.path.join(args.resume, "trainer_state.json")
+        if os.path.exists(tsp):
+            h0 = [h for h in json.load(open(tsp)).get("log_history", []) if "loss" in h]
+            prior = h0[-1]["loss"] if h0 else None
+        say(f"prior loss at resume: {prior} -- loss continuity is judged against this number")
+
+    if args.dry_run:
+        say(json.dumps({"dry_run": True, "train_samples": len(train_ds), "val_samples": len(val_ds),
+                        "steps_planned": steps_planned, "ckpt_every": ckpt_every,
+                        "params": pc, "fingerprint": fingerprint, "scaler": type(scaler).__name__}))
+        return
+
+    t0 = time.time()
+    tr.train(resume_from_checkpoint=args.resume or None)
+
+    elapsed = time.time() - t0
+    step = tr.state.global_step
+    # Did this segment really end on a pushed checkpoint? `stopped_at` is set when the step counter reaches
+    # the target, but performing the save at that step is the library's decision. If it ever ordered the
+    # stop ahead of the save, the session would exit 0 having pushed nothing, and every relaunch would
+    # redo the same interval while `latest` stayed put -- which looks like a stuck run, not a broken one.
+    if hub_cb is not None and args.stop_after_steps:
+        if args.stop_after_steps <= start_step:
+            raise SystemExit(f"--stop-after-steps {args.stop_after_steps} is not beyond the resume point "
+                             f"step {start_step}: this session would train nothing")
+        if step < args.stop_after_steps:
+            raise SystemExit(f"asked to stop at step {args.stop_after_steps}, stopped at {step}")
+        if args.stop_after_steps < steps_planned and args.stop_after_steps not in hub_cb.steps:
+            raise SystemExit(f"segment ended at step {step} but no checkpoint was pushed and verified "
+                             f"for it (pushed: {sorted(hub_cb.steps)}). The next session would resume "
+                             "from the previous pointer and re-read this interval.")
+    samples = samples_at(step, seqs_per_step)
+    consumed = samples * args.seq_len
+    consumed_this_run = max(0, step - start_step) * tokens_per_step
+    hist = [h for h in tr.state.log_history if "loss" in h]
+    final_loss = hist[-1]["loss"] if hist else None
+    say(f"TRAIN DONE step={step}/{steps_planned} elapsed={elapsed/60:.1f} min "
+        f"tokens_consumed={consumed:,} last_loss={final_loss}")
+    # E-006, again: a run that produced NaN still reports a confident tokens/sec. Never let that pass.
+    if final_loss is None or not math.isfinite(final_loss):
+        raise SystemExit(f"final loss is not finite ({final_loss}) -- not exporting anything")
+
+    # Validation first, and on every rank: `evaluate()` is collective, so gating it would deadlock.
+    ppl = None
+    val_err = None
+    try:
+        ev = tr.evaluate()
+        ppl = math.exp(min(20.0, ev["eval_loss"]))
+        say(f"validation loss={ev['eval_loss']:.4f} ppl={ppl:.2f} over {n_val} held-out windows")
+    except Exception as e:
+        # `say()` is a no-op on rank 1, and that silence is what made a rank-1 evaluation crash unreadable
+        # in probe 1 v1: the handler ran on both ranks, only rank 0 spoke, and the divergence looked like an
+        # unexplained exit 1 with `error_file: <N/A>`. Report on every rank, and carry the failure into the
+        # record so a missing PPL can never be read as a passing validation pass.
+        val_err = f"{type(e).__name__}: {str(e)[:300]}"
+        print(f"[rank {rank}] validation eval failed: {val_err}", flush=True)
+
+    # Sampled here, on every rank, before the rank-0-only export block: a collective reduction inside
+    # `if rank == 0:` would hang, and reading rank 0's allocator alone is what made the probe's memory
+    # verdict systematically optimistic -- it hides the other card's peak, the two CUDA contexts, and the
+    # reserved-but-unallocated blocks that are the real ceiling on a 14.56 GiB T4 (review finding B4).
+    mem = peak_stats()
+    say("memory:", json.dumps(mem, sort_keys=True))
+
+    # Everything below writes files, so it is rank 0's alone (E-029). Two ranks building the same export
+    # directory is a write race: `shutil.copy(src, out)` with `out` not yet created writes a regular FILE
+    # called `final` where rank 0 is making a directory of the same name, and the observed result was
+    # `training_args.bin` vanishing mid-hash -- which killed a fully trained, fully checkpointed leg on
+    # its last line. The variable is also no longer called `out`, which until now meant both the run
+    # directory and the export inside it.
+    exp = os.path.join(args.out, "final")
+    if rank == 0:
+        tr.save_model(exp)
+        # The tokenizer travels with the run so `final/` is loadable by AutoTokenizer on its own (§8). If
+        # it is not in the mix root, say so in the summary instead of skipping quietly and being surprised.
+        shipped = []
+        import shutil as _sh
+        os.makedirs(exp, exist_ok=True)
+        for fn in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json",
+                   "merges.txt"):
+            src = os.path.join(args.root, fn)
+            if os.path.exists(src):
+                _sh.copy(src, exp)
+                shipped.append(fn)
+        say(f"tokenizer files shipped with the export: {shipped or 'NONE -- Phase 5 must add them'}")
+        SD.save_cursor(exp, SD.cursor_dict(args.seq_len, samples, args.data_seed, fingerprint,
+                                           step=step, order_sha=order_sha))
+        json.dump({"step": step, "start_step": start_step, "tokens_consumed": consumed,
+                   "tokens_this_run": consumed_this_run, "session_seconds": round(elapsed, 1),
+                   "params": pc, "params_total": pc["sum_numel"], "fingerprint": fingerprint,
+                   "shuffle_perm_sha": order_sha,
+                   "final_loss": final_loss, "prior_loss": prior, "samples_consumed": samples,
+                   "steps_planned": steps_planned,
+                   # The recipe, written down. The Phase 5 card reads this file to state hyperparameters,
+                   # and it had none of them: every claim about sequence length, batch shape, throughput or
+                   # attention would have rendered as `None` because those lived only in the stdout RUN_JSON
+                   # line, a different schema (review finding A2). `seconds` was also renamed: it is THIS
+                   # session's elapsed, and a five-session run published as one number is wrong by ~5x (A3).
+                   "seq_len": args.seq_len, "micro_batch": args.micro_batch, "accum": args.accum,
+                   "world": world, "tokens_per_step": tokens_per_step,
+                   "tok_per_s": round(consumed_this_run / max(1, elapsed)),
+                   "attn": args.attn, "effective_attn": eff_attn, "grad_ckpt": bool(args.grad_ckpt),
+                   "optim": args.optim, "lr": args.lr, "warmup_frac": args.warmup_frac,
+                   "decay_frac": args.decay_frac, "seed": args.seed, "data_seed": args.data_seed,
+                   "val_ppl": ppl, "val_loss_error": val_err, "val_windows": n_val,
+                   "peak_gpu_gb": mem["rank0_allocated_gb"],
+                   "peak_reserved_gb": mem["rank0_reserved_gb"],
+                   "peak_gpu_gb_max_rank": mem["max_allocated_gb"],
+                   "peak_reserved_gb_max_rank": mem["max_reserved_gb"]},
+                  open(os.path.join(exp, "run_summary.json"), "w"), indent=1, sort_keys=True)
+    segment = bool(hub_cb and hub_cb.stopped_at and step < steps_planned)
+    if api is not None and not segment and rank == 0:
+        r = hubckpt.push_and_prune(args.hub_repo, exp, "final", api, repo_type=args.hub_repo_type,
+                                   token=os.environ.get("HF_TOKEN"), prune=args.prune)
+        # .get, not [k]: `pruned` and `free_after_gb` are only present on some paths through
+        # push_and_prune, and a KeyError here would throw away a completed run (E-023).
+        say("final push:", json.dumps({k: r.get(k) for k in
+                                      ("verify", "pruned", "free_before_gb", "free_after_gb",
+                                       "prune_blocked_reason")}, default=str)[:900])
+        # The run is over and `latest.json` has to say so. `on_save` only ever rolls it at a 381-step
+        # boundary, so without this the pointer stops at 3,810, the launcher's `start_step >= horizon`
+        # check never fires, and every relaunch after completion trains four more steps and pushes again.
+        if step >= steps_planned:
+            hubckpt.write_latest_pointer(args.hub_repo, {
+                "path_in_repo": "final", "step": step, "samples_consumed": samples,
+                "tokens_consumed": consumed, "fingerprint": fingerprint, "seed": args.data_seed},
+                api, repo_type=args.hub_repo_type, token=os.environ.get("HF_TOKEN"))
+            say(f"latest.json rolled to the completed run: final at step {step}")
+    elif segment:
+        # The session stopped on a checkpoint interval, and HubPush already pushed AND verified that
+        # checkpoint and rolled latest.json to it. Re-pushing it as "final" would add a 1.7 GB LFS
+        # revision per session for a copy the Hub already holds.
+        say(f"segment boundary reached at step {step} of {steps_planned}: the checkpoint HubPush just "
+            "verified is the resume point, so no separate final push")
+    say("RUN_JSON " + json.dumps({"final_step": step, "val_ppl": ppl, "tokens_consumed": consumed,
+                                 "segment_stop": segment, "steps_planned": steps_planned,
+                                 "tok_per_s": round(consumed_this_run / max(1, elapsed)),
+                                 "final_loss": final_loss, "prior_loss": prior,
+                                 "micro_batch": args.micro_batch, "seq_len": args.seq_len,
+                                 "accum": args.accum, "world": world,
+                                 "params": pc["sum_numel"], "attn": args.attn,
+                                 "grad_ckpt": bool(args.grad_ckpt),
+                                 "optim": args.optim, "gc_nonreentrant": bool(args.gc_nonreentrant),
+                                 "torch_compile": bool(args.torch_compile),
+                                 "effective_attn": eff_attn,
+                                 "peak_gpu_gb": mem["rank0_allocated_gb"],
+                                 "peak_reserved_gb": mem["rank0_reserved_gb"],
+                                 "peak_gpu_gb_max_rank": mem["max_allocated_gb"],
+                                 "peak_reserved_gb_max_rank": mem["max_reserved_gb"],
+                                 # So a caller cannot report success for a validation pass that never
+                                 # happened: a null val_ppl with no error here would be a bug, and a null
+                                 # with an error is a measured failure.
+                                 "val_error": val_err, "val_windows": n_val}, sort_keys=True))
+
+
+def peak_stats():
+    """Peak CUDA memory, from every rank, with reserved bytes and the cross-rank maximum.
+
+    An OOM mid-run is a lost session, so this is part of every preflight cell's evidence rather than
+    something to discover at step 300. `max_memory_allocated` on rank 0 alone understates what the card is
+    actually holding: the allocator reserves more than it hands out, and the other rank's peak is a separate
+    number on a separate GPU that the run's own limit has to respect.
+    """
+    out = {"cuda": torch.cuda.is_available(), "rank0_allocated_gb": None, "rank0_reserved_gb": None,
+           "max_allocated_gb": None, "max_reserved_gb": None, "device": None}
+    if not out["cuda"]:
+        return out
+    try:
+        alloc = torch.cuda.max_memory_allocated()
+        res = torch.cuda.max_memory_reserved()
+        out["rank0_allocated_gb"] = round(alloc / 1e9, 2)
+        out["rank0_reserved_gb"] = round(res / 1e9, 2)
+        out["device"] = torch.cuda.get_device_name(0)
+        t = torch.tensor([float(alloc), float(res)], device=alloc and "cuda" or "cpu")
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.MAX)
+            out["max_allocated_gb"] = round(float(t[0]) / 1e9, 2)
+            out["max_reserved_gb"] = round(float(t[1]) / 1e9, 2)
+        else:
+            out["max_allocated_gb"] = out["rank0_allocated_gb"]
+            out["max_reserved_gb"] = out["rank0_reserved_gb"]
+    except Exception as e:
+        out["error"] = f"{type(e).__name__}: {str(e)[:120]}"
+    return out
+
+
+if __name__ == "__main__":
+    main()