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,
) -> tuple[Float[torch.Tensor, "r"], float, float, bool]:
    """Per-module top-k removal. Returns (g_proj, cos_pre, cos_post, fired).

    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).
    """
    gn = g.norm()
    if gn < 1e-12:
        return g, 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, cos_pre, cos_pre, False
    g_proj = g - c_use @ V                          # [r]
    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, 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",
) -> 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 (vanilla arm diagnostic).

    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)
    """
    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)
        cos_pre_list.append(cos_pre)
        cos_post_list.append(cos_post)
        if fired:
            if not measure_only:
                info["delta_S"].grad = g_proj
            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,
    }