add read-only capture diagnostic for the 4096-cap parity bug

diag/diagnose_parity_capture.py

@@ -124,6 +124,79 @@ def repeat_fraction(ids: list[int], lag: int) -> float:
     return sum(1 for i in range(len(ids) - lag) if ids[i] == ids[i + lag]) / (len(ids) - lag)
 
 
+def stored_pair_analysis(ids, mask, lengths, behavior_all, trainer_flat, band: int,
+                         max_shift: int) -> dict[str, Any]:
+    """Structure of the delta that actually tripped the guard, with NO model.
+
+    `batch.pt` carries both sides of the comparison: `stu_behavior_log_probs` (what
+    SGLang returned) and `actual_log_probs` (what the trainer computed, masked-flat
+    in row-major order). Their difference IS the guard's delta. That means the
+    structural questions -- which row, where, contiguous or strided, what period --
+    are answerable from the capture alone, with no checkpoint, no GPU and no HF.
+
+    This matters because the capture holds a 6.4 GB student checkpoint and the bug
+    note flags a cleanup policy for it: if it is gone, a model-dependent diagnostic
+    produces nothing at all, while this path still reports the primary evidence.
+    """
+    tv: list[float] = []          # stored trainer logprobs, aligned order
+    bv: list[float] = []          # stored behaviour logprobs, aligned order
+    positions: list[int] = []
+    rows: list[int] = []
+    cursor = 0
+    for r in range(ids.shape[0]):
+        n = int(lengths[r])
+        for p in mask[r, :n].nonzero().squeeze(-1).tolist():
+            if cursor >= len(trainer_flat):
+                break
+            t = float(trainer_flat[cursor])
+            b = float(behavior_all[r, p])
+            cursor += 1
+            if not math.isfinite(b):
+                continue
+            tv.append(t)
+            bv.append(b)
+            positions.append(p)
+            rows.append(r)
+    delta = [abs(a - b) for a, b in zip(tv, bv)]
+    out: dict[str, Any] = {
+        "note": "stored behaviour vs stored trainer logprobs; this is the guard's delta",
+        "compared_tokens": len(delta),
+        "guard_stored_pair": stats(delta, positions),
+        "per_row": {},
+        "per_band": bands(delta, band),
+        "shift_test_stored_pair": {},
+        "periodicity": {},
+        "outlier_token_id_histogram_top": {},
+    }
+    for r in range(ids.shape[0]):
+        sel = [i for i, rr in enumerate(rows) if rr == r]
+        if sel:
+            out["per_row"][str(r)] = stats([delta[i] for i in sel],
+                                           [positions[i] for i in sel])
+    # Alignment test between the two STORED arrays: compare trainer[i] against
+    # behaviour[i + k]. If the engine's values were shifted relative to the
+    # trainer's, a non-zero shift would beat shift 0 here. (Comparing the deltas to
+    # each other, as a first draft did, is meaningless.)
+    for k in range(-max_shift, max_shift + 1):
+        pairs = [abs(tv[i] - bv[i + k]) for i in range(len(tv)) if 0 <= i + k < len(bv)]
+        out["shift_test_stored_pair"][str(k)] = stats(pairs)
+    for r in range(ids.shape[0]):
+        n = int(lengths[r])
+        row_ids = [int(v) for v in ids[r, :n].tolist()]
+        out["periodicity"][str(r)] = {
+            str(lag): round(repeat_fraction(row_ids, lag), 4)
+            for lag in (1, 2, 3, 4, 5, 6, 7, 8, 12, 16, 24, 32)}
+    outliers = [i for i, v in enumerate(delta) if v > 0.5]
+    if outliers:
+        tids = [int(ids[rows[i], positions[i]].item()) for i in outliers]
+        hist: dict[str, int] = {}
+        for t in tids:
+            hist[str(t)] = hist.get(str(t), 0) + 1
+        out["outlier_token_id_histogram_top"] = dict(
+            sorted(hist.items(), key=lambda kv: -kv[1])[:10])
+    return out
+
+
 def score_row(model, row_ids: list[int], temperature: float) -> dict[str, list[float]]:
     """Per-position logprobs for positions 1..L-1 of one unpadded row."""
     import torch
@@ -169,6 +242,9 @@ def main() -> None:
                     help="also rescore the capture's exact tokens with SGLang under the "
                          "production flags, to separate 'the stored values were corrupted' "
                          "from 'the engine is wrong today'")
