Correct donor capture to exact generation-prefill patch sites; verify all 26 self-restorations

interpretability_patchfix/train_interpretability_patchfix_pilot.py

@@ -0,0 +1,144 @@
+"""Match donor capture to generation prefill, including the final layer before normalization."""
+
+import importlib.util
+import json
+import re
+from pathlib import Path
+
+import torch
+from huggingface_hub import HfApi, hf_hub_download
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+REPO = "itsalissonsilva/Restriction-Learning-Experiment"
+REVISION = "3a3cb6bef28c71ffc822f0a68645e2adc4b6aed2"
+ADAPTER_REVISION = "efc3f02e5775d3d8428b4961f1fd94aedc6ca34f"
+
+
+def main():
+    torch.manual_seed(1729)
+    source = hf_hub_download(REPO, "train_restriction_pilot.py", revision=ADAPTER_REVISION)
+    spec = importlib.util.spec_from_file_location("helper", source)
+    helper = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(helper)
+    report = json.loads(
+        Path(hf_hub_download(REPO, "interpretability/results.json", revision=REVISION)).read_text()
+    )
+    tokenizer = AutoTokenizer.from_pretrained(helper.BASE, revision=helper.REVISION)
+    tokenizer.padding_side = "left"
+    base = AutoModelForCausalLM.from_pretrained(
+        helper.BASE,
+        revision=helper.REVISION,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="eager",
+    ).to("cuda")
+    model = PeftModel.from_pretrained(
+        base, REPO, subfolder="selective", revision=ADAPTER_REVISION
+    ).eval()
+    layers = [(n, m) for n, m in model.named_modules() if re.search(r"\.layers\.\d+$", n)]
+    private = report["patch_pairs"]["private"]
+    public = report["patch_pairs"]["public"]
+
+    @torch.no_grad()
+    def generate(rows, max_tokens=80):
+        inputs = tokenizer(
+            [helper.render_prompt(tokenizer, r) for r in rows], padding=True, return_tensors="pt"
+        ).to("cuda")
+        generated = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=False,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+        texts = tokenizer.batch_decode(
+            generated[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True
+        )
+        return [helper.grading(r, t) for r, t in zip(rows, texts, strict=True)]
+
+    def capture(rows):
+        stored = {}
+        handles = []
+        for i, (_, layer) in enumerate(layers):
+
+            def hook(module, args, output, index=i):
+                h = output[0] if isinstance(output, tuple) else output
+                if h.shape[1] > 1:
+                    stored[index] = h[:, -1].detach().clone()
+
+            handles.append(layer.register_forward_hook(hook))
+        try:
+            generate(rows, 1)
+        finally:
+            for handle in handles:
+                handle.remove()
+        assert len(stored) == 26
+        return stored
+
+    hp, hc = capture(private), capture(public)
+    baseline = generate(private)
+    result = dict(
+        source_revision=REVISION,
+        source_patch_exclusion_reason="Earlier capture used ordinary forward position IDs and post-normalized final hidden state; superseded by identical generation-prefill hooks at every patch site",
+        patching=[],
+        patch_pairs=dict(
+            private=private,
+            public=public,
+            baseline_private=baseline,
+            baseline_public=generate(public),
+        ),
+        seed=1729,
+    )
+    for layer in range(26):
+        delta = hc[layer].float() - hp[layer].float()
+        noise = torch.randn_like(delta)
+        noise *= delta.norm(dim=1, keepdim=True) / noise.norm(dim=1, keepdim=True).clamp_min(1e-20)
+        for donor_kind, donor in [
+            ("public", hc[layer]),
+            ("self_restoration", hp[layer]),
+            ("random_norm_matched", hp[layer] + noise.to(hp[layer].dtype)),
+        ]:
+
+            def patch(module, args, output):
+                h = output[0] if isinstance(output, tuple) else output
+                if h.shape[1] <= 1:
+                    return output
+                replacement = h.clone()
+                replacement[:, -1] = donor
+                return (replacement, *output[1:]) if isinstance(output, tuple) else replacement
+
+            handle = layers[layer][1].register_forward_hook(patch)
+            try:
+                outputs = generate(private)
+            finally:
+                handle.remove()
+            matches = sum(o["raw"] == b["raw"] for o, b in zip(outputs, baseline, strict=True))
+            if donor_kind == "self_restoration":
+                assert matches == 8, (layer, matches)
+            result["patching"].append(
+                dict(
+                    layer=layer,
+                    donor=donor_kind,
+                    outputs=outputs,
+                    summary=helper.summarize(private, outputs),
+                    raw_matches_private_baseline=matches,
+                )
+            )
+        print(
+            json.dumps(dict(stage="patching", layer=layer, self_restoration_matches=8)), flush=True
+        )
+    assert generate(private) == baseline
+    folder = Path("patchfix")
+    folder.mkdir(exist_ok=True)
+    (folder / "patching.json").write_text(json.dumps(result, indent=2))
+    (folder / "train_interpretability_patchfix_pilot.py").write_bytes(Path(__file__).read_bytes())
+    HfApi().upload_folder(
+        repo_id=REPO,
+        folder_path=str(folder),
+        path_in_repo="interpretability_patchfix",
+        commit_message="Correct donor capture to exact generation-prefill patch sites; verify all 26 self-restorations",
+    )
+    print(json.dumps(dict(stage="complete", layers=26, self_restoration_matches=208)), flush=True)
+
+
+if __name__ == "__main__":
+    main()