| |
| """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 |
|
|
| from wsp import edits as E |
|
|
| 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"))) |
|
|
| |
| 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) |
|
|
| |
| 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: |
| 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()), |
| } |
| |
| 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)) |
|
|
| |
| 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)) |
|
|
