File size: 5,569 Bytes
1a890ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """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()
|