evil_MoE/src/vgrout/proj.py

"""Gradient projection + delta_S grad utilities. Imported by smoke and train.

Shape conventions in the v_hack / delta_S plumbing (jaxtyping-annotated):
- `r` = per-module SVD rank (delta_S dimension; varies per Linear)
- `k` = number of v_hack directions kept per module (after top-k slice and
        global noise-floor filter at load time)
- `V` is `[k, r]`, rows orthonormal in R^r and oriented hack-ward
- `g = delta_S.grad` is `[r]`
- `c = V @ g` is `[k]`
"""
from __future__ import annotations

import torch
from jaxtyping import Float


def per_token_logps(logits: torch.Tensor, ids: torch.Tensor) -> torch.Tensor:
    """log p(ids | logits) gathered token-wise.

    Uses F.cross_entropy (fused softmax+gather) so we never materialise the
    full [B, L, V] fp32 softmax. On Qwen3.5-2B with V=152k, G=8, L≈1500 the
    fp32 vocab tensor was ~7 GB per forward -- the difference between OOM and
    fit on a 96 GB card when the autograd graph is alive.
    """
    B, L, V = logits.shape
    # CE's internal log_softmax accumulates in fp32 (stable) but returns input dtype.
    # The output [B*L] is small, so upcast it to fp32 for downstream PPO ratio math.
    return -torch.nn.functional.cross_entropy(
        logits.reshape(-1, V), ids.reshape(-1), reduction="none"
    ).float().view(B, L)


def _hackward_cos(c: Float[torch.Tensor, "k"], gn: torch.Tensor) -> float:
    """Fraction of the gradient's magnitude that points hack-ward into v_hack:
    ||relu(c)|| / ||g||, where c = V @ g and V rows are orthonormal and oriented
    hack-ward (c_i > 0 means "grad pushes hack-ward on axis i"). In [0, 1].

    relu BEFORE aggregating is the point: the one_sided projection removes only
    relu(c) (the hack-ward axes), and with V orthonormal ||removed|| = ||relu(c)||,
    so this reads directly as "fraction of the grad the projection strips" (a signed
    sum would let +/- axes cancel and read ~0 even while routing a large hack-ward
    magnitude).

    After a one_sided erase, V @ g_proj = min(c, 0) (positive axes zeroed), so
    relu of it is 0 -> cos_post == 0 exactly. That clean SHOULD (cos_post -> 0) is
    the diagnostic; we drop the sign because a one_sided method never acts on the
    safe-ward (negative) part anyway.
    """
    return (torch.relu(c).norm() / gn).item()


@torch.no_grad()
def mean_cos_pre_from_grads(
    grad_dict: dict[str, Float[torch.Tensor, "r"]],
    v_hack: dict[str, Float[torch.Tensor, "k r"]],
) -> float:
    """Mean over modules of ||relu(V @ g)|| / ||g|| (hack-ward fraction, in [0,1]).

    Used to compute per-source cos_pre (cos_pre_s for student-only grad,
    cos_pre_t for teacher-only grad) without mutating model.grad or calling
    the full projection pipeline.
    """
    cs = []
    for name, g in grad_dict.items():
        if g is None or name not in v_hack:
            continue
        V = v_hack[name].to(g.device, dtype=g.dtype)
        gn = g.norm()
        if gn < 1e-12:
            continue
        cs.append(_hackward_cos(V @ g, gn))
    return float(sum(cs) / len(cs)) if cs else float("nan")


