src/mxf/inject.py

"""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)