| """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 |
|
|
| def hook(mod, inp, o): |
| a = inp[0].float() |
| s = (a if idx is None else a[..., idx]) @ r |
| 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()} |
|
|
