File size: 7,431 Bytes
0985f91 | 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 | # keys, current residuals, and the gradient-optimised target vectors ("z*") that ROME/MEMIT use.
# z* is the only place a baseline backprops through the model. our own methods get it only as an oracle.
import os, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import *
# a few throwaway lead-ins so keys/values aren't glued to "subject is the first word"
PREFIXES = ["", "The weather was nice. ", "I think that ", "According to the report, ", "Yesterday someone said: ", "In other news, "]
def req(f):
"""what an editor is allowed to see. no paraphrases, no hops, no neighbours."""
return dict(id=f["id"], subject=f["subject"], target=f["target"], prompt=f["direct"], statement=f["statement"], rel=f["rel"])
def _hid(o): return o[0] if isinstance(o, tuple) else o
def POS_SUBJ(tok, p, r): return subj_pos(tok, p, r["subject"]) # last token of the subject
def POS_LAST(tok, p, r): return len(enc(tok, p)) - 1 # last token of the prompt
def keys_at(model, tok, reqs, layers, prefixes=PREFIXES, bs=24, pos_fn=POS_SUBJ):
"""mean over prefixes of the mlp key (input of dense_4h_to_h) at the last subject token. dict l -> [N, 8192]"""
acc = None
for pre in prefixes:
P = [pre + r["prompt"] for r in reqs]; pos = [pos_fn(tok, p, r) for p, r in zip(P, reqs)]
got = capture(model, tok, P, pos, {l: (W_out(model, l), "in") for l in layers}, bs=bs)
acc = got if acc is None else {l: acc[l] + got[l] for l in layers}
return {l: acc[l] / len(prefixes) for l in layers}
def resid_at(model, tok, reqs, L, prefixes=PREFIXES, bs=24, pos_fn=POS_SUBJ):
"""mean over prefixes of block-L output at the last subject token. [N, 2048]"""
acc = 0
for pre in prefixes:
P = [pre + r["prompt"] for r in reqs]; pos = [pos_fn(tok, p, r) for p, r in zip(P, reqs)]
acc = acc + capture(model, tok, P, pos, {"h": (blocks(model)[L], "out")}, bs=bs)["h"]
return acc / len(prefixes)
def compute_z(model, tok, reqs, L, steps=20, lr=0.5, clamp=0.75, klf=0.0625, wd=0.5, scale=256.0):
"""MEMIT's compute_z, several facts at once (each fact has its own delta, the losses don't mix).
returns absolute targets z [G, 2048] for block-L output at the last subject token."""
seqs, own, spos, tinfo, bare = [], [], [], [], []
for g, r in enumerate(reqs):
for k, pre in enumerate(PREFIXES):
pi, ti = split_target(tok, pre + r["prompt"], r["target"])
if k == 0: bare.append(len(seqs))
seqs.append(pi + ti[:-1]); own.append(g); spos.append(subj_pos(tok, pre + r["prompt"], r["subject"])); tinfo.append((len(pi) - 1, ti))
klp = r["subject"] + " is a"
seqs.append(enc(tok, klp)); own.append(g); spos.append(subj_pos(tok, klp, r["subject"])); tinfo.append(None)
x, m = pad(seqs); n = len(seqs); rows = torch.arange(n, device=DEV)
spos_t = torch.tensor(spos, device=DEV); own_t = torch.tensor(own, device=DEV)
G = len(reqs); d = model.config.hidden_size
delta = torch.zeros(G, d, device=DEV, requires_grad=True); opt = torch.optim.Adam([delta], lr=lr)
st = dict(init=None)
def hook(mod, inp, out):
h = _hid(out)
if st["init"] is None: st["init"] = h[rows, spos_t].detach().float()
h = h.clone(); h[rows, spos_t] = h[rows, spos_t] + delta[own_t].to(h.dtype)
return (h,) + tuple(out[1:]) if isinstance(out, tuple) else h
# flat index lists for the target-token log-probs
r_i, p_i, t_i, w_i, g_i = [], [], [], [], []
for s, info in enumerate(tinfo):
if info is None: continue
a, ti = info
for j, t in enumerate(ti):
r_i.append(s); p_i.append(a + j); t_i.append(t); w_i.append(1.0 / (len(ti) * len(PREFIXES))); g_i.append(own[s])
r_i, p_i, t_i, g_i = (torch.tensor(v, device=DEV) for v in (r_i, p_i, t_i, g_i)); w_i = torch.tensor(w_i, device=DEV)
kl_rows = torch.tensor([s for s, info in enumerate(tinfo) if info is None], device=DEV)
kl_pos = torch.tensor([len(seqs[s]) - 1 for s in kl_rows.tolist()], device=DEV)
kl_init, init_norm = None, None
with torch.enable_grad(), hooks([(blocks(model)[L], hook, "fwd")]):
for step in range(steps):
logits = model(input_ids=x, attention_mask=m).logits
lp = torch.log_softmax(logits[r_i, p_i].float(), -1)
nll = torch.zeros(G, device=DEV).index_add(0, g_i, -lp.gather(1, t_i[:, None])[:, 0] * w_i)
klp = torch.log_softmax(logits[kl_rows, kl_pos].float(), -1)
if kl_init is None:
kl_init = klp.detach(); init_norm = st["init"][torch.tensor(bare, device=DEV)].norm(dim=1)
kl = (kl_init.exp() * (kl_init - klp)).sum(-1)
loss = (nll + klf * kl + wd * (delta.norm(dim=1) / init_norm) ** 2).sum()
opt.zero_grad(); (loss * scale).backward()
if not torch.isfinite(delta.grad).all():
delta.grad = torch.nan_to_num(delta.grad, nan=0.0, posinf=0.0, neginf=0.0)
opt.step()
with torch.no_grad():
nrm = delta.norm(dim=1, keepdim=True); mx = (clamp * init_norm)[:, None]
delta.mul_(torch.clamp(mx / nrm.clamp_min(1e-8), max=1.0))
z = st["init"][torch.tensor(bare, device=DEV)] + delta.detach()
return z, nll.detach()
def z_cached(model, tok, reqs, L, group=8, tag="", **kw):
"""z* per fact id, cached on disk. MUST be called on the unedited model (the cache assumes it)."""
path = os.path.join(CK, f"z_L{L}{tag}.pt")
cache = torch.load(path) if os.path.exists(path) else {}
todo = [r for r in reqs if r["id"] not in cache]
t0 = time.time()
for a in range(0, len(todo), group):
g = todo[a:a + group]
while True:
try:
z, nll = compute_z(model, tok, g, L, **kw); break
except torch.OutOfMemoryError:
free(); assert len(g) > 1; g = g[:len(g) // 2]; print("[oom] z group ->", len(g), flush=True)
if len(g) < len(todo[a:a + group]): # finish the rest of this slice one small group at a time
rest = todo[a + len(g):a + group]
z2 = [compute_z(model, tok, rest[i:i + len(g)], L, **kw) for i in range(0, len(rest), len(g))]
z = torch.cat([z] + [q[0] for q in z2]); nll = torch.cat([nll] + [q[1] for q in z2]); g = todo[a:a + group]
for r, zz, nl in zip(g, z, nll): cache[r["id"]] = dict(z=zz.cpu(), nll=nl.item())
if (a // group) % 10 == 0:
torch.save(cache, path); print(f" z* {a + len(g)}/{len(todo)} nll~{nll.mean().item():.3f} {time.time() - t0:.0f}s", flush=True)
if todo: torch.save(cache, path)
return torch.stack([cache[r["id"]]["z"] for r in reqs]).to(DEV), [cache[r["id"]]["nll"] for r in reqs]
_CINV = {}
def cov_inv(l, kind="k", ridge=0.0):
"""inverse of the cached second-moment matrix for layer l (float64, on gpu). kind 'k' = mlp keys, 'x' = mlp inputs."""
key = (l, kind, ridge)
if key not in _CINV: # parked on cpu (64GB of ram, 8 of vram), shipped to the gpu per use
C = torch.load(os.path.join(CK, f"cov_{kind}_L{l}.pt"))["C"].to(DEV).double()
if ridge: C = C + ridge * C.diagonal().mean() * torch.eye(C.shape[0], device=DEV, dtype=torch.float64)
_CINV[key] = torch.linalg.inv(C).cpu()
del C; free()
return _CINV[key].to(DEV)
def drop_cov_cache():
_CINV.clear(); free()
|