train: add train_ounce100m (Phase 3/4 glue)

train/train_ounce100m.py

@@ -0,0 +1,351 @@
+# 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 build_config(seq_len, layers=None, hidden=None):
+    h = hidden or BASE["hidden"]
+    L = layers or BASE["layers"]
+    inter = BASE["ffn"] if hidden is None and layers is None else int(round(8 / 3 * h)) * 2 // 64 * 64
+    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=BASE["kv_heads"] if hidden is None else max(1, h // 64),
+                      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,
+                      pad_token_id=BASE["vocab"] - 1, eos_token_id=0)
+    # GQA ratio is part of the frozen shape; 9 heads with 3 kv heads, never MHA at this budget.
+    if hidden is None:
+        cfg.num_key_value_heads = BASE["kv_heads"]
+    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 sequences, not micro_batch (DDP shards the batch);
+#   - on resume, Trainer restores global_step from the checkpoint but our dataset was ALREADY sliced to
+#     start there, so the total is start_sample + step_within_this_run * seqs_per_step;
+#   - 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(step_in_run, start_sample, seqs_per_step):
+    return start_sample + int(step_in_run) * 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, start_sample, seqs_per_step,
+                 prune=True, min_free_gb=6.0):
+        self.repo, self.api, self.token = repo, api, token
+        self.seq_len, self.seed, self.fingerprint = seq_len, seed, fingerprint
+        self.start_sample, self.seqs_per_step = start_sample, seqs_per_step
+        self.prune, self.min_free_gb = prune, min_free_gb
+        self.n = 0
+        self.results = []
+
+    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.start_sample, self.seqs_per_step)
+        SD.save_cursor(d, SD.cursor_dict(self.seq_len, samples, self.seed, self.fingerprint))
+        t0 = time.time()
+        res = hubckpt.push_and_prune(self.repo, d, f"ckpt/checkpoint-{step}", self.api,
+                                     token=self.token, prune=self.prune,
+                                     min_free_gb=self.min_free_gb)
+        hubckpt.write_latest_pointer(self.repo, {
+            "path_in_repo": f"ckpt/checkpoint-{step}", "step": step, "samples_consumed": samples,
+            "tokens_consumed": samples * self.seq_len, "fingerprint": self.fingerprint,
+            "seed": self.seed}, self.api, token=self.token)
+        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.results.append(res)
+        return control
+
+
+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"])
+    ap.add_argument("--seq-len", type=int, default=2048)
+    ap.add_argument("--layers", type=int, default=0)
+    ap.add_argument("--hidden", type=int, default=0)
+    ap.add_argument("--micro-batch", type=int, default=2)
+    ap.add_argument("--accum", type=int, default=32)
+    ap.add_argument("--lr", type=float, default=6e-4)
+    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("--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")
+    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)
+    tokens_per_step = args.micro_batch * args.seq_len * args.accum * world
+    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
+        ptr = hubckpt.latest_pointer(args.hub_repo, repo_type=args.hub_repo_type)
+        if "error" in ptr:
+            say(f"no latest pointer on {args.hub_repo} ({ptr['error']}) -> starting from step 0")
+            args.resume = ""
+        else:
+            dest = os.path.join(args.out, 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)
+            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 = dest
+    cur = SD.load_cursor(args.resume) if args.resume else None
+    if args.resume:
+        say(f"resume from {args.resume} cursor={cur}")
+    start_sample = (cur or {}).get("samples_consumed", 0)
+    if cur:
+        for k, v in (("seq_len", args.seq_len), ("shuffle_seed", args.data_seed),
+                     ("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).")
+
+    train_ds = SD.make_dataset(store, args.seq_len, args.data_seed, start_sample=start_sample)
+    # 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)
+
+    cfg = build_config(args.seq_len, args.layers or None, args.hidden or None)
+    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")
+
+    seqs_per_step = args.micro_batch * world * args.accum
+    if seqs_per_step * args.seq_len != tokens_per_step:
+        raise SystemExit("internal: batch arithmetic disagrees, token/step accounting is unreliable")
+    ckpt_every = args.push_every_steps or max(1, steps_planned // 10)
+
+    api = None
+    if args.hub_repo:
+        import ounce100m_credentials
+        ounce100m_credentials.install()
+        from huggingface_hub import HfApi
+        api = HfApi()
+        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,
+        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,
+        gradient_checkpointing=False,   # measure before needing it; docs/01-plan.md §6
+        ddp_find_unused_parameters=False,
+        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},
+    )
+    if api is not None:
+        tr.add_callback(HubPush(args.hub_repo, api, os.environ.get("HF_TOKEN"), args.seq_len,
+                                args.data_seed, fingerprint, start_sample, seqs_per_step,
+                                prune=args.prune))
+
+    # 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)
+    say(f"precision: fp16={tacc.fp16} bf16={tacc.bf16} scaler={type(scaler).__name__} "
+        f"master_dtype={next(model.parameters()).dtype}")
+    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
+    samples = samples_at(step, start_sample, seqs_per_step)
+    consumed = samples * args.seq_len
+    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")
+
+    out = os.path.join(args.out, "final")
+    tr.save_model(out)
+    SD.save_cursor(out, SD.cursor_dict(args.seq_len, samples, args.data_seed, fingerprint))
+    json.dump({"step": step, "tokens_consumed": consumed, "seconds": round(elapsed, 1),
+               "params": pc, "fingerprint": fingerprint, "final_loss": final_loss,
+               "prior_loss": prior, "samples_consumed": samples, "steps_planned": steps_planned},
+              open(os.path.join(out, "run_summary.json"), "w"), indent=1, sort_keys=True)
+    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]}")
+    if api is not None:
+        r = hubckpt.push_and_prune(args.hub_repo, out, "final", api,
+                                   token=os.environ.get("HF_TOKEN"), prune=args.prune)
+        say("final push:", json.dumps({k: r[k] for k in ("verify", "pruned", "free_after_gb")})[:900])
+    say("RUN_JSON " + json.dumps({"final_step": step, "val_ppl": ppl, "tokens_consumed": consumed,
+                                 "tok_per_s": round(consumed / max(1, elapsed)),
+                                 "final_loss": final_loss, "prior_loss": prior}, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()