preflight: P2 parses RUN_JSON per cell and requires the frozen 22L/2048 cell to report

train/preflight.py

@@ -301,20 +301,40 @@ def torchrun(args, extra, timeout=7200):
 
 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."""
+    steady-state. The main run's whole schedule is division by this number. 22L/576 is the frozen shape;
+    20L is the H-column alternative from docs/01-plan.md §2.3."""
     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"],
+            micro = 4 if seq == 1024 else 2
+            r = torchrun(args, ["--layers", str(layers), "--seq-len", str(seq),
+                                "--micro-batch", str(micro), "--accum", "1", "--max-steps", "30",
+                                "--out", WORK + f"/t1_{layers}_{seq}", "--log-every", "10"],
                          timeout=5400)
-            out[f"L{layers}_s{seq}"] = {"rc": r["rc"], "seconds": r["seconds"],
-                                        "tail": r["out"][-1200:]}
+            cell = {"rc": r["rc"], "seconds": r["seconds"], "micro_batch": micro}
+            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("params:"):
+                    cell["params_line"] = line[:200]
+                elif line.startswith("precision:"):
+                    cell["precision"] = line[:200]
+            if "tok_per_s" not in cell:
+                cell["tail"] = r["out"][-1500:]
+                cell["err"] = r["err"][-800:]
+            out[f"L{layers}_s{seq}"] = cell
     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)
+    got = [c for c in out.values() if isinstance(c.get("tok_per_s"), (int, float))]
+    # Pass means all four cells ran AND the frozen cell (22L/seq2048) is reported, because that is the
+    # configuration the main run uses; a run where only the cheap cells succeeded proves nothing.
+    R["P2_pass"] = bool(len(out) == 4 and all(c.get("rc") == 0 for c in out.values())
+                        and len(got) == 4 and out["L22_s2048"].get("params") == 106194240)
+    R["P2_eta_hours_at_1B"] = ({k: round(1e9 / c["tok_per_s"] / 3600, 2)
+                                for k, c in out.items() if c.get("tok_per_s")})
+    print("VERDICT P2_pass=", R["P2_pass"], "eta_hours:", json.dumps(R["P2_eta_hours_at_1B"]), flush=True)
 
 
 def p3_cold_resume(args, R):