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

diag/diagnose_parity_capture.py

@@ -143,6 +143,7 @@ def stored_pair_analysis(ids, mask, lengths, behavior_all, trainer_flat, band: i
     positions: list[int] = []
     rows: list[int] = []
     cursor = 0
+    nan_per_row: dict[str, int] = {}
     for r in range(ids.shape[0]):
         n = int(lengths[r])
         for p in mask[r, :n].nonzero().squeeze(-1).tolist():
@@ -152,6 +153,10 @@ def stored_pair_analysis(ids, mask, lengths, behavior_all, trainer_flat, band: i
             b = float(behavior_all[r, p])
             cursor += 1
             if not math.isfinite(b):
+                # NaN marks the synthetic terminal event on a capped row; it is
+                # intentional and carries no behaviour probability. Count it so the
+                # number can be checked against the capture's own per_sample stats.
+                nan_per_row[str(r)] = nan_per_row.get(str(r), 0) + 1
                 continue
             tv.append(t)
             bv.append(b)
@@ -161,6 +166,8 @@ def stored_pair_analysis(ids, mask, lengths, behavior_all, trainer_flat, band: i
     out: dict[str, Any] = {
         "note": "stored behaviour vs stored trainer logprobs; this is the guard's delta",
         "compared_tokens": len(delta),
+        "nan_sentinel_tokens": sum(nan_per_row.values()),
+        "nan_sentinel_per_row": nan_per_row,
         "guard_stored_pair": stats(delta, positions),
         "per_row": {},
         "per_band": bands(delta, band),
@@ -294,6 +301,9 @@ def main() -> None:
     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']}")
+    print(f"  NaN terminal sentinels skipped: {stored['nan_sentinel_tokens']} "
+          f"per row {json.dumps(stored['nan_sentinel_per_row'])}"
+          "   (1 per capped row is expected)")
     if row.get("above_0p5_position_head"):
         pos = row["above_0p5_position_head"]
         print(f"  first outlier positions: {pos[:24]}")
@@ -408,6 +418,74 @@ def main() -> None:
     rows = [flat_row[i] for i in real]
 
     d_T = [abs(a - b) for a, b in zip(fresh_T, behavior)]
+
+    # ---------------------------------------------------------------------
+    # Padded-batch recomputation, replicating `replay_parity.py --layout batch`.
+    #
+    # On the real capture the two recomputations disagree about which stored array
+    # they reproduce, so this has to be measured rather than argued: a per-row
+    # forward matches `stu_behavior_log_probs` while a padded-batch forward is what
+    # the trainer's own path used. Running BOTH in one process decides which side
+    # is anomalous, and whether padding is the mechanism at all.
+    # ---------------------------------------------------------------------
+    with torch.inference_mode():
+        pos_p = attn.long().cumsum(-1) - 1
+        pos_p.masked_fill_(attn == 0, 1)
+        logits_p = model(input_ids=ids.cuda(), attention_mask=attn.cuda(),
+                         position_ids=pos_p.cuda(), use_cache=False).logits
+        padded_T_all: list[float] = []
+        padded_raw_all: list[float] = []
+        for r in range(ids.shape[0]):
+            n = int(lengths[r])
+            for p in mask[r, :n].nonzero().squeeze(-1).tolist():
+                j = p
+                if j < 0 or j >= n - 1:
+                    padded_T_all.append(float("nan"))
+                    padded_raw_all.append(float("nan"))
+                    continue
+                block = logits_p[r, j].float()
+                label = int(ids[r, j + 1].item())
+                padded_T_all.append(float(torch.log_softmax(block / temperature, -1)[label]))
+                padded_raw_all.append(float(torch.log_softmax(block, -1)[label]))
+        del logits_p
+        torch.cuda.empty_cache()
+    padded_T = sub(padded_T_all)
+    padded_raw = sub(padded_raw_all)
+    report["padded_vs_row_fresh_T"] = stats([abs(a - b) for a, b in zip(padded_T, fresh_T)])
+    report["padded_fresh_T_vs_behavior"] = stats([abs(a - b) for a, b in zip(padded_T, behavior)])
+    report["padded_fresh_T_vs_stored_trainer"] = stats(
+        [abs(a - b) for a, b in zip(padded_T, trainer)])
+    report["padded_fresh_raw_vs_stored_trainer"] = stats(
+        [abs(a - b) for a, b in zip(padded_raw, trainer)])
+    report["row_fresh_raw_vs_stored_trainer"] = stats(
+        [abs(a - b) for a, b in zip(fresh_raw, trainer)])
+    report["row_fresh_raw_vs_behavior"] = stats(
+        [abs(a - b) for a, b in zip(fresh_raw, behavior)])
+    print("\n== which stored array does each recomputation reproduce? ==")
+    for key in ("padded_vs_row_fresh_T", "padded_fresh_T_vs_behavior",
+                "padded_fresh_T_vs_stored_trainer", "row_fresh_T_vs_behavior",
+                "row_fresh_T_vs_stored_trainer"):
+        if key in report:
+            r_ = report[key]
+        elif key == "row_fresh_T_vs_behavior":
+            r_ = report["guard_fresh_T_vs_behavior"]
+        else:
+            r_ = report["trainer_fresh_T_vs_stored_trainer"]
+        print(f"  {key:36s} mean={r_['mean']:.6f} p99={r_['p99']:.4f} "
+              f"max={r_['max']:.4f} >0.5={r_['above_0p5']}")
+    print("  temperature probe on the TRAINER side (raw reference):")
+    for key in ("row_fresh_raw_vs_stored_trainer", "padded_fresh_raw_vs_stored_trainer"):
+        r_ = report[key]
+        print(f"    {key:34s} mean={r_['mean']:.6f} p99={r_['p99']:.4f} "
+              f"max={r_['max']:.4f} >0.5={r_['above_0p5']}")
+    close = {"padded": report["padded_fresh_T_vs_stored_trainer"]["max"],
+             "row": report["trainer_fresh_T_vs_stored_trainer"]["max"]}
+    report["which_recomputation_matches_the_trainer_array"] = (
+        "padded_batch" if close["padded"] < close["row"] / 3 else
+        "per_row" if close["row"] < close["padded"] / 3 else "neither_clearly")
+    print(f"  => the trainer's stored array is reproduced by: "
+          f"{report['which_recomputation_matches_the_trainer_array']}"
+          f"  (padded max={close['padded']:.4f}, row max={close['row']:.4f})")
     d_raw = [abs(a - b) for a, b in zip(fresh_raw, behavior)]
     d_trainer = [abs(a - b) for a, b in zip(fresh_T, trainer)]
     d_trainer_raw = [abs(a - b) for a, b in zip(fresh_raw, trainer)]