evil_MoE/src/projected_grpo/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 _signed_cos(c: Float[torch.Tensor, "k"], gn: torch.Tensor) -> float:
    """Signed scalar projection of g onto the hack-oriented span of V.

    c = V @ g (per-axis coefficients with V rows orthonormal and oriented
    hack-ward, so c_i > 0 means "grad pushes hack-ward on axis i").
    We return sum(c) / ||g||, which is bounded in [-||c||/||g||, +||c||/||g||]
    and is positive when the dominant per-axis components push toward hack,
    negative when they push toward safe.

    Replaces the older unsigned ||c||/||g|| ratio: that magnitude hid the
    direction (after a one_sided projection it stayed positive even though
    the residual was all safe-pointing), so we couldn't read the sign off
    a single column.
    """
    return (c.sum() / 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 sum(V @ g) / ||g||, signed.

    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(_signed_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 SIGNED scalars (sum of per-axis V @ g coefficients,
    normalized by ||g||). Positive = grad pushes toward hack; negative = grad
    pushes toward safe. Under one_sided projection cos_post should fall to
    zero or negative (we removed the positive part). Under no_gate cos_post
    is approximately zero by construction. Under reverse cos_post flips sign
    relative to cos_pre (we subtract 2*c@V, so V@g_proj = -V@g).

    `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 = _signed_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 = _signed_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,
    route: 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).
    `route`: erase AND park the removed hack-ward component in the quarantine
      knob delta_S_hack.grad (Gradient Routing, Cloud 2410.04332). delta_S gets
      the IDENTICAL g_proj as erase (same gate/preserve/overshoot), so the
      deployment model -- delta_S with delta_S_hack zeroed at eval -- evolves
      under the same update rule as the erase arm (each is its own AdamW param;
      the quarantine's separate optimizer state cannot perturb delta_S). That is
      the sense in which route ⊇ erase: erase == route with the quarantine
      discarded. CAVEAT (not an identity): the combined TRAINING forward
      delta_S + delta_S_hack does NOT reproduce a vanilla update -- AdamW steps
      the two knobs independently, so the sum over-moves hack-ward. That is
      intended (the model keeps hacking during training so the capability lands
      in the quarantine), and it only affects the training trajectory, never the
      ablated deployment. Mutually exclusive with measure_only.

    Diagnostics returned (per call, averaged over modules):
      mean_cos_pre   = mean over modules of sum(V @ g)/||g||, signed
      mean_cos_post  = same after projection
      frac_fired     = fraction of modules where at least one direction fired (c_i > 0)
    """
    if route and measure_only:
        raise ValueError("route and measure_only are mutually exclusive")
    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, removed, 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          # same update rule as erase
            if route:
                # quarantine the discarded hack-ward part; removed ∈ span(V),
                # ablated at eval so its magnitude/overshoot scaling is harmless.
                info["delta_S_hack"].grad = removed
        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,
    }