p4_stop_probe: E-044 repin (rank-0 stop assertions; memory gate moved onto the training-only leg)

kernels/p4_stop_probe.py

@@ -24,7 +24,7 @@ import hashlib, json, os, re, shutil, signal, subprocess, sys, threading, time
 
 os.chdir("/kaggle/working")
 sys.path.insert(0, "/kaggle/working")
-REV = "b6210a298619be7534b4a06e722bd31592f5d986"
+REV = "55b8fc47dd7799bf3fc08b7943421f578bac4a2c"
 WANT = {
     "ounce100m_credentials.py": ("ounce100m_credentials.py",
                                  "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
@@ -33,7 +33,7 @@ WANT = {
     "hubckpt.py": ("train/hubckpt.py",
                    "d4b50ed0928c678c94ce6764612a4655f2959c0a7b91d6fd588e8efddff298b2"),
     "train_ounce100m.py": ("train/train_ounce100m.py",
-                           "d795bb388dab42be5efedf133d0124265b3a0c60ddd8f5c375986a85de70b954"),
+                           "dcb0ac199c0616575423ddea0d6b76d2257c06089d3f73aae0e0359de958ff03"),
 }
 BASE = "https://huggingface.co/Cion-lab/ounce100m-code/resolve/" + REV
 for p, (rp, want) in sorted(WANT.items()):
@@ -276,17 +276,17 @@ res["checks"] = {
     and rj2.get("tokens_consumed") == STEPS * TPS,
     "params_are_the_frozen_model": rj1.get("params") == 106194240,
     "checkpointing_really_off": rj1.get("grad_ckpt") is False,
-    # 13.6 GB of ~14.56 usable leaves the run room for a fragmentation spike and for the two CUDA contexts;
-    # above that the answer to the user's question is "not on this hardware", not "keep going". Reported in
-    # both modes, gated only in the soak: a 20-step peak is not evidence about 3,814 steps.
-    # rank 0's allocated peak is not the constraint: reserved bytes, the CUDA contexts and the *other*
-    # rank's peak are what an OOM is actually about, so the max across ranks is what gets gated.
-    "peak_memory_below_13_6_gb": (not GATE_PEAK) or (
-        (rj1.get("peak_reserved_gb_max_rank") or rj1.get("peak_gpu_gb_max_rank") or 99) <= 13.6
-        and (rj2.get("peak_reserved_gb_max_rank") or rj2.get("peak_gpu_gb_max_rank") or 99) <= 13.6),
-    # Leg 1 must skip validation (a forced stop cannot enter the collective safely -- that is the crash
-    # this probe found), leg 2 must actually run it. Checking "val_ppl is not None" on both would demand
-    # the very thing that kills the rank.
+    # What gets gated is leg 1, because it is the only leg whose number means "training". `max_memory_reserved`
+    # is the high-water mark since the process started, and leg 2 runs the validation pass before sampling
+    # it, so leg 2's figure is training-plus-eval -- on smoke that came out 13.66 GB reserved against 12.7
+    # allocated, while the segment that never evaluated is the state the run sits in for 3,814 steps. Gating
+    # on the eval-inflated number would fail D-018 for a state the run never occupies. The other rank's peak
+    # and the reserved bytes are what an OOM is actually about, hence `max_rank` rather than rank 0 alone.
+    "peak_memory_during_training_below_13_6_gb": (not GATE_PEAK) or (
+        (rj1.get("peak_reserved_gb_max_rank") or rj1.get("peak_gpu_gb_max_rank") or 99) <= 13.6),
+    # Leg 1 skips validation by design (E-040's fix; the crash it was written for turned out to be E-044,
+    # but the skip still stands -- a mid-run PPL point is not worth an untested path in a billed session),
+    # and leg 2 must actually run it, because that is where the report's PPL comes from.
     "leg1_skipped_validation_by_design": rj1.get("val_skipped") is True and rj1.get("val_ppl") is None,
     "leg2_validation_actually_ran": (rj2.get("val_ppl") is not None and rj2.get("val_error") is None
                                      and rj2.get("val_skipped") is False),