File size: 3,997 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 | """Helpers for the hands-on labs: load the model in transformers, read what a layer computes, or change it.
Three things you can do to a running model with a hook:
record_outputs copy what a whole layer (attention or DeltaNet mixer + MLP) adds up to, per token
record_mlp_acts copy the 3,584 neuron activations that feed the MLP's down_proj, per token
steer_down change what down_proj writes, per token, without touching any weight
"""
from __future__ import annotations
import contextlib
import json
import os
import torch
LAB = os.environ.get("WSP_LAB", "/nvme/wsp-lab")
def model_dir(name: str) -> str:
if name == "base":
return os.path.join(LAB, "models", "Qwen3.5-0.8B")
return name if name.startswith("/") else os.path.join(LAB, "models", name)
def load(name="base", dtype=torch.bfloat16, device="cuda"):
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained(model_dir(name))
m = AutoModelForCausalLM.from_pretrained(model_dir(name), dtype=dtype).to(device).eval()
return m, tok
def blocks(m):
return m.model.layers
def down_proj(m, layer: int):
return blocks(m)[layer].mlp.down_proj
def encode(tok, text: str, device="cuda"):
return tok(text, return_tensors="pt", add_special_tokens=False).input_ids.to(device)
@contextlib.contextmanager
def record_outputs(m, layers):
"""out[l] = list of [tokens, 1024] tensors: the residual stream right after block l (one per forward call)."""
out, hs = {l: [] for l in layers}, []
for l in layers:
def hook(mod, inp, o, l=l):
o = o[0] if isinstance(o, tuple) else o
out[l].append(o.detach()[0].float().cpu())
hs.append(blocks(m)[l].register_forward_hook(hook))
try:
yield out
finally:
for h in hs:
h.remove()
@contextlib.contextmanager
def record_mlp_acts(m, layers):
"""out[l] = list of [tokens, 3584] tensors: the neuron activations that down_proj reads at layer l."""
out, hs = {l: [] for l in layers}, []
for l in layers:
def hook(mod, inp, l=l):
out[l].append(inp[0].detach()[0].float().cpu())
hs.append(down_proj(m, l).register_forward_pre_hook(hook))
try:
yield out
finally:
for h in hs:
h.remove()
@contextlib.contextmanager
def steer_down(m, layer: int, u: torch.Tensor, alpha: float, cols=None, record=None):
"""Make down_proj at `layer` write W a - alpha * u * s, where s = u . (W[:, cols] a[cols]).
This is exactly what the weight edit W' = W - alpha * u u^T W[:, cols] does, computed on the fly.
s is one number per token: how hard the chosen neurons are pushing along u right now.
If `record` is a list, each call appends s (a [tokens] tensor) to it."""
lin = down_proj(m, layer)
W = lin.weight
u = (u / u.norm()).to(W.device, torch.float32)
idx = None if cols is None else torch.as_tensor(list(cols), device=W.device)
WS = (W if idx is None else W[:, idx]).float()
r = WS.T @ u # one number per chosen neuron: how much it writes along u
def hook(mod, inp, o):
a = inp[0].float()
s = (a if idx is None else a[..., idx]) @ r # [batch, tokens]
if record is not None:
record.append(s.detach()[0].cpu())
return (o.float() - alpha * s[..., None] * u).to(o.dtype)
h = lin.register_forward_hook(hook)
try:
yield
finally:
h.remove()
def released_edit(path=None):
"""u, alpha and the 64 columns per layer, as recovered from the released weights by scripts/build_models.py."""
path = path or os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "results",
"edit_info.json")
info = json.load(open(path))
return {int(l): {"u": torch.tensor(v["u"]), "alpha": v["alpha"], "cols": v["cols"]}
for l, v in info["layers"].items()}
|