evil_MoE/src/projected_grpo/proj.py

"""Gradient projection + delta_S grad utilities. Imported by smoke and train."""
from __future__ import annotations

import torch


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)


@torch.no_grad()
def mean_cin_from_grads(
    grad_dict: dict[str, torch.Tensor],
    v_hack: dict[str, torch.Tensor],
) -> float:
    """Mean over modules of ||V g|| / ||g||, given a dict of per-module grads.

    Used to compute per-source cin (cin_s for student-only grad, cin_t for
    teacher-only grad) without mutating model.grad or calling the full
    projection pipeline. v_hack rows are orthonormal so ||V g|| <= ||g|| and
    the ratio is in [0,1].
    """
    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(((V @ g).norm() / gn).item())
    return float(sum(cs) / len(cs)) if cs else float("nan")


@torch.no_grad()
def project_delta_S_grad(
    wrappers: dict,
    v_hack: dict[str, torch.Tensor],
    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_in   = mean over modules of ||V g||/||g||  (subspace energy fraction in)
      mean_cos_out  = same after projection
      frac_fired    = fraction of modules where at least one direction fired (c_i > 0)
    """
    cos_in_list, cos_out_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]
        gn = g.norm()
        if gn < 1e-12:
            cos_in_list.append(0.0); cos_out_list.append(0.0); continue
        c = V @ g                                       # [k]
        cin = c.norm() / gn
        cos_in_list.append(cin.item())
        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())
        else:
            raise ValueError(f"unknown gate_mode={gate_mode!r}")
        if fired:
            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)
            cout = (V @ g_proj).norm() / g_proj.norm().clamp_min(1e-12)
            cos_out_list.append(cout.item())
            if not measure_only:
                info["delta_S"].grad = g_proj
            n_fired += 1
        else:
            cos_out_list.append(cin.item())
    cin_t = torch.tensor(cos_in_list); cout_t = torch.tensor(cos_out_list)
    return {
        "mean_cos_in":  cin_t.mean().item(),
        "min_cos_in":   cin_t.min().item() if cin_t.numel() else float("nan"),
        "max_cos_in":   cin_t.max().item() if cin_t.numel() else float("nan"),
        "mean_cos_out": cout_t.mean().item(),
        "min_cos_out":  cout_t.min().item() if cout_t.numel() else float("nan"),
        "max_cos_out":  cout_t.max().item() if cout_t.numel() else float("nan"),
        "frac_fired":   n_fired / len(cos_in_list) if cos_in_list else 0.0,
    }