Cleanup init commit
This commit is contained in:
Extraltodeus
@@ -0,0 +1,380 @@
|
||||
import itertools
|
||||
import threading
|
||||
|
||||
import torch
|
||||
|
||||
# Default layer slice for a new rule, as fractions of the model's layer count:
|
||||
# e.g. 56 layers -> from int(56*3/5)=33 to int(56*4/5)=44.
|
||||
DEFAULT_LAYERS_FRAC_LO = 3 / 5
|
||||
DEFAULT_LAYERS_FRAC_HI = 4 / 5
|
||||
|
||||
|
||||
def default_layers(n_layers):
|
||||
lo = int(n_layers * DEFAULT_LAYERS_FRAC_LO)
|
||||
hi = min(int(n_layers * DEFAULT_LAYERS_FRAC_HI), n_layers - 1)
|
||||
return list(range(lo, hi + 1))
|
||||
|
||||
|
||||
def effective_coeffs(mode, factor, g):
|
||||
"""Effective coefficients ``(alpha, beta)`` of a rule's effect under the
|
||||
global multiplier ``g``: ``delta = alpha·(v̂_A·h)·v̂_A + beta·(v̂_A·h)·v̂_B``
|
||||
(``beta = 0`` in scale mode).
|
||||
|
||||
Saturates the over-correction: at g=1 the effect is exactly that of the
|
||||
factor; beyond it, it converges to full removal of the component (or to the
|
||||
explicitly requested inversion if factor < 0) WITHOUT overshooting it.
|
||||
Without this bound, g·(factor-1) < -1 makes the component negative — a
|
||||
chaotic anti-direction (measured: "zap Paris" at scale 4 → "Paris Paris
|
||||
Paris..." in a loop).
|
||||
"""
|
||||
if mode == "scale":
|
||||
alpha = g * (factor - 1.0)
|
||||
if factor < 1.0:
|
||||
# final component 1+alpha bounded to min(factor, 0)
|
||||
alpha = max(alpha, min(factor, 0.0) - 1.0)
|
||||
return alpha, 0.0
|
||||
# replace: saturated removal of A (never anti-A), addition of B linear in g
|
||||
return -min(g, 1.0), g * factor
|
||||
|
||||
|
||||
def abliteration_direction(weight_u, rule):
|
||||
"""Residual directions of a rule for the abliteration mode (global
|
||||
pure-weight edit).
|
||||
|
||||
``weight_u``: the un-embedding matrix W_U (lm_head), [vocab, d_model]. The
|
||||
directions live in the residual space (the basis W_U reads). Returns
|
||||
``(v_a, v_b)`` (float, CPU, normalized); ``v_b`` is None in scale mode. The
|
||||
effect applied to each residual write ``h`` is
|
||||
``h += alpha·(v̂_A·h)·v̂_A + beta·(v̂_A·h)·v̂_B`` with ``(alpha, beta)`` given
|
||||
by :func:`effective_coeffs` (which folds in the global scale).
|
||||
"""
|
||||
v_a = weight_u[rule["token_id"]].detach().float().cpu()
|
||||
v_a = v_a / v_a.norm().clamp_min(1e-8)
|
||||
v_b = None
|
||||
if rule["mode"] != "scale":
|
||||
v_b = weight_u[rule["replacement_id"]].detach().float().cpu()
|
||||
v_b = v_b / v_b.norm().clamp_min(1e-8)
|
||||
return v_a, v_b
|
||||
|
||||
|
||||
# Rule application modes:
|
||||
# standard — layer-by-layer residual steering (hook on the output of the
|
||||
# chosen layers). The most expressive live, but no layer write
|
||||
# carries the "skip": not faithfully exportable.
|
||||
# readthrough — change of basis of the downstream READS (cf. core/rebase):
|
||||
# the preview hooks the RMSNorm output with the same transform
|
||||
# as the bake → preview = exported checkpoint.
|
||||
# exact — readthrough + counter-transform of the downstream writes
|
||||
# (reproduces a hook applied exactly once; regularized inverse
|
||||
# near a full zap → reserved for soft factors).
|
||||
# abliteration — global W_U projection on every residual write (embed + all
|
||||
# block outputs); bake = the same projections on the writes.
|
||||
# The pure-weights path for architectures the rebase does not
|
||||
# support (write norms, Gemma style). Faithful for full
|
||||
# zaps/replaces; a rule's layers are ignored (global).
|
||||
MODES = ("standard", "readthrough", "exact", "abliteration")
|
||||
|
||||
|
||||
class Interventions:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._counter = itertools.count(1)
|
||||
self._rules = []
|
||||
self._handles = []
|
||||
self._scale = 1.0
|
||||
self._mode = "standard"
|
||||
|
||||
@property
|
||||
def active(self):
|
||||
return bool(self._rules)
|
||||
|
||||
@property
|
||||
def global_scale(self):
|
||||
return self._scale
|
||||
|
||||
@property
|
||||
def mode(self):
|
||||
return self._mode
|
||||
|
||||
def set_scale(self, scale):
|
||||
with self._lock:
|
||||
self._scale = float(scale)
|
||||
return self._scale
|
||||
|
||||
def set_mode(self, mode):
|
||||
if mode not in MODES:
|
||||
raise ValueError(f"unknown intervention mode: {mode}")
|
||||
with self._lock:
|
||||
self._mode = mode
|
||||
return self._mode
|
||||
|
||||
def rules_full(self):
|
||||
return list(self._rules)
|
||||
|
||||
def active_rules_full(self):
|
||||
"""Full rules (with directions) actually applied — for export: a disabled
|
||||
rule or one without layers must not be baked."""
|
||||
return list(self._active_rules())
|
||||
|
||||
def _active_rules(self):
|
||||
"""Rules actually applied: non-empty layers AND not disabled. The
|
||||
`enabled` flag lets you switch a rule off without losing its layer
|
||||
selection (the "layers=[]" gesture stays possible but clears the selection)."""
|
||||
return [r for r in self._rules if r["layers"] and r.get("enabled", True)]
|
||||
|
||||
def summary(self):
|
||||
return [
|
||||
{
|
||||
"id": rule["id"],
|
||||
"token_id": rule["token_id"],
|
||||
"token": rule["token"],
|
||||
"mode": rule["mode"],
|
||||
"factor": rule["factor"],
|
||||
"replacement_id": rule["replacement_id"],
|
||||
"replacement": rule["replacement"],
|
||||
"layers": rule["layers"],
|
||||
"enabled": rule.get("enabled", True),
|
||||
}
|
||||
for rule in self._rules
|
||||
]
|
||||
|
||||
def _direction(self, lens, weight, token_id, layers):
|
||||
row = weight[token_id].float()
|
||||
dirs = {}
|
||||
for layer in layers:
|
||||
J = lens.jacobians.get(layer)
|
||||
if J is None:
|
||||
# layer not fitted by the lens: direct logit lens (J = I),
|
||||
# a good approximation near the output
|
||||
v = row
|
||||
else:
|
||||
v = row @ J.float().to(weight.device)
|
||||
dirs[layer] = v / v.norm().clamp_min(1e-8)
|
||||
return dirs
|
||||
|
||||
def add(self, lens_manager, jl, *, token_id, mode="scale", factor=0.0,
|
||||
replacement_id=None, layers=None, enabled=True):
|
||||
with self._lock:
|
||||
lens = lens_manager.lens
|
||||
if lens is None:
|
||||
raise ValueError("no lens loaded")
|
||||
if mode not in ("scale", "replace"):
|
||||
raise ValueError(f"invalid mode: {mode}")
|
||||
if mode == "replace" and replacement_id is None:
|
||||
raise ValueError("replacement_id required in replace mode")
|
||||
n_layers = len(jl.layers)
|
||||
if layers is None:
|
||||
layers = default_layers(n_layers)
|
||||
# layers=[] is valid: rule recorded but inactive
|
||||
layers = sorted({int(l) for l in layers if 0 <= int(l) < n_layers})
|
||||
weight = jl._lm_head.weight
|
||||
if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
||||
raise ValueError("interventions unavailable on a quantized model")
|
||||
tokenizer = jl.tokenizer
|
||||
rule = {
|
||||
"id": next(self._counter),
|
||||
"token_id": int(token_id),
|
||||
"token": tokenizer.decode([int(token_id)]),
|
||||
"mode": mode,
|
||||
"factor": float(factor),
|
||||
"replacement_id": int(replacement_id) if replacement_id is not None else None,
|
||||
"replacement": tokenizer.decode([int(replacement_id)]) if replacement_id is not None else None,
|
||||
"layers": [int(l) for l in layers],
|
||||
"enabled": bool(enabled),
|
||||
"dirs_a": self._direction(lens, weight, int(token_id), layers),
|
||||
"dirs_b": self._direction(lens, weight, int(replacement_id), layers)
|
||||
if replacement_id is not None
|
||||
else None,
|
||||
}
|
||||
self._rules.append(rule)
|
||||
return self.summary()
|
||||
|
||||
def update(self, rule_id, *, factor=None, layers=None, enabled=None,
|
||||
token_id=None, replacement_id=None, mode=None,
|
||||
lens_manager=None, jl=None):
|
||||
with self._lock:
|
||||
for rule in self._rules:
|
||||
if rule["id"] != rule_id:
|
||||
continue
|
||||
if factor is not None:
|
||||
rule["factor"] = float(factor)
|
||||
if enabled is not None:
|
||||
rule["enabled"] = bool(enabled)
|
||||
# token / replacement / mode / layers change the directions →
|
||||
# the lens and model are required to re-resolve them
|
||||
needs_dirs = any(x is not None for x in (layers, token_id, replacement_id, mode))
|
||||
if not needs_dirs:
|
||||
return self.summary()
|
||||
if lens_manager is None or jl is None:
|
||||
raise ValueError("model and lens required to edit the rule")
|
||||
lens = lens_manager.lens
|
||||
if lens is None:
|
||||
raise ValueError("no lens loaded")
|
||||
tokenizer = jl.tokenizer
|
||||
if mode is not None:
|
||||
if mode not in ("scale", "replace"):
|
||||
raise ValueError(f"invalid mode: {mode}")
|
||||
rule["mode"] = mode
|
||||
if token_id is not None:
|
||||
rule["token_id"] = int(token_id)
|
||||
rule["token"] = tokenizer.decode([int(token_id)])
|
||||
if replacement_id is not None:
|
||||
rule["replacement_id"] = int(replacement_id)
|
||||
rule["replacement"] = tokenizer.decode([int(replacement_id)])
|
||||
if rule["mode"] == "scale":
|
||||
rule["replacement_id"] = None
|
||||
rule["replacement"] = None
|
||||
elif rule["replacement_id"] is None:
|
||||
raise ValueError("replacement_id required in replace mode")
|
||||
if layers is not None:
|
||||
n_layers = len(jl.layers)
|
||||
# new_layers=[] is valid: rule kept but inactive
|
||||
rule["layers"] = sorted({int(l) for l in layers if 0 <= int(l) < n_layers})
|
||||
weight = jl._lm_head.weight
|
||||
rule["dirs_a"] = self._direction(lens, weight, rule["token_id"], rule["layers"])
|
||||
rule["dirs_b"] = (
|
||||
self._direction(lens, weight, rule["replacement_id"], rule["layers"])
|
||||
if rule["replacement_id"] is not None
|
||||
else None
|
||||
)
|
||||
return self.summary()
|
||||
raise ValueError(f"unknown rule {rule_id}")
|
||||
|
||||
def remove(self, rule_id=None):
|
||||
with self._lock:
|
||||
self.detach()
|
||||
if rule_id is None:
|
||||
self._rules = []
|
||||
else:
|
||||
self._rules = [r for r in self._rules if r["id"] != rule_id]
|
||||
return self.summary()
|
||||
|
||||
def attach(self, jl):
|
||||
if not self._rules:
|
||||
return
|
||||
if self._mode == "abliteration":
|
||||
self._attach_abliteration(jl)
|
||||
return
|
||||
if self._mode in ("readthrough", "exact"):
|
||||
self._attach_rebase(jl, exact=self._mode == "exact")
|
||||
return
|
||||
by_layer = {}
|
||||
for rule in self._active_rules():
|
||||
for layer in rule["layers"]:
|
||||
by_layer.setdefault(layer, []).append(rule)
|
||||
|
||||
def make_hook(layer, rules):
|
||||
def hook(module, inputs, output):
|
||||
h = output[0] if isinstance(output, tuple) else output
|
||||
g = self._scale
|
||||
for rule in rules:
|
||||
alpha, beta = effective_coeffs(rule["mode"], rule["factor"], g)
|
||||
vA = rule["dirs_a"][layer].to(h.device, h.dtype)
|
||||
coef = (h * vA).sum(-1, keepdim=True)
|
||||
h = h + alpha * coef * vA
|
||||
if beta:
|
||||
vB = rule["dirs_b"][layer].to(h.device, h.dtype)
|
||||
h = h + beta * coef * vB
|
||||
if isinstance(output, tuple):
|
||||
return (h,) + tuple(output[1:])
|
||||
return h
|
||||
|
||||
return hook
|
||||
|
||||
self._handles = [
|
||||
jl.layers[layer].register_forward_hook(make_hook(layer, rules))
|
||||
for layer, rules in by_layer.items()
|
||||
]
|
||||
|
||||
def _attach_abliteration(self, jl):
|
||||
# Abliteration-mode preview: the SAME projection on every residual write
|
||||
# (embed + each block's output), mirroring the pure-weight bake. A rule's
|
||||
# layers make no sense here (global projection), but layers=[] stays THE
|
||||
# "rule disabled" gesture: we honor it too.
|
||||
active = self._active_rules()
|
||||
if not active:
|
||||
return
|
||||
weight_u = jl._lm_head.weight
|
||||
dirs = [(abliteration_direction(weight_u, r), r) for r in active]
|
||||
|
||||
def apply(h):
|
||||
g = self._scale
|
||||
for (v_a, v_b), rule in dirs:
|
||||
alpha, beta = effective_coeffs(rule["mode"], rule["factor"], g)
|
||||
va = v_a.to(h.device, h.dtype)
|
||||
coef = (h * va).sum(-1, keepdim=True)
|
||||
h = h + alpha * coef * va
|
||||
if beta:
|
||||
h = h + beta * coef * v_b.to(h.device, h.dtype)
|
||||
return h
|
||||
|
||||
def emb_hook(module, inputs, output):
|
||||
return apply(output)
|
||||
|
||||
def blk_hook(module, inputs, output):
|
||||
h = output[0] if isinstance(output, tuple) else output
|
||||
h = apply(h)
|
||||
return (h,) + tuple(output[1:]) if isinstance(output, tuple) else h
|
||||
|
||||
self._handles = [jl._embed_tokens.register_forward_hook(emb_hook)]
|
||||
self._handles += [blk.register_forward_hook(blk_hook) for blk in jl.layers]
|
||||
|
||||
def _attach_rebase(self, jl, exact):
|
||||
# readthrough/exact preview: the SAME transform as the bake (core/rebase),
|
||||
# applied by hooks on the OUTPUT of the reading RMSNorms (and, in exact
|
||||
# mode, on the downstream writes) — the preview and the exported
|
||||
# checkpoint differ only by rounding.
|
||||
from core import rebase # local import (rebase imports effective_coeffs from here)
|
||||
|
||||
active = self._active_rules()
|
||||
if not active:
|
||||
return
|
||||
n_layers = len(jl.layers)
|
||||
cums = rebase.cumulative(active, self._scale, n_layers)
|
||||
if not cums:
|
||||
return
|
||||
|
||||
def read_hook_for(norm, U, V):
|
||||
Ug, Vg = rebase.gamma_pair(norm, U, V)
|
||||
weight = norm.weight
|
||||
Ug = Ug.to(weight.device, weight.dtype)
|
||||
Vg = Vg.to(weight.device, weight.dtype)
|
||||
|
||||
def hook(module, inputs, output):
|
||||
return output + (output @ Vg) @ Ug.T
|
||||
|
||||
return hook
|
||||
|
||||
def write_hook_for(module, U_inv, V):
|
||||
weight = module.weight
|
||||
U_inv = U_inv.to(weight.device, weight.dtype)
|
||||
V = V.to(weight.device, weight.dtype)
|
||||
|
||||
def hook(module, inputs, output):
|
||||
return output - (output @ V) @ U_inv.T
|
||||
|
||||
return hook
|
||||
|
||||
handles = []
|
||||
for m in sorted(k for k in cums if k < n_layers):
|
||||
U, V = cums[m]
|
||||
block = jl.layers[m]
|
||||
norms = {}
|
||||
for _suffix, _module, norm in rebase.iter_reads(block):
|
||||
norms[id(norm)] = norm
|
||||
for norm in norms.values():
|
||||
handles.append(norm.register_forward_hook(read_hook_for(norm, U, V)))
|
||||
if exact:
|
||||
U_inv, Vw, _regularized = rebase.inverse_uv(U, V)
|
||||
for _suffix, module in rebase.iter_writes(block):
|
||||
handles.append(module.register_forward_hook(write_hook_for(module, U_inv, Vw)))
|
||||
U, V = cums[n_layers]
|
||||
handles.append(
|
||||
jl._final_norm.register_forward_hook(read_hook_for(jl._final_norm, U, V))
|
||||
)
|
||||
self._handles = handles
|
||||
|
||||
def detach(self):
|
||||
for handle in self._handles:
|
||||
handle.remove()
|
||||
self._handles = []
|
||||
Reference in New Issue
Block a user