File size: 4,516 Bytes
bb102c0 | 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 | #!/usr/bin/env python
"""Lab 1: find the edit by comparing two weight files. No GPU needed.
. /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
python labs/lab1_diff.py # under a minute on CPU
python labs/lab1_diff.py --edited qwen35-all2 # diff any other model in /nvme/wsp-lab/models instead;
# writes to results/lab1/try-<name>, main results untouched
1. Compare every tensor of the base model and the released model. Which ones changed?
2. In each changed matrix: how many entries and which columns (neurons) changed?
3. Take the change apart with an SVD: is it rank one? If so, the left vector is the direction u, and a
least-squares fit gives the strength alpha in W' = W - alpha * u (u^T W[:, cols]).
"""
import argparse
import json
import os
import sys
import numpy as np
import torch
from safetensors import safe_open
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from wsp import hf # noqa: E402
from wsp.edits import recover, weight_path # noqa: E402
ap = argparse.ArgumentParser()
ap.add_argument("--edited", default="qwen35-p50c", help="model folder (in /nvme/wsp-lab/models) to compare with base")
a = ap.parse_args()
TRY = None if a.edited == "qwen35-p50c" else "try-" + os.path.basename(a.edited.rstrip("/"))
OUT = os.path.join(ROOT, "results", "lab1", *([TRY] if TRY else []))
FIG = OUT if TRY else os.path.join(ROOT, "results", "figures")
os.makedirs(OUT, exist_ok=True)
os.makedirs(FIG, exist_ok=True)
base_dir, edit_dir = hf.model_dir("base"), hf.model_dir(a.edited)
# 1. every tensor, bit for bit
changed, n_total, n_params = [], 0, 0
with safe_open(weight_path(base_dir), "pt") as fb, safe_open(weight_path(edit_dir), "pt") as fe:
for name in fb.keys():
tb, te = fb.get_tensor(name), fe.get_tensor(name)
n_total += 1
n_params += tb.numel()
if not torch.equal(tb, te):
changed.append((name, tb, te))
print(f"{len(changed)} of {n_total} tensors differ ({n_params / 1e6:.0f}M numbers in all):")
res = {"tensors_total": n_total, "params_total": n_params, "changed": {}}
col_norms, svals = {}, {}
for name, W0, W1 in changed:
L = int(name.split("layers.")[1].split(".")[0])
D = W1.float() - W0.float()
entries = int((W1 != W0).sum())
cols = torch.nonzero(D.abs().sum(0) > 0).flatten().tolist()
r = recover(W0, W1)
s = r["singular_values"]
col_norms[L] = D.norm(dim=0).numpy()
svals[L] = s.numpy()
res["changed"][name] = {"shape": list(W0.shape), "entries_changed": entries, "n_cols": len(cols),
"cols": cols, "rank1_share": float(s[0] ** 2 / (s ** 2).sum()), "alpha": r["alpha"],
"relative_change": float(D.norm() / W0.float().norm())}
print(f" {name} shape {list(W0.shape)}")
print(f" entries changed: {entries:,} of {W0.numel():,}, all inside {len(cols)} columns (neurons)")
print(f" singular values of the change: {[round(float(x), 4) for x in s[:4]]} ... "
f"rank-one share {res['changed'][name]['rank1_share']:.6f}")
print(f" fitted alpha = {r['alpha']:.4f}; the matrix moved by {res['changed'][name]['relative_change']:.1%}")
res["changed"][name]["u"] = r["u"].tolist()
names = list(res["changed"])
if len(names) == 2:
u1, u2 = (torch.tensor(res["changed"][n]["u"]) for n in names)
res["abs_cos_between_the_two_u"] = abs(float(u1 @ u2))
print(f"|cos| between the two directions: {res['abs_cos_between_the_two_u']:.3f} "
f"(1 = same direction, 0 = unrelated; the SVD sign is arbitrary, Lab 3 fixes it)")
json.dump(res, open(os.path.join(OUT, "diff.json"), "w"), indent=1)
import matplotlib # noqa: E402
matplotlib.use("Agg")
import matplotlib.pyplot as plt # noqa: E402
layers = sorted(col_norms)
fig, ax = plt.subplots(2, len(layers), figsize=(6.5 * len(layers), 6.5), squeeze=False)
for k, L in enumerate(layers):
ax[0, k].plot(col_norms[L], color="#cf222e", lw=0.8)
ax[0, k].set_title(f"layer {L}: size of the change in each of 3,584 columns")
ax[0, k].set_xlabel("neuron (column of down_proj)")
ax[1, k].bar(np.arange(1, 11), svals[L][:10], color="#0969da")
ax[1, k].set_yscale("log")
ax[1, k].set_title(f"layer {L}: singular values of the change (log scale)")
ax[1, k].set_xlabel("rank")
fig.tight_layout()
fig.savefig(os.path.join(FIG, "lab1_diff.png"), dpi=130)
print("wrote", OUT)
|