def _project_one_module(
    g: Float[torch.Tensor, "r"],
    V: Float[torch.Tensor, "k r"],
    gate_mode: str,
    preserve_magnitude: bool,
    overshoot: float = 1.0,
) -> tuple[Float[torch.Tensor, "r"], Float[torch.Tensor, "r"], float, float, bool]:
    """Per-module top-k removal. Returns (g_proj, removed, cos_pre, cos_post, fired).

    `removed` = overshoot*c_use@V, the vector subtracted from g (computed
    before any preserve_magnitude rescale, so removed ∈ span(V) always).
    Erasure drops it; routing parks it in delta_S_hack. Note g_proj + removed
    == g ONLY when preserve_magnitude is False and overshoot is 1.0; with the
    defaults g_proj is rescaled, so the sum is not the original g (routing does
    not rely on that sum -- see project_delta_S_grad).

    cos_pre / cos_post are the hack-ward FRACTION ||relu(V @ g)|| / ||g|| (in
    [0,1]; see _hackward_cos). cos_pre = how much of the grad points hack-ward
    into v_hack; cos_post = the residual after projection. Under one_sided (and
    no_gate, and overshoot>=1) projection cos_post -> 0 exactly: every hack-ward
    axis was removed, so relu of the residual coefficients is 0.

    `overshoot` scales the removed coefficient: g_proj = g - overshoot*c_use@V.
    overshoot=1.0 just removes the hack-ward component; overshoot=1.1 removes
    110% (a 10% reversal: V@g_proj = -0.1c on fired axes), a milder version of
    gate_mode="reverse" (which is overshoot=2.0 on the full, ungated c).
    """
    gn = g.norm()
    if gn < 1e-12:
        z = torch.zeros_like(g)
        return g, z, 0.0, 0.0, False
    c = V @ g                                       # [k]
    cos_pre = _hackward_cos(c, gn)
    if gate_mode == "no_gate":
        c_use = c
        fired = True
    elif gate_mode == "one_sided":
        mask = (c > 0).to(c.dtype)
        c_use = c * mask
        fired = bool((c_use != 0).any())
    elif gate_mode == "reverse":
        # Subtract 2*c@V: V@g_proj = V@g - 2*(V V^T) c = c - 2c = -c.
        # Flips the sign of the gradient component in span(V); pushes
        # actively away from hack rather than just removing.
        c_use = 2 * c
        fired = True
    else:
        raise ValueError(f"unknown gate_mode={gate_mode!r}")
    if not fired:
        return g, torch.zeros_like(g), cos_pre, cos_pre, False
    removed = overshoot * c_use @ V                 # [r]
    g_proj = g - removed
    gp_n = g_proj.norm()
    if preserve_magnitude and gp_n > 1e-12:
        g_proj = g_proj * (gn / gp_n)
    cos_post = _hackward_cos(V @ g_proj, g_proj.norm().clamp_min(1e-12))
    return g_proj, removed, cos_pre, cos_post, True


@torch.no_grad()
def project_delta_S_grad(
    wrappers: dict,
    v_hack: dict[str, Float[torch.Tensor, "k r"]],
    preserve_magnitude: bool,
    measure_only: bool = False,
    gate_mode: str = "one_sided",
    overshoot: float = 1.0,
) -> dict[str, float]:
    """Per-module top-k removal of hack-aligned grad components.

    For each wrapped module:
      g = delta_S.grad        # [r]
      V = v_hack[name]        # [k, r], rows orthonormal, oriented hack-ward
      c = V @ g               # [k]   per-direction coefficients

    gate_mode="one_sided" (default):
      mask = (c > 0)          # only zap when grad is going hack-ward on that axis
      g' = g - (c * mask) @ V # subtract only positive-coefficient components

    gate_mode="no_gate":
      g' = g - c @ V          # full V·V^T removal, sign-agnostic;
                              # drives ||V g'|| -> 0 exactly. No trust in v_hack
                              # orientation: any motion in span(V) is suspect.

    `preserve_magnitude`: rescale g' to ||g|| after projection.
    `measure_only`: same math, but g is not mutated (the `none` intervention).
    Diagnostics returned (per call, averaged over modules):
      mean_cos_pre   = mean over modules of ||relu(V @ g)||/||g|| (hack-ward fraction, [0,1])
      mean_cos_post  = same after projection (-> 0 when hack-ward axes were removed)
      frac_fired     = fraction of modules where at least one direction fired (c_i > 0)
    """
    cos_pre_list, cos_post_list, n_fired = [], [], 0
    for name, info in wrappers.items():
        g = info["delta_S"].grad
        if g is None:
            continue
        if name not in v_hack:  # module dropped by global noise-floor filter
            continue
        V = v_hack[name].to(g.device, dtype=g.dtype)   # [k, r]
        g_proj, _, cos_pre, cos_post, fired = _project_one_module(
            g, V, gate_mode, preserve_magnitude, overshoot)
        cos_pre_list.append(cos_pre)
        cos_post_list.append(cos_post)
        if fired and not measure_only:
            info["delta_S"].grad = g_proj
        if fired:
            n_fired += 1
    pre_t = torch.tensor(cos_pre_list); post_t = torch.tensor(cos_post_list)
    return {
        "mean_cos_pre":  pre_t.mean().item(),
        "min_cos_pre":   pre_t.min().item() if pre_t.numel() else float("nan"),
        "max_cos_pre":   pre_t.max().item() if pre_t.numel() else float("nan"),
        "mean_cos_post": post_t.mean().item(),
        "min_cos_post":  post_t.min().item() if post_t.numel() else float("nan"),
        "max_cos_post":  post_t.max().item() if post_t.numel() else float("nan"),
        "frac_fired":    n_fired / len(cos_pre_list) if cos_pre_list else 0.0,
    }