Upload train/preflight.py with huggingface_hub

train/preflight.py

@@ -29,8 +29,16 @@ 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 {})))
+    try:
+        p = subprocess.run(argv, cwd=WORK, capture_output=True, text=True, timeout=timeout,
+                           env=dict(os.environ, **(env or {})))
+    except subprocess.TimeoutExpired as ex:
+        # A timeout is a result, not an exception to propagate: letting it raise discarded every cell
+        # that had already completed and orphaned both T4s (review finding).
+        print(f"   TIMEOUT after {time.time() - t0:.0f}s", flush=True)
+        return {"rc": -9, "out": (ex.stdout or b"").decode()[-200000:],
+                "err": (ex.stderr or b"").decode()[-3000:],
+                "seconds": round(time.time() - t0, 1), "timeout": True}
     for line in (p.stdout or "").splitlines():
         print("   |", line[:240], flush=True)
     if p.returncode != 0:
@@ -93,6 +101,10 @@ kw = dict(output_dir="/kaggle/working/argcheck", per_device_train_batch_size=2,
           adam_beta2=0.95, adam_epsilon=1e-8, max_grad_norm=1.0, lr_scheduler_type="constant",
           warmup_ratio=0.0, num_train_epochs=1, max_steps=10, fp16=True, bf16=False,
           gradient_checkpointing=False, ddp_find_unused_parameters=False, dataloader_num_workers=2,
+          # the three knobs the tokens/sec sweep turns, checked here on CPU for free before any GPU cell
+          # spends billed seconds discovering transformers 5 renamed or rejected one
+          gradient_checkpointing_kwargs={"use_reentrant": False}, optim="adamw_torch_fused",
+          torch_compile=False, per_device_eval_batch_size=4,
           dataloader_pin_memory=False, remove_unused_columns=False, ignore_data_skip=True,
           save_strategy="steps", save_steps=5, save_total_limit=1, logging_steps=5, report_to=[],
           seed=1, data_seed=1, accelerator_config={"dispatch_batches": False},
@@ -372,6 +384,99 @@ def p2_throughput(args, R):
     print("VERDICT P2_pass=", R["P2_pass"], "eta_hours:", json.dumps(R["P2_eta_hours_at_1B"]), flush=True)
 
 
+def p_opts(args, R):
+    """Tokens/sec, again -- because P3 came in at 5.99 GB peak on a 15.4 GB card and every number used to
+    freeze D-011 was measured at accum 1, not the run's accum 32.
+
+    Gradient checkpointing exists to buy memory, and it costs a full extra forward pass per layer. At
+    5.99 GB of ~15.4 GB it is buying headroom nobody is spending, so the honest question is whether the
+    frozen config is over-paying for it. Cells below change only *how* the same maths is executed --
+    micro-batch/accum multiply to the same 262,144 tokens/step in every row, and the optimiser, LR,
+    schedule, seed and data are identical -- so a winner can be adopted at freeze time without touching
+    what the run is a test of.
+
+    A baseline cell runs FIRST and the driver reports gains against it, because comparing across sessions
+    on shared T4s is how the 4,071-vs-11,062 confusion happened in the first place.
+    """
+    cells = [
+        ("base_m4_gc",         20, ["--micro-batch", "4", "--accum", "32", "--grad-ckpt"],
+         {"micro_batch": 4, "accum": 32, "grad_ckpt": True}),
+        ("m4_nogc",            20, ["--micro-batch", "4", "--accum", "32", "--no-grad-ckpt"],
+         {"micro_batch": 4, "accum": 32, "grad_ckpt": False}),
+        ("m8_gc_acc16",        20, ["--micro-batch", "8", "--accum", "16", "--grad-ckpt"],
+         {"micro_batch": 8, "accum": 16, "grad_ckpt": True}),
+        ("m8_nogc_acc16",      20, ["--micro-batch", "8", "--accum", "16", "--no-grad-ckpt"],
+         {"micro_batch": 8, "accum": 16, "grad_ckpt": False}),
+        ("m16_gc_acc8",        20, ["--micro-batch", "16", "--accum", "8", "--grad-ckpt"],
+         {"micro_batch": 16, "accum": 8, "grad_ckpt": True}),
+        ("m4_gc_nonreentrant", 20, ["--micro-batch", "4", "--accum", "32", "--grad-ckpt",
+                                    "--gc-nonreentrant"],
+         {"micro_batch": 4, "accum": 32, "grad_ckpt": True, "gc_nonreentrant": True}),
+        ("m4_gc_fusedoptim",   20, ["--micro-batch", "4", "--accum", "32", "--grad-ckpt",
+                                    "--optim", "adamw_torch_fused"],
+         {"micro_batch": 4, "accum": 32, "grad_ckpt": True, "optim": "adamw_torch_fused"}),
+    ]
+    common = ["--seq-len", str(args.seq_len), "--tokens", str(args.tokens), "--attn", args.attn,
+              "--hub-repo", "", "--log-every", "10", "--val-tokens", "200000"]
+    out = {}
+    base_rate = None
+    for label, steps, extra, expect in cells:
+        # Every row keeps the frozen 262,144 tokens/step: micro_batch x accum x 2 cards is constant, so a
+        # winner changes how the same maths is executed, never what is being trained.
+        if expect["micro_batch"] * expect["accum"] * 2 * args.seq_len != 262144:
+            raise SystemExit(f"cell {label} does not preserve the frozen batch size in tokens")
+        # An OOM in one row is a result, not a systematic defect: unlike P2, keep going and let the table
+        # show which configurations this hardware can hold. --push-every-steps = the whole cell means one
+        # checkpoint write per cell instead of ten, which would otherwise dominate elapsed and fill the
+        # 19.5 GB working directory (review finding).
+        r = torchrun(args, common + extra + ["--max-steps", str(steps),
+                                             "--push-every-steps", str(steps),
+                                             "--out", WORK + "/opt_" + label], timeout=1500)
+        cell = {"rc": r["rc"], "seconds": r["seconds"], "argv": " ".join(extra),
+                "expect": expect, "tokens_per_step": expect["micro_batch"] * expect["accum"] * 2
+                * args.seq_len}
+        for line in r["out"].splitlines():
+            if line.startswith("RUN_JSON "):
+                try:
+                    cell.update(json.loads(line[len("RUN_JSON "):]))
+                except Exception as e:
+                    cell["run_json_unparseable"] = str(e)[:120]
+            elif line.startswith("precision:"):
+                cell["precision"] = line[:240]
+            elif "OutOfMemory" in line or "CUDA out of memory" in line:
+                cell["oom"] = True
+        # Trust but verify: the cell is only the configuration it claims to be if the trainer says so.
+        cell["expect_mismatch"] = {k: [v, cell.get(k)] for k, v in expect.items()
+                                   if cell.get(k) is not None and cell.get(k) != v}
+        if "tok_per_s" not in cell:
+            cell["tail"] = (r["out"] or "")[-1200:]
+            cell["err_first_error"] = ((r["err"] or "")[-1500:])
+        if label == "base_m4_gc":
+            base_rate = cell.get("tok_per_s")
+        if base_rate and cell.get("tok_per_s"):
+            cell["gain_vs_base_pct"] = round(100.0 * (cell["tok_per_s"] / base_rate - 1.0), 1)
+        out[label] = cell
+        print(f"VERDICT opt {label} rc={cell['rc']} tok_per_s={cell.get('tok_per_s')} "
+              f"peak_gb={cell.get('peak_gpu_gb')} gain={cell.get('gain_vs_base_pct')} "
+              f"mismatch={cell['expect_mismatch'] or 'none'}", flush=True)
+        # Nothing here needs to survive the cell, and 8 cells x ~3 GB of export would fill the disk and
+        # take the whole session down with it (§3.13's free-space discipline).
+        sh(["bash", "-c", "rm -rf " + WORK + "/opt_" + label + " && df -h " + WORK + " | tail -1"],
+           label="cleanup " + label)
+    R["OPT"] = out
+    got = {k: v["tok_per_s"] for k, v in out.items() if v.get("tok_per_s")}
+    best = max(got, key=got.get) if got else None
+    R["OPT_best"] = best
+    R["OPT_best_tok_per_s"] = got.get(best)
+    R["OPT_best_eta_hours_at_1B"] = (round(1e9 / got[best] / 3600, 2) if best else None)
+    # "Passing" here means the table is trustworthy, not that anything improved: every cell must have
+    # either run or OOMed (a real answer), and none may have died in a way that hides a defect.
+    R["OPT_pass"] = bool(len(out) == len(cells) and base_rate and
+                         all(("tok_per_s" in v) or v.get("oom") for v in out.values())
+                         and not any(v["expect_mismatch"] for v in out.values()))
+    print("VERDICT OPT_pass=", R["OPT_pass"], "best=", best, got, flush=True)
+
+
 def p3_cold_resume(args, R):
     """Three legs, each ending in one checkpoint, each started from the Hub with nothing local but the mix.
 
@@ -458,7 +563,9 @@ def p3_cold_resume(args, R):
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--stage", required=True,
-                    choices=["pargs", "p0", "p1", "p2", "p3", "summary", "all"])
+                    choices=["pargs", "p0", "p1", "p2", "p3", "opts", "summary", "all"],
+                    help="opts is the tokens/sec re-audit at the run's real accumulation; it is not a "
+                         "Gate 3 stage and does not enter GATE_3_READY")
     ap.add_argument("--seq-len", type=int, default=1024)   # D-011: the config the main run will use
     ap.add_argument("--attn", default="eager", choices=["eager", "sdpa"],
                     help="passed through to the trainer by P3; named so a default change cannot silently"
@@ -489,6 +596,8 @@ def main():
         p2_throughput(args, R)
     if "p3" in todo:
         p3_cold_resume(args, R)
+    if "opts" in todo:
+        p_opts(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 ("PARGS", "P0", "P1", "P2", "P3")}