preflight: assert the trapezoid SHAPE (off-by-one at the last step is expected, not a bug)

train/preflight.py

@@ -109,8 +109,9 @@ mults = []
 for i in range(steps):
     mults.append(sch.lr_lambdas[0](i))
     sch.optimizer.param_groups[0]["lr"] = 6e-4 * mults[-1]
-R["lr_shape"] = {"peak": max(mults), "at_warm_end": mults[76], "at_80pct": mults[int(steps*0.8)],
+R["lr_shape"] = {"peak": max(mults), "at_step76": mults[76], "at_80pct": mults[int(steps*0.8)],
                  "at_90pct": mults[int(steps*0.9)], "final": mults[-1],
+                 "hits_zero_only_past_last_step": mults[-1] > 0.0,
                  "plateau_fraction_steps": round(sum(1 for x in mults if x > 0.999) / steps, 3)}
 # one real forward/backward on CPU with a tiny model to prove the collator output feeds the model.
 # The report is printed BEFORE this optional probe: v3 of pargs lost every finding it had made because
@@ -145,16 +146,25 @@ _emit(R)
     R["PARGS_n_blocks"] = max(0, len(blocks) - 1)
     R["PARGS_rc"] = r["rc"]
     j = j or {}
+    lr = j.get("lr_shape") or {}
+    # Test the SHAPE of the schedule, not one value: flat plateau, then linear to ~0 at the final
+    # update. (v4 asserted final == 0.0 and failed on 0.0013 -- LambdaLR reaches 0 at step T, which is
+    # one past the last update at T-1. Every cosine implementation shares that off-by-one; the honest
+    # assertion is "effectively zero at the last step", not "exactly zero".)
+    shape_ok = (lr.get("peak") == 1.0 and lr.get("at_80pct") == 1.0
+                and 0.45 < lr.get("at_90pct", -1) < 0.55
+                and 0.0 <= lr.get("final", 1.0) < 0.005
+                and 0.70 < lr.get("plateau_fraction_steps", 0) < 0.85)
     R["PARGS_pass"] = bool(j.get("PARGS_DONE") and j.get("kwargs") == "accepted"
                            and j.get("params", {}).get("sum_numel") == 106194240
                            # 20L must be 99,114,048 per docs/01-plan.md §2.3's closed-form table, and its
                            # FFN must stay 1536; a doubled FFN would make test T1 meaningless.
                            and j.get("shapes_20L", {}).get("sum_numel") == 99114048
-                           and j.get("shape20_ffn") == 1536
-                           and (j.get("lr_shape") or {}).get("final") == 0.0
+                           and j.get("shape20_ffn") == 1536 and shape_ok
                            and (j.get("forward_backward") or {}).get("finite"))
-    print("VERDICT PARGS_pass=", R["PARGS_pass"], "blocks=", R["PARGS_n_blocks"],
-          json.dumps({k: j.get(k) for k in ("transformers", "torch", "kwargs", "forward_backward")})[:400],
+    print("VERDICT PARGS_pass=", R["PARGS_pass"], "blocks=", R["PARGS_n_blocks"], "shape_ok=", shape_ok,
+          json.dumps({k: j.get(k) for k in ("transformers", "torch", "kwargs", "params", "shapes_20L",
+                                            "lr_shape", "forward_backward")}, default=str)[:700],
           flush=True)