File size: 4,826 Bytes
84bcfb1 | 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 | #!/usr/bin/env python
"""Build every model folder the labs use.
1. qwen35-p50c: the paper's released edit, rebuilt with the paper's own apply_edit.py (hash-checked).
2. results/edit_info.json: for each edited layer, the direction u, the strength alpha and the 64 columns,
read back out of the weights (Lab 1 explains how).
3. A self-check: our formula with the recovered u, alpha and columns must give the released bits.
4. One folder per variant in configs/variants.json (qwen35-<name>).
"""
import argparse
import json
import os
import subprocess
import sys
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
import torch # noqa: E402
from wsp import edits as E # noqa: E402
ap = argparse.ArgumentParser()
ap.add_argument("--base", default="/nvme/wsp-lab/models/Qwen3.5-0.8B")
ap.add_argument("--models", default="/nvme/wsp-lab/models")
ap.add_argument("--exp8", default="/nvme/wsp-lab/code/external/lm_understand-logs/exp8_math")
ap.add_argument("--variants", default=os.path.join(ROOT, "configs", "variants.json"))
ap.add_argument("--info", default=os.path.join(ROOT, "results", "edit_info.json"))
ap.add_argument("--only", nargs="*", default=None)
a = ap.parse_args()
LAYERS = (15, 19)
rel_dir = os.path.join(a.models, "qwen35-p50c")
man = json.load(open(os.path.join(a.exp8, "runs/p50c/edit/manifest.json")))
# 1. the released edit
if not os.path.exists(E.weight_path(rel_dir)):
os.makedirs(rel_dir, exist_ok=True)
tmp = E.weight_path(rel_dir) + ".tmp.safetensors"
subprocess.run([sys.executable, os.path.join(a.exp8, "code/apply_edit.py"), E.weight_path(a.base),
os.path.join(a.exp8, "runs/p50c/edit"), tmp], check=True)
os.replace(tmp, E.weight_path(rel_dir))
E.link_base_files(a.base, rel_dir)
json.dump({"name": "p50c", "about": "the paper's released edit, rebuilt with its apply_edit.py",
"edited_sha256": {t["name"]: t["edited_sha256"] for t in man["tensors"]}},
open(os.path.join(rel_dir, "edit_note.json"), "w"), indent=1)
# 2. read the edit back out of the weights
info = {"layers": {}}
for L in LAYERS:
n = E.NAME.format(L)
W0, W1 = E.load_down(a.base, L), E.load_down(rel_dir, L)
want = next(t["edited_sha256"] for t in man["tensors"] if t["name"] == n)
assert E.sha256(W1) == want, f"layer {L}: released tensor hash mismatch"
r = E.recover(W0, W1)
u = r["u"]
if u[u.abs().argmax()] < 0: # SVD's sign is arbitrary; fix one convention (labs/03 finds the loop sign)
u = -u
D = (W1.float() - W0.float())
info["layers"][str(L)] = {
"u": [round(float(x), 8) for x in u], "alpha": r["alpha"], "cols": r["cols"],
"singular_values_top5": [float(x) for x in r["singular_values"][:5]],
"rank1_share": float(r["singular_values"][0] ** 2 / (r["singular_values"] ** 2).sum()),
"relative_change": float(D.norm() / W0.float().norm()),
"changed_entries": int((W1.view(torch.int16) != W0.view(torch.int16)).sum()),
}
# 3. self-check: our formula with what we recovered vs the released bits
for alpha in (r["alpha"], 6.0):
W2 = E.projection_edit(W0, u, alpha, r["cols"])
diff = (W2.view(torch.int16) != W1.view(torch.int16))
info["layers"][str(L)][f"rebuild_alpha_{alpha:.4f}_entries_differing"] = int(diff.sum())
info["layers"][str(L)][f"rebuild_alpha_{alpha:.4f}_max_abs_diff"] = float((W2.float() - W1.float()).abs().max())
print(f"layer {L}: alpha={r['alpha']:.4f}, columns={len(r['cols'])}, rank-1 share="
f"{info['layers'][str(L)]['rank1_share']:.6f}, relative change={info['layers'][str(L)]['relative_change']:.5f}")
u15, u19 = (torch.tensor(info["layers"][str(L)]["u"]) for L in LAYERS)
info["cos_u15_u19"] = float(u15 @ u19 / u15.norm() / u19.norm())
os.makedirs(os.path.dirname(a.info), exist_ok=True)
json.dump(info, open(a.info, "w"), indent=1)
print(json.dumps({k: v for L, d in info["layers"].items() for k, v in d.items() if k.startswith("rebuild")}, indent=1))
# 4. variants
released = E.released_spec(a.info)
variants = json.load(open(a.variants))
for name, spec in variants.items():
if name.startswith("_") or (a.only and name not in a.only):
continue
out = os.path.join(a.models, f"qwen35-{name}")
note_p = os.path.join(out, "edit_note.json")
weights_spec = {k: v for k, v in spec.items() if k != "about"}
old = json.load(open(note_p)).get("spec", {}) if os.path.exists(note_p) else None
if old is not None and {k: v for k, v in old.items() if k != "about"} == weights_spec:
print("exists", name)
continue
new, desc = E.build_from_spec(a.base, spec, released)
E.write_variant(a.base, out, new, {"name": name, "spec": spec, "built": desc})
print("built", name, json.dumps(desc))
|