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

diag/diagnose_padding_trigger.py

@@ -119,6 +119,47 @@ def main() -> None:
         return p
 
     pos_stored = positions_for(attn)
+
+    # ---------------------------------------------------------------------
+    # Per-row delta for the stored configuration, for EVERY row.
+    #
+    # Discriminator between two mechanisms: if rows AFTER the first are affected
+    # (row 1 sees row 0, row 3 sees rows 0-2) then the trainer's forward couples
+    # rows -- a concatenation/packing-style defect. If only the LONGEST row is
+    # affected it is a length/width effect. The two have different fixes.
+    # ---------------------------------------------------------------------
+    with torch.inference_mode():
+        padded_all = model(input_ids=ids.cuda(), attention_mask=attn.cuda(),
+                           position_ids=pos_stored.cuda(), use_cache=False).logits
+    per_row: dict[str, Any] = {}
+    for i in range(ids.shape[0]):
+        n_i = int(lengths[i])
+        with torch.inference_mode():
+            ref_i = model(input_ids=ids[i:i + 1, :n_i].cuda(), use_cache=False).logits[0]
+        per_row[str(i)] = max_logit_delta(padded_all[i, :n_i], ref_i)
+        del ref_i
+        torch.cuda.empty_cache()
+        print("TRIGGER_PER_ROW=" + json.dumps(
+            {"row": i, "length": n_i, **per_row[str(i)]}, sort_keys=True), flush=True)
+    del padded_all
+    torch.cuda.empty_cache()
+    report["per_row_stored_config"] = per_row
+
+    # Structure of the padded region only: is it a single constant pad value? No
+    # token ids are printed, just counts.
+    pad_region = []
+    for i in range(ids.shape[0]):
+        L = int(lengths[i])
+        if L >= width:
+            pad_region.append({"row": i, "padded_tokens": 0})
+            continue
+        seg = ids[i, L:]
+        uniq = int(torch.unique(seg).numel())
+        pad_region.append({"row": i, "padded_tokens": int(seg.numel()),
+                           "distinct_pad_values": uniq, "all_equal": bool(uniq == 1)})
+    report["pad_region_structure"] = pad_region
+    print("TRIGGER_PAD_REGION=" + json.dumps(pad_region), flush=True)
+
     prefix = torch.zeros_like(attn)
     for i, L in enumerate(lengths):
         prefix[i, :L] = 1