| |
| """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 |
| from wsp.edits import recover, weight_path |
|
|
| 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) |
|
|
| |
| 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 |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| 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) |
|
|
