Upload kernels/phase4_session.py with huggingface_hub

kernels/phase4_session.py

@@ -0,0 +1,209 @@
+"""One Phase 4 session of the ounce100m main run: resume, train to a checkpoint boundary, stop.
+
+Why this file exists as a *launcher with a planner* rather than a plain torchrun line. The run needs
+35.9 GPU-hours and Kaggle gives 30 h/week on sessions whose cap I have not measured, so it will be
+interrupted -- and an interruption that lands between two checkpoints throws away everything since the
+last one. §2 fixes the checkpoint cadence at one per 10 % of tokens (381 steps, ~3.6 h at the measured
+rate), and §3.1 forbids restarting the run once checkpoint 1 exists. The only lever left is *when a
+session chooses to stop*, so this script computes the largest whole number of checkpoint intervals that
+fits the session's remaining GPU hours and passes that as --stop-after-steps. Every session then ends
+sitting on a pushed, Hub-verified, locally-pruned checkpoint, and no billed work is ever discarded.
+
+Failure modes are deliberately loud: if the budget cannot cover one checkpoint interval it refuses to
+start rather than burning up to 3.6 hours of quota on work that cannot be recovered, and if the Hub
+pointer is unreadable for any reason other than "no checkpoint yet" it aborts (a resume that silently
+starts from step 0 would train the same data twice and destroy the run's honesty).
+
+No credentials in this file: they are fetched at run time by ounce100m_credentials (D-006).
+"""
+
+import json
+import os
+import subprocess
+import sys
+import time
+
+WORK = "/kaggle/working"
+MIX_REPO = "Cion-lab/ounce100m-mix-v1"
+CKPT_REPO = "Cion-lab/ounce100m-ckpt"
+
+# The frozen recipe (D-011, D-012). Anything here that a future session might want to change is a
+# change to the run itself, so it is asserted against the trainer's own numbers below rather than trusted.
+SEQ_LEN = 1024
+MICRO_BATCH = 4
+ACCUM = 32
+WORLD = 2
+TOKENS = 1_000_000_000
+TOKENS_PER_STEP = MICRO_BATCH * SEQ_LEN * ACCUM * WORLD      # 262,144
+HORIZON_STEPS = int(TOKENS / TOKENS_PER_STEP)                 # 3,814 -- the trainer's own int() division
+CKPT_EVERY = max(1, HORIZON_STEPS // 10)                      # 381 -- §2's one-per-10 % cadence
+
+# Planning rate: overwritten from P3's measured frozen-geometry cell (`P3_frozen_rate_tok_per_s`) at launch.  
+PLANNING_TOK_PER_S = float(os.environ.get("PLANNING_TOK_PER_S", "7732"))
+SESSION_GPU_HOURS = float(os.environ.get("SESSION_GPU_HOURS", "11.0"))
+QUOTA_LEFT_HOURS = float(os.environ.get("QUOTA_LEFT_HOURS", "30.0"))
+RESERVE_HOURS = float(os.environ.get("RESERVE_HOURS", "0.6"))  # startup, val eval, the last push
+
+
+def sh(argv, label, timeout=None, env=None):
+    print(f"=== {label}", flush=True)
+    e = dict(os.environ)
+    e["PYTHONPATH"] = WORK
+    e.update(env or {})
+    t0 = time.time()
+    try:
+        p = subprocess.run(argv, cwd=WORK, capture_output=True, text=True, timeout=timeout, env=e)
+    except subprocess.TimeoutExpired as ex:
+        print(f"{label} TIMEOUT after {time.time() - t0:.0f}s", flush=True)
+        return -9, (ex.stdout or b"").decode()[-200000:], (ex.stderr or b"").decode()[-4000:]
+    out, err = p.stdout or "", p.stderr or ""
+    for line in out.splitlines():
+        print("   |", line[:300], flush=True)
+    if p.returncode != 0:
+        lines = err.splitlines()
+        keep = []
+        for i, ln in enumerate(lines):            # first traceback, not a wrapper's tail (E-023)
+            if "Traceback" in ln or "Error" in ln:
+                keep = lines[max(0, i - 2):i + 30]
+                break
+        print("   STDERR_HEAD:\n" + "\n".join(keep or lines[:45])[-4000:], flush=True)
+    print(f"{label}_RC {p.returncode} seconds {round(time.time() - t0, 1)}", flush=True)
+    return p.returncode, out
+
+
+def fetch(rev, want):
+    """Download the run's code at a pinned revision and assert every sha256 before executing any of it."""
+    base = f"https://huggingface.co/Cion-lab/ounce100m-code/resolve/{rev}"
+    import hashlib
+    for local, (repo_path, sha) in sorted(want.items()):
+        rc = subprocess.run(["curl", "-sfL", f"{base}/{repo_path}", "-o", local],
+                            capture_output=True, text=True).returncode
+        assert rc == 0, f"fetch failed: {repo_path} at {rev}"
+        got = hashlib.sha256(open(os.path.join(WORK, local), "rb").read()).hexdigest()
+        assert got == sha, f"SHA MISMATCH {repo_path}: hub {got[:16]} != expected {sha[:16]}"
+        print("OK", local, got[:12], flush=True)
+    print("REV_OK", rev, flush=True)
+
+
+def plan(start_step, free_gb):
+    """Largest whole number of checkpoint intervals that fits this session's remaining hours."""
+    seconds_per_ckpt = CKPT_EVERY * TOKENS_PER_STEP / PLANNING_TOK_PER_S
+    usable = min(SESSION_GPU_HOURS, max(0.0, QUOTA_LEFT_HOURS - RESERVE_HOURS))
+    n = int(usable * 3600.0 // seconds_per_ckpt)
+    remaining_intervals = max(0, (HORIZON_STEPS - start_step + CKPT_EVERY - 1) // CKPT_EVERY)
+    n = min(n, remaining_intervals)
+    stop = min(HORIZON_STEPS, start_step + n * CKPT_EVERY) if n else 0
+    return {"start_step": start_step, "horizon_steps": HORIZON_STEPS, "ckpt_every": CKPT_EVERY,
+            "tokens_per_step": TOKENS_PER_STEP, "planning_tok_per_s": PLANNING_TOK_PER_S,
+            "hours_per_checkpoint_interval": round(seconds_per_ckpt / 3600.0, 3),
+            "session_gpu_hours": SESSION_GPU_HOURS, "quota_left_hours": QUOTA_LEFT_HOURS,
+            "reserve_hours": RESERVE_HOURS, "usable_hours": round(usable, 3),
+            "intervals_this_session": n, "stop_after_steps": stop,
+            "tokens_at_stop": stop * TOKENS_PER_STEP,
+            "pct_of_run_at_stop": round(100.0 * stop * TOKENS_PER_STEP / TOKENS, 2),
+            "free_gb_before_start": free_gb}
+
+
+def main():
+    os.chdir(WORK)
+    sys.path.insert(0, WORK)
+    rev = os.environ.get("OUNCE100M_REV") or sys.argv[1]
+    fetch(rev, {
+        "ounce100m_credentials.py": ("ounce100m_credentials.py",
+                                     "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
+        "shard_dataset.py": ("train/shard_dataset.py",
+                             "4111832e79445fb0ffce6a73a02da4c97255d6c7886762682e04a417f2695a78"),
+        "hubckpt.py": ("train/hubckpt.py",
+                       "7a250e8232dc99d39108913f169c290e1ef03744217d71bc2dcf833f5d784c53"),
+        "train_ounce100m.py": ("train/train_ounce100m.py",
+                               "c8462cac0cabebc65bcfdcd62be3ed53b03e08120e3a6eb9ca2c4dc235cace42"),
+    })
+
+    import ounce100m_credentials
+    print("creds:", json.dumps(ounce100m_credentials.install(verify=True)), flush=True)
+    rc, out = sh(["bash", "-c", "nvidia-smi --query-gpu=name,memory.used,memory.total "
+                     "--format=csv,noheader; df -h /kaggle/working | tail -1; free -g | head -2"],
+                 "env")
+    free_gb = None
+    for line in out.splitlines():
+        if "overlay" in line or line.startswith("/dev/"):
+            parts = line.split()
+            if len(parts) >= 4:
+                free_gb = parts[3]
+
+    # The Hub pointer is the only truth about where the run is (never local state -- §3.13).
+    import hubckpt
+    tok = os.environ["HF_TOKEN"]
+    ptr = hubckpt.latest_pointer(CKPT_REPO, repo_type="dataset", token=tok)
+    if ptr.get("error"):
+        print(f"FATAL: cannot read latest.json from {CKPT_REPO}: {ptr['error'][:300]}\n"
+              "Refusing to guess a start step: the alternative is training data that was already read.",
+              flush=True)
+        raise SystemExit(6)
+    start_step = int(ptr.get("step", 0)) if not ptr.get("not_found") else 0
+    print(f"hub says: {'no checkpoint yet -- starting from step 0' if start_step == 0 else ptr}",
+          flush=True)
+    if start_step >= HORIZON_STEPS:
+        print("VERDICT RUN_ALREADY_COMPLETE at step", start_step, flush=True)
+        return
+
+    # Data: the published mix, fetched by anonymous read so this works on any instance. The checks run
+    # here rather than inside the downloader so a quoting mistake cannot turn into a silent pass.
+    rc, _ = sh([sys.executable, "-c",
+                "import sys; sys.path.insert(0, '/kaggle/working');\n"
+                "from huggingface_hub import snapshot_download\n"
+                f"p = snapshot_download(repo_id='{MIX_REPO}', repo_type='dataset',\n"
+                "                        local_dir='/kaggle/working/mixroot')\n"
+                "print('downloaded to', p)\n"],
+               "fetch_mix", 3600)
+    if rc != 0:
+        raise SystemExit(7)
+    man = json.load(open(os.path.join(WORK, "mixroot", "manifest.json")))
+    windows = int(man["total_tokens"]) // SEQ_LEN
+    needed_windows = HORIZON_STEPS * MICRO_BATCH * ACCUM * WORLD
+    print(f"mix: {man['total_tokens']:,} train tokens in {man['n_shards']} shards, "
+          f">= {man['distinct_sources_per_shard_min']} sources per shard, "
+          f"{windows:,} windows of {SEQ_LEN}; the horizon needs {needed_windows:,} "
+          f"({100.0 * windows / needed_windows - 100:.1f} % margin)", flush=True)
+    if int(man["total_tokens"]) < TOKENS or windows < needed_windows:
+        # Training past the end of the dataset would end the epoch early and report a token count that
+        # was never reached. The token target is a claim in the final report, so it is checked here.
+        print(f"VERDICT REFUSED_TO_START: mix holds {man['total_tokens']:,} tokens / {windows:,} "
+              f"windows but the run needs {TOKENS:,} / {needed_windows:,}. Either the target drops to "
+              "0.9 B (D-012's fallback) or more sources are staged -- not a decision to make "
+              "implicitly by launching anyway.", flush=True)
+        raise SystemExit(9)
+
+    pl = plan(start_step, free_gb)
+    print("PLAN_JSON", json.dumps(pl, indent=1), flush=True)
+    if not pl["intervals_this_session"]:
+        print(f"VERDICT REFUSED_TO_START: {pl['usable_hours']} usable hours will not reach one "
+              f"{pl['hours_per_checkpoint_interval']} h checkpoint interval. Launching anyway would "
+              "burn quota on work that cannot be resumed from. Raise SESSION_GPU_HOURS or wait for "
+              "quota.", flush=True)
+        raise SystemExit(8)
+
+    argv = [sys.executable, "-u", "train_ounce100m.py",
+            "--root", f"{WORK}/mixroot", "--out", f"{WORK}/run",
+            "--hub-repo", CKPT_REPO, "--prune",
+            "--seq-len", str(SEQ_LEN), "--attn", "eager", "--grad-ckpt",
+            "--micro-batch", str(MICRO_BATCH), "--accum", str(ACCUM),
+            "--tokens", str(TOKENS), "--lr", "6e-4", "--warmup-frac", "0.02", "--decay-frac", "0.80",
+            "--log-every", "20", "--val-tokens", "2000000",
+            "--resume", "auto", "--stop-after-steps", str(pl["stop_after_steps"])]
+    rc, out = sh(["torchrun", "--nproc_per_node=2"] + argv, "TRAIN",
+                 timeout=int((pl["intervals_this_session"] * pl["hours_per_checkpoint_interval"] + 1.5)
+                             * 3600))
+    for line in out.splitlines():
+        if line.startswith(("CKPT ", "RUN_JSON ", "resume from", "auto-resume", "precision:",
+                            "params:", "validation loss", "segment boundary")):
+            print("KEY>", line[:400], flush=True)
+    print("VERDICT PHASE4_SESSION rc", rc, "steps", start_step, "->", pl["stop_after_steps"],
+          "tokens_at_stop", pl["tokens_at_stop"], flush=True)
+    sh(["bash", "-c", "df -h /kaggle/working | tail -1"], "disk_after")
+    # The session's own wall clock is the measurement that sharpens every later plan.
+    raise SystemExit(rc)
+
+
+if __name__ == "__main__":
+    main()