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 = []