File size: 3,368 Bytes
8505f8e e6dc020 8505f8e | 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 | """Norm-matched activation injection (activation-oracle formula) + residual read hook.
resid[b, pos] += unit(v[b]) * ‖resid[b, pos]‖ * coeff (v detached; grad flows via resid)
"""
import contextlib
import torch
def get_layer(model, layer: int):
"""The decoder block at `layer`, unwrapping DDP + PEFT."""
m = model.module if hasattr(model, "module") else model
base = m.get_base_model() if hasattr(m, "get_base_model") else m
return base.model.layers[layer]
def make_inject_hook(vecs, positions, coeff, device, dtype):
"""vecs: list of [1, d] unit-ish directions (one per batch row). positions: list[list[int]]."""
normed = [torch.nn.functional.normalize(v.to(device, dtype), dim=-1) for v in vecs]
def hook(_module, _inp, out):
h = out[0] if isinstance(out, tuple) else out
if h.shape[1] <= 1: # decode step (KV-cache): marker already injected at prefill
return out
if h.shape[0] != len(normed):
raise RuntimeError(f"inject batch {h.shape[0]} != {len(normed)} vectors")
for b, pos in enumerate(positions):
p = torch.tensor(pos, device=h.device)
base = h[b, p] # [k, d]
scale = base.norm(dim=-1, keepdim=True) * coeff
h[b, p] = base + (normed[b] * scale).to(h.dtype).detach()
return out
return hook
def make_packed_inject_hook(vecs, rows, cols, coeff, device, dtype):
"""Packed-block variant: vecs [K, d]; direction j injected at (rows[j], cols[j]) — several
markers per batch row, one direction per marker. Norm-matched formula identical to
make_inject_hook. Training-forward only (seq_len == pack_len > 1), so the decode-step guard
below never triggers; kept for symmetry."""
normed = torch.nn.functional.normalize(vecs.to(device, dtype), dim=-1) # [K, d]
rows, cols = rows.to(device), cols.to(device)
def hook(_module, _inp, out):
h = out[0] if isinstance(out, tuple) else out
if h.shape[1] <= 1: # decode step (KV-cache): marker already injected at prefill
return out
base = h[rows, cols] # [K, d]
scale = base.norm(dim=-1, keepdim=True) * coeff
h[rows, cols] = base + (normed * scale).to(h.dtype).detach()
return out
return hook
@contextlib.contextmanager
def hooked(module, hook):
handle = module.register_forward_hook(hook)
try:
yield
finally:
handle.remove()
class _Stop(Exception):
pass
@torch.no_grad()
def read_resid(model, layer, batch, pool="mean"):
"""Layer-`layer` residual for a tokenized batch. pool: 'mean'|'last'|'all'. No injection, base model."""
captured = {}
def cap(_m, _i, out):
captured["h"] = (out[0] if isinstance(out, tuple) else out).float()
raise _Stop
h = None
handle = get_layer(model, layer).register_forward_hook(cap)
try:
model(**batch)
except _Stop:
h = captured["h"]
finally:
handle.remove()
mask = batch["attention_mask"].bool()
if pool == "all":
return h, mask
if pool == "last":
idx = mask.sum(1) - 1
return h[torch.arange(h.shape[0]), idx]
summed = (h * mask.unsqueeze(-1)).sum(1)
return summed / mask.sum(1, keepdim=True).clamp(min=1)
|