+    ap.add_argument("--no-model", action="store_true",
+                    help="skip every checkpoint-dependent check and report only the "
+                         "stored-pair structure (works from batch.pt alone)")
     args = ap.parse_args()
 
     import torch
@@ -198,12 +274,70 @@ def main() -> None:
                                 "attention_backend", "source_commit")}, sort_keys=True),
         flush=True)
 
+    lengths = attn.sum(-1).tolist()
+
+    # ---------------------------------------------------------------------
+    # Model-free structural analysis, ALWAYS run first.
+    #
+    # `batch.pt` holds both sides of the guard's comparison, so the structure of
+    # the failure is available with no checkpoint, no GPU and no HF. The capture
+    # carries a 6.4 GB student checkpoint and the bug note flags a cleanup policy
+    # for it; if it is gone, a model-dependent diagnostic reports nothing at all,
+    # while this path still reports the primary evidence.
+    # ---------------------------------------------------------------------
+    trainer_flat = batch["actual_log_probs"].tolist()
+    stored = stored_pair_analysis(ids, mask, lengths, behavior_all, trainer_flat,
+                                  args.band, args.max_shift)
+    report["stored_pair"] = stored
+    print("\n== STORED-PAIR structure (no model needed) ==")
+    row = stored["guard_stored_pair"]
+    print(f"  guard delta from the capture itself: n={row['tokens']} "
+          f"mean={row['mean']:.6f} p99={row['p99']:.4f} max={row['max']:.4f} "
+          f">0.5={row['above_0p5']}")
+    if row.get("above_0p5_position_head"):
+        pos = row["above_0p5_position_head"]
+        print(f"  first outlier positions: {pos[:24]}")
+        # stats() already computes the gap histogram over the first 300 outliers.
+        print(f"  gap histogram          : "
+              f"{json.dumps(row.get('above_0p5_gap_histogram') or {})}")
+        mod: dict[str, int] = {}
+        for p in pos:
+            mod[str(p % 8)] = mod.get(str(p % 8), 0) + 1
+        print(f"  mod-8 histogram        : {json.dumps(mod)}")
+        print(f"  span                   : {row.get('above_0p5_span_positions')}")
+    for r, s in stored["per_row"].items():
+        print(f"  row {r}: n={s['tokens']} mean={s['mean']:.6f} p99={s['p99']:.4f} "
+              f"max={s['max']:.4f} >0.5={s['above_0p5']}")
+    print("  per band (mean/max): " + json.dumps(
+        [[b[0], b[1], round(b[3], 5), round(b[4], 4)] for b in stored["per_band"]]))
+    print("  shift test (stored pair): " + json.dumps(
+        {k: [round(v["mean"], 6), round(v["max"], 4)]
+         for k, v in stored["shift_test_stored_pair"].items()}))
+    print("  periodicity: " + json.dumps(stored["periodicity"]))
+    if stored["outlier_token_id_histogram_top"]:
+        print("  outlier token-id histogram (top): "
+              + json.dumps(stored["outlier_token_id_histogram_top"]))
+
+    have_model = not args.no_model and (args.capture / "student").is_dir()
+    if not have_model:
+        reason = ("--no-model was requested" if args.no_model
+                  else f"no student checkpoint at {args.capture / 'student'}")
+        report["model_skipped"] = reason
+        print(f"\nPARITY_DIAG_MODEL=skipped ({reason})")
+        print("  The HF cross-checks (temperature detector, entropy/margin, padding,"
+              " engine) need the checkpoint and were not run.")
+        if args.out:
+            args.out.write_text(json.dumps(report, indent=2, sort_keys=True,
+                                           default=str) + "\n")
+            print("PARITY_DIAG_OUT=" + str(args.out), flush=True)
+        print("PARITY_DIAG_DONE=1", flush=True)
+        return
+
     model = AutoModelForCausalLM.from_pretrained(
         args.capture / "student", dtype=getattr(torch, args.dtype),
         attn_implementation="eager", local_files_only=True,
         trust_remote_code=False).cuda()
     model.train(bool(meta.get("model_training")))
-    lengths = attn.sum(-1).tolist()
 
     # ---- per-row fresh recomputation (unpadded) -------------------------------
     per_row: list[dict[str, Any]] = []