train: a segment that stopped short of the horizon skips the validation pass instead of entering a collective the two ranks can reach a batch apart (probe 1 v1/v4 found rank 1 dying there with no traceback, while every horizon leg passed the same call)

train/train_ounce100m.py

@@ -612,20 +612,32 @@ def main():
     if final_loss is None or not math.isfinite(final_loss):
         raise SystemExit(f"final loss is not finite ({final_loss}) -- not exporting anything")
 
-    # Validation first, and on every rank: `evaluate()` is collective, so gating it would deadlock.
+    # A segment that stopped short of the horizon does not run the validation pass, on every rank.
+    #
+    # Probe 1 v1 and v4 showed the same split twice: the leg stopped early by `--stop-after-steps` died inside
+    # `evaluate()` on rank 1 with no Python traceback at all, while the leg that reached its horizon ran the
+    # same call and exported cleanly. `should_training_stop` is raised from `on_step_end`, so the two ranks
+    # can leave the inner loop one batch apart, and a collective entered in that state kills a rank without
+    # raising -- which is exactly the signature. The report's PPL comes from the horizon leg, and the
+    # mid-run legs keep their training loss, which is what §5 asks to watch on every wake.
+    forced_stop = bool(hub_cb and hub_cb.stopped_at and step < steps_planned)
     ppl = None
     val_err = None
-    try:
-        ev = tr.evaluate()
-        ppl = math.exp(min(20.0, ev["eval_loss"]))
-        say(f"validation loss={ev['eval_loss']:.4f} ppl={ppl:.2f} over {n_val} held-out windows")
-    except Exception as e:
+    if forced_stop:
+        say(f"validation skipped: segment ended at step {step} of {steps_planned} (evaluate() after a "
+            "forced stop diverges between ranks -- probe 1 v1/v4)")
+    else:
+        try:
+            ev = tr.evaluate()
+            ppl = math.exp(min(20.0, ev["eval_loss"]))
+            say(f"validation loss={ev['eval_loss']:.4f} ppl={ppl:.2f} over {n_val} held-out windows")
+        except Exception as e:
         # `say()` is a no-op on rank 1, and that silence is what made a rank-1 evaluation crash unreadable
         # in probe 1 v1: the handler ran on both ranks, only rank 0 spoke, and the divergence looked like an
         # unexplained exit 1 with `error_file: <N/A>`. Report on every rank, and carry the failure into the
         # record so a missing PPL can never be read as a passing validation pass.
-        val_err = f"{type(e).__name__}: {str(e)[:300]}"
-        print(f"[rank {rank}] validation eval failed: {val_err}", flush=True)
+            val_err = f"{type(e).__name__}: {str(e)[:300]}"
+            print(f"[rank {rank}] validation eval failed: {val_err}", flush=True)
 
     # Sampled here, on every rank, before the rank-0-only export block: a collective reduction inside
     # `if rank == 0:` would hang, and reading rank 0's allocator alone is what made the probe's memory
@@ -674,13 +686,14 @@ def main():
                    "attn": args.attn, "effective_attn": eff_attn, "grad_ckpt": bool(args.grad_ckpt),
                    "optim": args.optim, "lr": args.lr, "warmup_frac": args.warmup_frac,
                    "decay_frac": args.decay_frac, "seed": args.seed, "data_seed": args.data_seed,
-                   "val_ppl": ppl, "val_loss_error": val_err, "val_windows": n_val,
+                   "val_ppl": ppl, "val_loss_error": val_err, "val_skipped": forced_stop,
+                   "val_windows": n_val,
                    "peak_gpu_gb": mem["rank0_allocated_gb"],
                    "peak_reserved_gb": mem["rank0_reserved_gb"],
                    "peak_gpu_gb_max_rank": mem["max_allocated_gb"],
                    "peak_reserved_gb_max_rank": mem["max_reserved_gb"]},
                   open(os.path.join(exp, "run_summary.json"), "w"), indent=1, sort_keys=True)
-    segment = bool(hub_cb and hub_cb.stopped_at and step < steps_planned)
+    segment = forced_stop        # the same predicate, computed before the validation decision above
     if api is not None and not segment and rank == 0:
         r = hubckpt.push_and_prune(args.hub_repo, exp, "final", api, repo_type=args.hub_repo_type,
                                    token=os.environ.get("HF_TOKEN"), prune=args.prune)
@@ -722,7 +735,8 @@ def main():
                                  # So a caller cannot report success for a validation pass that never
                                  # happened: a null val_ppl with no error here would be a bug, and a null
                                  # with an error is a measured failure.
-                                 "val_error": val_err, "val_windows": n_val}, sort_keys=True))
+                                 "val_error": val_err, "val_skipped": forced_stop, "val_windows": n_val},
+                                 sort_keys=True))
 
 
 def peak_stats():