Upload train/preflight.py with huggingface_hub

train/preflight.py

@@ -388,6 +388,7 @@ def p3_cold_resume(args, R):
     # exercises the real step, the real checkpoint size and the real cursor, not a cheaper stand-in.
     common = ["--seq-len", str(args.seq_len), "--tokens", str(args.tokens),
               "--accum", str(args.accum), "--micro-batch", "4",
+              "--attn", args.attn, "--grad-ckpt",
               "--hub-repo", args.ckpt_repo, "--prune", "--log-every", "5"]
     legs, cursors, losses = [], [], []
     for i in range(3):
@@ -416,20 +417,40 @@ def p3_cold_resume(args, R):
                "cursor_line": cur,
                "tail": out[-1200:] if r["rc"] else None,
                "err": (r["err"] or "")[-900:] if r["rc"] else None}
+        for line in out.splitlines():
+            if line.startswith("RUN_JSON "):
+                try:
+                    leg.update(json.loads(line[len("RUN_JSON "):]))
+                except Exception as e:
+                    leg["run_json_unparseable"] = str(e)[:120]
         legs.append(leg)
-        print(f"VERDICT P3 leg {'abc'[i]} rc={r['rc']} auto_resume={leg['auto_resume']}", flush=True)
+        print(f"VERDICT P3 leg {'abc'[i]} rc={r['rc']} auto_resume={leg['auto_resume']} "
+              f"tok_per_s={leg.get('tok_per_s')} peak_gb={leg.get('peak_gpu_gb')}", flush=True)
         if r["rc"] != 0:
             break
     R["P3"] = {"legs": legs, "ckpt_lines": cursors, "prior_loss_lines": losses}
+    # The last leg is the one that has run 60 steps of the frozen geometry, so it is also the throughput
+    # cell D-012's tripwire is read from. Neither earlier 1024 A/B cell used micro 4 *with* grad-ckpt.
+    frozen = legs[-1] if legs else {}
+    R["P3_frozen_rate_tok_per_s"] = frozen.get("tok_per_s")
+    R["P3_frozen_peak_gpu_gb"] = frozen.get("peak_gpu_gb")
+    R["P3_frozen_params"] = frozen.get("params")
+    R["P3_eta_hours_at_1B"] = (round(1e9 / frozen["tok_per_s"] / 3600, 2)
+                               if frozen.get("tok_per_s") else None)
+    R["P3_D012_fallback_fires"] = (R["P3_frozen_rate_tok_per_s"] < 5150
+                                  if R["P3_frozen_rate_tok_per_s"] else None)
     seen = [int(s.split("samples=")[1].split()[0].replace(",", "")) for s in cursors
             if "samples=" in s]
     R["P3_cursor_sequence"] = seen
     ok = (len(legs) == 3 and all(l["rc"] == 0 for l in legs)
           and all(l["auto_resume"] for l in legs[1:])
           and len(cursors) >= 3 and len(seen) >= 3
-          and seen == sorted(seen) and len(set(seen)) == len(seen))
+          and seen == sorted(seen) and len(set(seen)) == len(seen)
+          # ...and the thing that resumed was the frozen model, not a smaller stand-in.
+          and R["P3_frozen_params"] == 106194240 and bool(R["P3_frozen_rate_tok_per_s"]))
     R["P3_pass"] = bool(ok)
-    print("VERDICT P3_pass=", R["P3_pass"], "cursors=", seen, flush=True)
+    print("VERDICT P3_pass=", R["P3_pass"], "cursors=", seen,
+          "tok_per_s=", R["P3_frozen_rate_tok_per_s"], flush=True)
     R["P3_hub_postclean"] = (sh(["python", "-c", _DEL % args.ckpt_repo],
                                 label="P3 hub postclean")["out"] or "").strip()[-200:]
 
@@ -439,6 +460,9 @@ def main():
     ap.add_argument("--stage", required=True,
                     choices=["pargs", "p0", "p1", "p2", "p3", "summary", "all"])
     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"
+                         " alter what the frozen-geometry measurement tested")
     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=20,