Restriction-Learning-Experiment / interpretability_patchfix /train_interpretability_patchfix_pilot.py
Correct donor capture to exact generation-prefill patch sites; verify all 26 self-restorations
8cb2b5d verified | """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"] | |
| 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() | |