| """Weight edits on Qwen3.5-0.8B's MLP down projections, and model folders vLLM can load. |
| |
| The edit the paper shipped, per layer L in {15, 19}: |
| |
| W' = W - alpha * u (u^T W) restricted to a set S of columns (alpha = 6, |S| = 64) |
| |
| W is down_proj.weight (shape [1024, 3584]): column j is what MLP neuron j adds to the |
| residual stream when it fires with activation 1. u is a unit vector in the 1024-d |
| residual stream ("the loop direction"). For a selected column, u^T w_j' = (1 - alpha) u^T w_j: |
| alpha = 1 removes the column's loop component, alpha = 6 flips it and makes it 5x larger. |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| NAME = "model.language_model.layers.{}.mlp.down_proj.weight" |
| WEIGHT_FILE = "model.safetensors-00001-of-00001.safetensors" |
|
|
|
|
| def weight_path(model_dir: str | os.PathLike) -> str: |
| return os.path.join(model_dir, WEIGHT_FILE) |
|
|
|
|
| def load(model_dir, names): |
| from safetensors import safe_open |
| out = {} |
| with safe_open(weight_path(model_dir), framework="pt") as f: |
| for n in names: |
| out[n] = f.get_tensor(n) |
| return out |
|
|
|
|
| def load_down(model_dir, layer: int) -> torch.Tensor: |
| return load(model_dir, [NAME.format(layer)])[NAME.format(layer)] |
|
|
|
|
| def sha256(t: torch.Tensor) -> str: |
| return hashlib.sha256(t.contiguous().view(torch.uint16).numpy().tobytes()).hexdigest() |
|
|
|
|
| def projection_edit(W: torch.Tensor, u: torch.Tensor, alpha: float, cols=None) -> torch.Tensor: |
| """W - alpha * outer(u, mask * (u @ W)), computed in float32, returned in W's dtype.""" |
| W32 = W.float() |
| u = (u.float() / u.float().norm()) |
| g = u @ W32 |
| if cols is not None: |
| mask = torch.zeros_like(g) |
| mask[torch.as_tensor(list(cols), dtype=torch.long)] = 1.0 |
| g = g * mask |
| return (W32 - alpha * torch.outer(u, g)).to(W.dtype) |
|
|
|
|
| def recover(W0: torch.Tensor, W1: torch.Tensor): |
| """From a base and an edited matrix: the changed columns, the direction u (top left |
| singular vector of the change), the strength alpha, and the singular values.""" |
| D = W1.double() - W0.double() |
| cols = torch.nonzero(D.abs().sum(0) > 0).flatten() |
| U, S, _ = torch.linalg.svd(D[:, cols], full_matrices=False) |
| u = U[:, 0] |
| P = torch.outer(u, u @ W0.double()[:, cols]) |
| alpha = float(-(D[:, cols] * P).sum() / (P * P).sum()) |
| return {"cols": cols.tolist(), "u": u.float(), "alpha": alpha, "singular_values": S.float()} |
|
|
|
|
| def random_orthogonal(u: torch.Tensor, seed: int) -> torch.Tensor: |
| g = torch.Generator().manual_seed(seed) |
| r = torch.randn(u.shape[0], generator=g, dtype=torch.float64) |
| u64 = u.double() / u.double().norm() |
| r = r - (r @ u64) * u64 |
| return (r / r.norm()).float() |
|
|
|
|
| def link_base_files(base_dir, out_dir): |
| """Symlink every file of the base folder except the weights (config, tokenizer, index...).""" |
| base_dir, out_dir = Path(base_dir), Path(out_dir) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| for p in base_dir.iterdir(): |
| if p.name == WEIGHT_FILE or p.name.startswith(".") or p.name == "edit_note.json": |
| continue |
| q = out_dir / p.name |
| if q.is_symlink() or q.exists(): |
| q.unlink() |
| q.symlink_to(p.resolve()) |
|
|
|
|
| def write_variant(base_dir, out_dir, new_tensors: dict, note: dict): |
| """A model folder = symlinks to the base folder's files + a weight file with |
| `new_tensors` swapped in. Every other tensor is copied bit for bit.""" |
| from safetensors.torch import load_file, save_file |
| base_dir, out_dir = Path(base_dir), Path(out_dir) |
| link_base_files(base_dir, out_dir) |
| tensors = load_file(str(base_dir / WEIGHT_FILE)) |
| for n, t in new_tensors.items(): |
| assert tensors[n].shape == t.shape and tensors[n].dtype == t.dtype, n |
| tensors[n] = t.contiguous() |
| save_file(tensors, str(out_dir / WEIGHT_FILE), metadata={"format": "pt"}) |
| note = dict(note) |
| note["edited_sha256"] = {n: sha256(t) for n, t in new_tensors.items()} |
| (out_dir / "edit_note.json").write_text(json.dumps(note, indent=1)) |
| return out_dir |
|
|
|
|
| def released_spec(edit_info_path): |
| """The released edit, as recovered by labs/01: per layer, u, alpha and columns.""" |
| d = json.loads(Path(edit_info_path).read_text()) |
| return {int(l): {"u": torch.tensor(v["u"]), "alpha": v["alpha"], "cols": v["cols"]} |
| for l, v in d["layers"].items()} |
|
|
|
|
| def build_from_spec(base_dir, spec: dict, released: dict): |
| """spec: {"layers": {"15": {"alpha": 6, "dir": "released" | "random_orth:SEED" | [floats], |
| "cols": "released" | "all" | "random:SEED" | [ints], |
| "norm_match": false}, ...}} |
| Returns {tensor name: new bf16 tensor} and a description.""" |
| new, desc = {}, {} |
| for l_str, s in spec["layers"].items(): |
| l = int(l_str) |
| W = load_down(base_dir, l) |
| rel = released.get(l) |
| d = s.get("dir", "released") |
| c = s.get("cols", "released") |
| if rel is None and (isinstance(d, str) or (isinstance(c, str) and c != "all") or s.get("norm_match")): |
| raise ValueError(f"layer {l}: the released edit does not touch this layer, " |
| "so give an explicit dir (a list of floats) and cols (a list or 'all')") |
| if d == "released": |
| u = rel["u"] |
| elif isinstance(d, str) and d.startswith("random_orth:"): |
| u = random_orthogonal(rel["u"], int(d.split(":")[1])) |
| else: |
| u = torch.tensor(d, dtype=torch.float32) |
| u = u / u.norm() |
| if c == "released": |
| cols = rel["cols"] |
| elif c == "all": |
| cols = None |
| elif isinstance(c, str) and c.startswith("random:"): |
| rng = np.random.default_rng(int(c.split(":")[1])) |
| pool = np.setdiff1d(np.arange(W.shape[1]), np.array(rel["cols"])) |
| cols = sorted(rng.choice(pool, size=len(rel["cols"]), replace=False).tolist()) |
| else: |
| cols = list(c) |
| alpha = float(s["alpha"]) |
| if s.get("norm_match"): |
| |
| W32 = W.float() |
| ref = (rel["u"] / rel["u"].norm()) @ W32[:, rel["cols"]] |
| mine = u @ (W32 if cols is None else W32[:, cols]) |
| alpha = alpha * float(ref.norm() / mine.norm()) |
| W1 = projection_edit(W, u, alpha, cols) |
| new[NAME.format(l)] = W1 |
| desc[l] = {"alpha": alpha, "n_cols": W.shape[1] if cols is None else len(cols), |
| "dir": d if isinstance(d, str) else "custom", |
| "relative_change": float((W1.float() - W.float()).norm() / W.float().norm())} |
| return new, desc |
|
|
