train: Gate 3 preflight driver (reader, hub cycle, T1 throughput, cold resume)

train/preflight.py

@@ -0,0 +1,261 @@
+# Gate 3 preflight driver: proves the pipeline, not the model, and writes the evidence itself.
+
+# Every test here answers a question §5 Phase 3 lists, and each one emits a machine-readable verdict
+# line so a future session can read the outcome without re-reading a log. GPU stages are few and short:
+# the whole file is designed to fit inside the 6 GPU-hour lifetime test cap (memory/QUOTA.md), and it
+# prints what it spent so the ledger can be updated from the log.
+
+# Stages, in the cheap-first order they should run:
+#   P0 (CPU)  reader: cursor math, resume slicing, determinism, val split really held out
+#   P1 (CPU)  hub cycle: push_and_prune with a real-sized 100M checkpoint + optimizer state
+#   P2 (GPU)  throughput: 20L vs 22L at seq 1024/2048 -- test T1, the number the main-run ETA is built on
+#   P3 (GPU)  short train -> kill -> cold resume from the Hub on an empty disk -> loss continuity
+#   P4 (GPU)  resume twice in sequence; the cursor must advance monotonically, never re-read
+#
+# P3/P4 are the tests that matter most and the ones that cannot be faked: a run that resumes from what it
+# left on disk proves nothing, because every real interruption takes the disk with it (§3.13).
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+
+WORK = "/kaggle/working"
+REV_DEFAULT = ""  # filled from the launcher; only used for reporting
+
+
+def sh(argv, timeout=None, env=None, label=""):
+    print(f"=== {label or ' '.join(argv[:3])}", flush=True)
+    t0 = time.time()
+    p = subprocess.run(argv, cwd=WORK, capture_output=True, text=True, timeout=timeout,
+                       env=dict(os.environ, **(env or {})))
+    for line in (p.stdout or "").splitlines():
+        print("   |", line[:240], flush=True)
+    if p.returncode != 0:
+        print("   STDERR:", (p.stderr or "")[-3000:], flush=True)
+    return {"rc": p.returncode, "out": (p.stdout or "")[-200000:],
+            "err": (p.stderr or "")[-3000:], "seconds": round(time.time() - t0, 1)}
+
+
+def last_json(text, begin, end):
+    if begin not in text or end not in text:
+        return None
+    body = text.rsplit(begin, 1)[1].split(end, 1)[0]
+    try:
+        return json.loads(body)
+    except Exception:
+        return None
+
+
+# ------------------------------------------------------------------------ P0 reader (CPU, free)
+def p0_reader(args, R):
+    """The reader is where an unrecoverable main run would hide: if two runs at the same cursor read
+    different tokens, every resume silently trains on a subset. Checked on bytes, not on feelings."""
+    src = r'''
+import json, os, sys, numpy as np
+sys.path.insert(0, "/kaggle/working")
+import shard_dataset as SD
+man = json.load(open("/kaggle/working/mixroot/manifest.json"))
+store = SD.PackedTokenStore("/kaggle/working/mixroot", man)
+L = 256
+full = SD.make_dataset(store, L, 1234)
+n = len(full)
+part = SD.make_dataset(store, L, 1234, start_sample=n - 40)
+same = all(bool((full[n - 40 + j]["input_ids"] == part[j]["input_ids"]).all()) for j in range(40))
+labels_ok = bool((full[0]["labels"][:-1] == full[0]["input_ids"][1:]).all()
+                 and int(full[0]["labels"][-1]) == int(full[0]["input_ids"][0]))
+det2 = SD.make_dataset(store, L, 1234)
+deterministic = bool((full[7]["input_ids"] == det2[7]["input_ids"]).all())
+diffseed = SD.make_dataset(store, L, 999)
+seed_matters = not bool((full[7]["input_ids"] == diffseed[7]["input_ids"]).all())
+vstore = SD.PackedTokenStore("/kaggle/working/mixroot", man,
+                             files=[s["file"] for s in man["val_shards"]])
+tv = set(); vv = set()
+vds = SD.make_dataset(vstore, L, 0, shuffle=False, count=min(400, vstore.total_tokens // L))
+for i in range(min(len(full), 4000)):
+    tv.add(full[i]["input_ids"][:64].numpy().tobytes())
+for i in range(len(vds)):
+    vv.add(vds[i]["input_ids"][:64].numpy().tobytes())
+print("READER_JSON_BEGIN")
+print(json.dumps({
+  "total_tokens": store.total_tokens, "samples_at_L256": n,
+  "resume_matches_uninterrupted": same, "labels_are_next_token": labels_ok,
+  "same_seed_same_order": deterministic, "different_seed_different_order": seed_matters,
+  "val_tokens": vstore.total_tokens, "val_windows_checked": len(vds),
+  "train_val_prefix_collision": len(tv & vv),
+  "tokens_per_shard_min": min(s["tokens"] for s in man["shards"]),
+  "n_shards": len(man["shards"]),
+}))
+print("READER_JSON_END")
+'''
+    r = sh([sys.executable, "-c", src], label="P0 reader")
+    R["P0"] = last_json(r["out"], "READER_JSON_BEGIN", "READER_JSON_END")
+    R["P0_rc"] = r["rc"]
+    j = R["P0"] or {}
+    R["P0_pass"] = bool(r["rc"] == 0 and j.get("resume_matches_uninterrupted")
+                        and j.get("labels_are_next_token") and j.get("same_seed_same_order")
+                        and j.get("different_seed_different_order")
+                        and j.get("train_val_prefix_collision") == 0
+                        and j.get("val_tokens", 0) > 0)
+    print("VERDICT P0_pass=", R["P0_pass"], flush=True)
+
+
+# ------------------------------------------------------------------------ P1 hub cycle (CPU, free)
+def p1_hub(args, R):
+    """§3.13 at real size: a 106M-parameter checkpoint with optimizer state, pushed, verified from the
+    Hub by re-listing and re-hashing, then pruned. Cheap because the weights are random; the BYTES are
+    what is being timed."""
+    src = r'''
+import json, os, sys, time, numpy as np
+sys.path.insert(0, "/kaggle/working")
+import ounce100m_credentials, hubckpt
+ounce100m_credentials.install()
+from huggingface_hub import HfApi
+api = HfApi(); tok = os.environ["HF_TOKEN"]
+repo = "Cion-lab/ounce100m-ckptbench-DELETEME"
+try:
+    api.create_repo(repo_id=repo, repo_type="dataset", exist_ok=True, token=tok)
+except Exception as e:
+    print("repo create:", type(e).__name__, str(e)[:120])
+d = "/kaggle/working/ckptbench/checkpoint-1"
+os.makedirs(d, exist_ok=True)
+# 106,194,240 params x 4 B x 4 arrays (fp16-saved weights + fp32 master + Adam m and v) is the real
+# checkpoint footprint, and it is the size D-007 measured at 1.6 GB. Timing a 425 MB mock would flatter
+# the main run.
+n = 106194240
+t0 = time.time()
+for name in ("weights.fp32", "master.fp32", "adam_m.fp32", "adam_v.fp32"):
+    a = np.lib.format.open_memmap(os.path.join(d, "state." + name.replace(".", "_") + ".npy"),
+                                  dtype=np.float32, mode="w+", shape=(n,))
+    a[:] = np.float32(0.0)
+    a.flush()
+    del a
+json.dump({"step": 1, "samples_consumed": 381500}, open(os.path.join(d, "cursor.json"), "w"))
+up_t0 = time.time()
+try:
+    res = hubckpt.push_and_prune(repo, d, "ckpt/checkpoint-1", api, token=tok, prune=True)
+    ver = res["verify"]
+    # second, independent proof: pull it back into a clean directory and compare hashes
+    back = "/kaggle/working/ckptbench/restored"
+    got = hubckpt.download_checkpoint(repo, "ckpt/checkpoint-1", back, api, token=tok)
+    same = json.load(open(os.path.join(back, "cursor.json"))) == {"step": 1,
+                                                                  "samples_consumed": 381500}
+    # The repo is a timing rig, not an artifact: leaving a 1.7 GB public blob invites a future session
+    # to mistake it for a checkpoint. Everything measurable is in the log by this point.
+    api.delete_repo(repo_id=repo, repo_type="dataset", token=tok)
+except Exception as e:
+    res, ver, got, same = {"error": f"{type(e).__name__}: {str(e)[:300]}"}, {}, {}, False
+print("HUB_JSON_BEGIN")
+print(json.dumps({"push_verify_seconds": round(time.time()-up_t0,1),
+                  "cycle": {k: res.get(k) for k in ("verify","pruned","free_before_gb","free_after_gb")},
+                  "downloaded": got, "readback_ok": bool(ver.get("ok")) and same,
+                  "gen_seconds": round(time.time()-t0,1)}, default=str))
+print("HUB_JSON_END")
+'''
+    r = sh([sys.executable, "-c", src], timeout=5400, label="P1 hub cycle (425 MB mock checkpoint)")
+    R["P1"] = last_json(r["out"], "HUB_JSON_BEGIN", "HUB_JSON_END")
+    R["P1_rc"] = r["rc"]
+    j = R["P1"] or {}
+    cyc = ((j.get("cycle") or {}).get("verify") or {})
+    R["P1_pass"] = bool(r["rc"] == 0 and cyc.get("ok") and j.get("readback_ok")
+                        and (j.get("cycle") or {}).get("pruned"))
+    print("VERDICT P1_pass=", R["P1_pass"], json.dumps(cyc)[:300], flush=True)
+    R["P1_cleanup"] = "delete Cion-lab/ounce100m-ckptbench-DELETEME when done reading it"
+
+
+# ------------------------------------------------------------------------ P2..P4 (GPU)
+def torchrun(args, extra, timeout=7200):
+    return sh(["torchrun", "--nproc_per_node=2", "train_ounce100m.py", "--root",
+               WORK + "/mixroot", "--seq-len", str(args.seq_len)] + extra,
+              timeout=timeout, label="torchrun " + " ".join(extra[:6]))
+
+
+def p2_throughput(args, R):
+    """T1: 20L vs 22L at two sequence lengths, on the real data, for enough steps that the number is
+    steady-state. The main run's whole schedule is division by this number."""
+    out = {}
+    for layers in (20, 22):
+        for seq in (1024, 2048):
+            micro = 8 if seq == 1024 else 2
+            r = torchrun(args, ["--layers", str(layers), "--hidden", "576", "--seq-len", str(seq),
+                                "--micro-batch", str(micro), "--accum", "1", "--max-steps", "25",
+                                "--out", WORK + f"/t1_{layers}_{seq}", "--log-every", "5"],
+                         timeout=5400)
+            out[f"L{layers}_s{seq}"] = {"rc": r["rc"], "seconds": r["seconds"],
+                                        "tail": r["out"][-1200:]}
+    R["P2"] = out
+    R["P2_pass"] = all(v["rc"] == 0 for v in out.values()) and len(out) == 4
+    print("VERDICT P2_pass=", R["P2_pass"], flush=True)
+
+
+def p3_cold_resume(args, R):
+    """Short run -> wipe every local trace -> resume with --resume auto, which must recover the exact
+    sample position from the Hub. Loss continuity is the pass condition: a restart that silently
+    re-initialised would jump the loss."""
+    common = ["--tokens", str(args.tokens), "--accum", str(args.accum), "--micro-batch", "2",
+              "--hub-repo", args.ckpt_repo, "--prune", "--log-every", "5"]
+    a = torchrun(args, common + ["--max-steps", str(args.steps_a), "--out", WORK + "/p3"],
+                 timeout=9000)
+    # erase the instance's memory of the run: local checkpoints, the run dir, and the downloaded mix?
+    # No -- the mix is the dataset and a real interruption keeps it. Only the run state goes.
+    wipe = sh(["bash", "-c", f"rm -rf {WORK}/p3 {WORK}/run; df -h {WORK} | tail -1"],
+              label="P3 wipe local run state")
+    b = torchrun(args, common + ["--max-steps", str(args.steps_a + args.steps_b),
+                                 "--out", WORK + "/p3b", "--resume", "auto"], timeout=9000)
+    R["P3"] = {"first": {"rc": a["rc"], "tail": a["out"][-2500:]},
+               "wipe": wipe["out"][-400:],
+               "resumed": {"rc": b["rc"], "tail": b["out"][-2500:]}}
+    la = [l for l in a["out"].splitlines() if "'loss'" in l or "loss=" in l]
+    lb = [l for l in b["out"].splitlines() if "'loss'" in l or "loss=" in l]
+    R["P3_loss_last_before"] = la[-1][:200] if la else None
+    R["P3_loss_first_after"] = lb[0][:200] if lb else None
+    R["P3_pass"] = bool(a["rc"] == 0 and b["rc"] == 0 and "auto-resume: hub says step" in b["out"]
+                        and lb)
+    print("VERDICT P3_pass=", R["P3_pass"], flush=True)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--stage", required=True,
+                    choices=["p0", "p1", "p2", "p3", "summary", "all"])
+    ap.add_argument("--seq-len", type=int, default=2048)
+    ap.add_argument("--accum", type=int, default=32)
+    ap.add_argument("--tokens", type=int, default=1_000_000_000)
+    ap.add_argument("--steps-a", type=int, default=60)
+    ap.add_argument("--steps-b", type=int, default=60)
+    ap.add_argument("--ckpt-repo", default="Cion-lab/ounce100m-ckptbench-DELETEME")
+    ap.add_argument("--out-json", default=WORK + "/preflight.json")
+    args = ap.parse_args()
+
+    R = {}
+    if os.path.exists(args.out_json):
+        try:
+            R = json.load(open(args.out_json))
+        except Exception:
+            print("existing preflight.json unreadable; starting fresh", flush=True)
+    todo = ["p0", "p1", "p2", "p3"] if args.stage == "all" else [args.stage]
+    t0 = time.time()
+    if "p0" in todo:
+        p0_reader(args, R)
+    if "p1" in todo:
+        p1_hub(args, R)
+    if "p2" in todo:
+        p2_throughput(args, R)
+    if "p3" in todo:
+        p3_cold_resume(args, R)
+    R["seconds_this_invocation"] = round(time.time() - t0, 1)
+    R["gpu_hours_this_invocation"] = round(R["seconds_this_invocation"] / 3600.0, 3)
+    R["PASSES"] = {k: R.get(k + "_pass") for k in ("P0", "P1", "P2", "P3")}
+    R["GATE_3_READY"] = all(v is True for v in R["PASSES"].values())
+    with open(args.out_json, "w") as f:
+        json.dump(R, f, indent=1, default=str)
+    print("PREFLIGHT_JSON")
+    print(json.dumps({"PASSES": R["PASSES"], "GATE_3_READY": R["GATE_3_READY"],
+                      "gpu_hours_this_invocation": R["gpu_hours_this_invocation"]}))
+    print("/PREFLIGHT_JSON")
+
+
+if __name__ == "__main__":
+    main()