train/train_ounce100m.py

# 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)
        # E-042: `prune=False` above means `push_and_prune`'s own guard never runs, so this path had to
        # carry it. An unverified checkpoint must stop the session *before* the pointer moves: a pointer
        # rolled onto bytes nobody checked, with the local copy then deleted, is a run that cannot be
        # resumed and does not know it.
        ver = res["verify"]
        if not ver["ok"]:
            raise RuntimeError(
                f"checkpoint {step} did not verify on the Hub "
                f"(missing={len(ver['missing'])} mismatched={len(ver['mismatch'])} "
                f"extra={len(ver['extra'])} readback_ok={ver['readback_ok']}). Keeping the local copy in "
                f"{d} and refusing to roll latest.json -- the next session would otherwise resume from "
                "bytes that were never confirmed to exist.")
        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_verified(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 and rank == 0:
        # Rank 0 only, because every fact below is a rank-0 fact: `on_save` returns early on the other
        # ranks, so `hub_cb.steps` is empty on rank 1 by construction and this guard -- written to catch
        # a segment that ended without a pushed checkpoint -- killed *every* correctly checkpointed
        # segment on rank 1 with `pushed: []`. That SystemExit, not a divergent collective, is what
        # probe v1/v4/v5/v6/v7 saw as "rank 1 exitcode 1 with error_file <N/A>": rank 1 died on its own
        # assertion while rank 0 sat in the next all_reduce until torchrun SIGTERMed it (E-044).
        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
    # Every rank, with `print` and not `say`: probe v6 died on the leg that stopped early and the only
    # evidence torchrun passed up was "rank 1 exitcode 1, error_file <N/A>", because the lines that would
    # have localised it are rank-0-only. These four marks say which statement in the post-training
    # sequence a rank never reached.
    print(f"[rank {rank}] post-train: step={step}/{steps_planned} log_history_loss_entries={len(hist)} "
          f"final_loss={final_loss} start_step={start_step}", flush=True)
    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")

    # A segment that stopped short of the horizon does not run the validation pass, on every rank.
    #
    # Written for what probes v1/v4 saw ("rank 1 exitcode 1, error_file <N/A>, right after a forced stop"),
    # which was attributed to `evaluate()` being entered from a diverged collective state. E-044 later found
    # the actual cause -- the rank-0-only push assertion firing on rank 1 -- so this skip is not what makes
    # the forced-stop leg work. It stays, on its own merits: the held-out PPL of a 762-step segment is not
    # worth a forward pass over 2 M tokens inside a billed session, and §5's per-wake monitoring is served by
    # the training loss. The report's PPL comes from the horizon leg, which is where it is measured.
    forced_stop = bool(hub_cb and hub_cb.stopped_at and step < steps_planned)
    print(f"[rank {rank}] past the stop guard: forced_stop={forced_stop} "
          f"stopped_at={getattr(hub_cb, 'stopped_at', None)}", flush=True)
    ppl = None
    val_err = None
    if forced_stop:
        say(f"validation skipped: segment ended at step {step} of {steps_planned} (evaluate() after a "
            "forced stop diverges between ranks -- probe 1 v1/v4)")
    else:
        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).
    print(f"[rank {rank}] peak_stats: entering (this is a collective)", flush=True)
    mem = peak_stats()
    print(f"[rank {rank}] peak_stats: done, max_reserved={mem['max_reserved_gb']} GiB", flush=True)
    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_skipped": forced_stop,
                   "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)
    print(f"[rank {rank}] past the export and push block", flush=True)
    segment = forced_stop        # the same predicate, computed before the validation decision above
    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_skipped": forced_stop, "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()