add padding-trigger ablation for the 4096-cap parity bug

diag/diagnose_padding_trigger.py

@@ -0,0 +1,165 @@
+"""Find which property of the capture's own arrays makes a padded forward differ.
+
+Confirmed on `failure-zusi05xj`: a padded-batch eager forward reproduces the trainer's
+stored logprobs BIT-EXACTLY (max 0.0) while a per-row forward reproduces the engine's
+(max 0.33); the two differ by 1.7597 max on 150 positions. A synthetic sweep that copied
+the capture's geometry (lengths 506/224/4096/644, longest at index 2, pad id 0, train
+mode, position_ids = cumsum-1 with pads forced to 1) reports max|delta logit| = 0.00000
+for every variant. So the trigger is a property of the capture's DATA, not its shape.
+
+This ablates the real arrays one input at a time and reports the content-independent
+statistic `max |delta logit|` between the padded forward and the per-row forward:
+
+  base            ids + attention_mask exactly as stored, position_ids = cumsum-1
+  base_nopos      same, but no position_ids passed
+  pad_zero        padded region of every row replaced by token id 0
+  attn_prefix     attention_mask rebuilt as a plain prefix-ones mask from its own sum
+  attn_all_ones   attention_mask replaced by all ones
+  long_only_pad   single-row batch of the long row, padded to the stored width
+  long_only       the long row alone at its own length (no padding at all)
+
+If `attn_prefix` or `pad_zero` collapses the delta, that names the trigger.
+
+Read-only and structure-only: it prints mask shape statistics and deltas, never prompt
+text or token ids.
+
+  CUDA_VISIBLE_DEVICES=0 bash $SRC/experiments/runai/python-b200-host.sh \
+    /tmp/diagnose_padding_trigger.py "$CAPDIR" --out /tmp/pad_trigger.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any
+
+
+def max_logit_delta(a, b, chunk: int = 256) -> dict[str, Any]:
+    """max |a - b| over the vocabulary, per position, in chunks."""
+    n = int(min(a.shape[0], b.shape[0]))
+    worst, worst_pos = 0.0, -1
+    means: list[float] = []
+    for lo in range(0, n, chunk):
+        hi = min(lo + chunk, n)
+        d = (a[lo:hi].float() - b[lo:hi].float()).abs()
+        means.append(float(d.mean()))
+        m = float(d.max())
+        if m > worst:
+            worst, worst_pos = m, lo + int(d.max(dim=1).values.argmax())
+        del d
+    return {"positions": n, "max_abs_dlogit": worst, "worst_position": worst_pos,
+            "mean_abs_dlogit": sum(means) / len(means) if means else 0.0}
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("capture", type=Path)
+    ap.add_argument("--out", type=Path, default=None)
+    ap.add_argument("--dtype", choices=("bfloat16", "float32"), default="bfloat16")
+    ap.add_argument("--row", type=int, default=-1, help="row to compare; -1 = longest")
+    args = ap.parse_args()
+
+    import torch
+    from transformers import AutoModelForCausalLM
+
+    meta = json.loads((args.capture / "metadata.json").read_text())
+    batch = torch.load(args.capture / "batch.pt", map_location="cpu", weights_only=True)
+    ids = batch["stu_input_ids"].long()
+    attn = batch["stu_attn_mask"].long()
+    lengths = [int(v) for v in attn.sum(-1).tolist()]
+    width = int(ids.shape[1])
+    r = args.row if args.row >= 0 else max(range(len(lengths)), key=lambda i: lengths[i])
+
+    report: dict[str, Any] = {
+        "capture": str(args.capture), "rows": int(ids.shape[0]), "width": width,
+        "lengths": lengths, "row": r,
+        "model_training": meta.get("model_training"),
+        "attention_backend": meta.get("attention_backend"),
+    }
+    print("TRIGGER_META=" + json.dumps(
+        {k: report[k] for k in ("rows", "width", "lengths", "row", "model_training")},
+        sort_keys=True), flush=True)
+
+    # Structural description of the masks only: counts and whether the ones form a
+    # plain prefix. No token ids and no text leave the machine.
+    struct = []
+    for i in range(ids.shape[0]):
+        a = attn[i]
+        ones = int(a.sum())
+        first_zero = int((a == 0).nonzero()[0]) if bool((a == 0).any()) else -1
+        struct.append({"row": i, "ones": ones, "first_zero": first_zero,
+                       "is_prefix_ones": bool(int(a[:ones].sum()) == ones)})
+    report["mask_structure"] = struct
+    print("TRIGGER_MASK_STRUCTURE=" + json.dumps(struct), flush=True)
+
+    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")))
+
+    n_r = int(lengths[r])
+    lone_ids = ids[r:r + 1, :n_r].clone()
+
+    with torch.inference_mode():
+        ref = model(input_ids=lone_ids.cuda(), use_cache=False).logits[0]
+
+    def forward_row(x, a, positions):
+        with torch.inference_mode():
+            out = model(input_ids=x.cuda(),
+                        attention_mask=(a.cuda() if a is not None else None),
+                        position_ids=(positions.cuda() if positions is not None else None),
+                        use_cache=False).logits
+        return out[r] if out.dim() == 3 else out[0]
+
+    def positions_for(a):
+        p = a.long().cumsum(-1) - 1
+        p.masked_fill_(a == 0, 1)
+        return p
+
+    pos_stored = positions_for(attn)
+    prefix = torch.zeros_like(attn)
+    for i, L in enumerate(lengths):
+        prefix[i, :L] = 1
+    pad_ids = torch.zeros((1, width), dtype=torch.long)
+    pad_attn = torch.zeros((1, width), dtype=torch.long)
+    pad_ids[0, :n_r] = ids[r, :n_r]
+    pad_attn[0, :n_r] = 1
+
+    plans: list[tuple[str, Any, Any, Any]] = [
+        ("base", ids, attn, pos_stored),
+        ("base_nopos", ids, attn, None),
+        ("pad_zero", torch.where(attn.bool(), ids, torch.zeros_like(ids)), attn, pos_stored),
+        ("attn_prefix", ids, prefix, None),
+        ("attn_all_ones", ids, torch.ones_like(attn), None),
+        ("long_only_pad", pad_ids, pad_attn, None),
+        ("long_only", lone_ids, torch.ones_like(lone_ids), None),
+    ]
+    variants: dict[str, dict[str, Any]] = {}
+    for name, x, a, p in plans:
+        got = forward_row(x, a, p)
+        if name == "long_only":
+            variants[name] = max_logit_delta(got, ref)
+        else:
+            variants[name] = max_logit_delta(got, ref)
+        del got
+        torch.cuda.empty_cache()
+        print("TRIGGER_VARIANT=" + json.dumps({"variant": name, **variants[name]},
+                                             sort_keys=True), flush=True)
+
+    report["variants"] = variants
+    ranked = sorted((v["max_abs_dlogit"], k) for k, v in variants.items())
+    report["smallest_delta_variant"] = ranked[0][1] if ranked else None
+    print("TRIGGER_RANKED=" + json.dumps(
+        {"smallest": ranked[0][1] if ranked else None,
+         "deltas": {k: round(v["max_abs_dlogit"], 6) for k, v in variants.items()}},
+        sort_keys=True), flush=True)
+    if args.out:
+        args.out.write_text(json.dumps(report, indent=2, sort_keys=True, default=str) + "\n")
+        print("TRIGGER_OUT=" + str(args.out), flush=True)
+    print("TRIGGER_DONE=1", flush=True)
+
+
+if __name__ == "__main__":
+    main()