| |
| """Lab 7: what else does the edit change? |
| |
| . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code |
| CUDA_VISIBLE_DEVICES=7 python labs/lab7_side_effects.py # a few minutes on one GPU |
| |
| For the base model and every edited model built in models/: |
| 1. Ordinary text. Perplexity on the first 64 x 1,024 tokens of WikiText-2 (test), and KL(base || edited): |
| how far the edited model's next-token prediction moves from the base model's, averaged over tokens. |
| 2. The agent's 8 texts. KL on its productive texts vs its looping texts. A good edit moves predictions a lot |
| where the model loops and barely anywhere else. |
| Every edited model differs from the base only in the down_proj weights of the layers its edit_note.json lists |
| (15 and 19 for most of them), so the script loads the base once and swaps those matrices in and out. |
| |
| Experiments. Any flag other than the defaults writes to results/lab7/try-NAME/ (figure included), so the main |
| results stay as they are: |
| python labs/lab7_side_effects.py --chunks 8 # quick check on 8 x 1,024 WikiText tokens |
| python labs/lab7_side_effects.py --last 36 # KL on only the last 36 tokens of the 8 texts |
| python labs/lab7_side_effects.py --models p50c capA-a6 # only some models (after Lab 8: your own edit) |
| """ |
| import argparse |
| import json |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
| import torch |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, ROOT) |
| from wsp import hf |
| from wsp.edits import NAME, load |
|
|
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--models", nargs="+", metavar="NAME", |
| help="only these edited models, named like their folders without 'qwen35-' (default: all of them)") |
| ap.add_argument("--chunks", type=int, default=64, help="how many 1,024-token WikiText-2 chunks to score (default 64)") |
| ap.add_argument("--last", type=int, default=0, |
| help="measure KL on the 8 texts over only their last N tokens (default 0: every token)") |
| ap.add_argument("--try-name", help="folder name for an experiment; results go to results/lab7/try-NAME/") |
| a = ap.parse_args() |
|
|
| t_start = time.time() |
| CHUNK = 1024 |
| parts = [] |
| if a.chunks != 64: |
| parts.append(f"chunks-{a.chunks}") |
| if a.last: |
| parts.append(f"last-{a.last}") |
| if a.models: |
| parts.append("models-" + ("-".join(a.models) if len(a.models) <= 4 else str(len(a.models)))) |
| try_name = a.try_name or "-".join(parts) |
| OUT = os.path.join(ROOT, "results", "lab7", f"try-{try_name}" if try_name else "") |
| FIG = OUT if try_name else os.path.join(ROOT, "results", "figures") |
| os.makedirs(OUT, exist_ok=True) |
| os.makedirs(FIG, exist_ok=True) |
|
|
| MODELS = os.path.join(hf.LAB, "models") |
| found = sorted(d[len("qwen35-"):] for d in os.listdir(MODELS) |
| if d.startswith("qwen35-") and os.path.exists(os.path.join(MODELS, d, "edit_note.json"))) |
| missing = [n for n in a.models or [] if n not in found] |
| if missing: |
| sys.exit(f"no edited model folder for {missing}; found {found}") |
|
|
|
|
| def edited_layers(n): |
| note = json.load(open(os.path.join(MODELS, "qwen35-" + n, "edit_note.json"))) |
| return {int(L) for L in note["spec"]["layers"]} if "spec" in note else {15, 19} |
|
|
|
|
| |
| LAYERS = sorted(set().union(*(edited_layers(n) for n in a.models or found))) |
| m, tok = hf.load("base") |
| mods = {L: hf.down_proj(m, L) for L in LAYERS} |
| base_w = {L: mods[L].weight.detach().clone() for L in LAYERS} |
| variants = {} |
| for n in a.models or found: |
| t = load(os.path.join(MODELS, "qwen35-" + n), [NAME.format(L) for L in LAYERS]) |
| variants[n] = {L: t[NAME.format(L)].to("cuda") for L in LAYERS} |
| print("models:", ["base"] + list(variants), "| down_proj swapped at layers", LAYERS) |
|
|
|
|
| def set_weights(w): |
| for L in LAYERS: |
| mods[L].weight.data.copy_(w[L]) |
|
|
|
|
| from datasets import load_dataset |
|
|
| wiki = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="test") |
| ids = tok("".join(wiki["text"]), add_special_tokens=False).input_ids |
| if a.chunks > len(ids) // CHUNK: |
| sys.exit(f"WikiText-2 test has only {len(ids) // CHUNK} chunks of {CHUNK} tokens") |
| chunks = torch.tensor(ids[:CHUNK * a.chunks]).view(a.chunks, CHUNK) |
| pairs = json.load(open(os.path.join(ROOT, "labs", "data", "astra_pairs.json"))) |
| pair_ids = [hf.encode(tok, t) for t in pairs["texts"]] |
|
|
| names = ["base"] + list(variants) |
| nll = {n: [] for n in names} |
| kl = {n: [] for n in names} |
| kl_pairs = {n: {"productive": [], "looping": []} for n in names} |
|
|
|
|
| def logprobs(x): |
| return m(x).logits[0, :-1].float().log_softmax(-1) |
|
|
|
|
| def kl_text(lb, lv): |
| k = (lb.exp() * (lb - lv)).sum(-1) |
| return (k[-a.last:] if a.last else k).mean().item() |
|
|
|
|
| with torch.no_grad(): |
| for c in range(a.chunks): |
| x = chunks[c:c + 1].cuda() |
| set_weights(base_w) |
| lb = logprobs(x) |
| tgt = x[0, 1:, None] |
| nll["base"].append(-lb.gather(1, tgt).mean().item()) |
| kl["base"].append(0.0) |
| for n, w in variants.items(): |
| set_weights(w) |
| lv = logprobs(x) |
| nll[n].append(-lv.gather(1, tgt).mean().item()) |
| kl[n].append((lb.exp() * (lb - lv)).sum(-1).mean().item()) |
| for i, x in enumerate(pair_ids): |
| kind = "looping" if i % 2 else "productive" |
| set_weights(base_w) |
| lb = logprobs(x) |
| kl_pairs["base"][kind].append(0.0) |
| for n, w in variants.items(): |
| set_weights(w) |
| lv = logprobs(x) |
| kl_pairs[n][kind].append(kl_text(lb, lv)) |
| set_weights(base_w) |
|
|
| res = {} |
| for n in names: |
| res[n] = {"wikitext_ppl": float(np.exp(np.mean(nll[n]))), |
| "wikitext_kl": float(np.mean(kl[n])), |
| "kl_productive_texts": float(np.mean(kl_pairs[n]["productive"])), |
| "kl_looping_texts": float(np.mean(kl_pairs[n]["looping"]))} |
| for n in names: |
| res[n]["ppl_change_pct"] = 100 * (res[n]["wikitext_ppl"] / res["base"]["wikitext_ppl"] - 1) |
| json.dump({"tokens": CHUNK * a.chunks, "last": a.last, "models": res}, |
| open(os.path.join(OUT, "side_effects.json"), "w"), indent=1) |
| print(f"{'model':10s} {'ppl':>8s} {'ppl +%':>7s} {'KL wiki':>9s} {'KL prod':>9s} {'KL loop':>9s}") |
| for n in names: |
| r = res[n] |
| print(f"{n:10s} {r['wikitext_ppl']:8.3f} {r['ppl_change_pct']:7.2f} {r['wikitext_kl']:9.5f} " |
| f"{r['kl_productive_texts']:9.5f} {r['kl_looping_texts']:9.5f}") |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| order = [n for n in names if n != "base"] |
| fig, ax = plt.subplots(1, 2, figsize=(14, 4.5)) |
| ax[0].bar(order, [res[n]["ppl_change_pct"] for n in order], color="#0969da") |
| ax[0].set_ylabel("WikiText-2 perplexity change vs base (%)") |
| ax[0].set_title("Damage to ordinary text") |
| xi = np.arange(len(order)) |
| ax[1].bar(xi - 0.2, [res[n]["kl_productive_texts"] for n in order], 0.4, label="productive texts", color="#1a7f37") |
| ax[1].bar(xi + 0.2, [res[n]["kl_looping_texts"] for n in order], 0.4, label="looping texts", color="#cf222e") |
| ax[1].set_xticks(xi, order) |
| ax[1].set_yscale("log") |
| ax[1].set_ylabel("KL(base || edited), mean per token") |
| ax[1].set_title("Where the edit moves predictions") |
| ax[1].legend() |
| for a_ in ax: |
| a_.tick_params(axis="x", rotation=60) |
| fig.tight_layout() |
| fig.savefig(os.path.join(FIG, "lab7_side_effects.png"), dpi=130) |
| print(f"wrote {OUT} in {time.time() - t_start:.0f} s") |
|
|
