File size: 6,770 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """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"):
# same Frobenius size of change as the released edit at this alpha
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
|