evil_MoE/src/vgrout/lora2r.py

"""lora2r adapter: one rank-2r LoRA per target Linear, A and B both trainable,
partitioned into a deployed block [:r] and a quarantine block [r:].

    y += B@(A@x) - B0@(A0@x)

A0/B0 are FROZEN copies of the (seeded Gaussian) init, subtracted so the net
delta is exactly 0 at init while h = A@x is alive. Both factors must be nonzero
at init: the gate reads c.grad = h ⊙ (Bᵀδ) and A.grad = (Bᵀδ)xᵀ, both
identically zero under standard zero-B LoRA -- no extraction, no band
calibration, no trainable A at step 0. Init values beyond "nonzero" are NOT
load-bearing (we previously used PiSSA; see docs/spec/20260610_lora2r_v2_plan.md
T2 for why it was dropped).

[B|B_q] @ ([A;A_q]@x) has no cross terms (column b_k only ever multiplies row
a_k), so the two blocks are independent adapters. Per-rollout block masks on this
tensor mirror the retain/forget parameter partition used by SGTM (Cloud et al.).
"""
from __future__ import annotations

import torch
import torch.nn.functional as F
from jaxtyping import Float
from loguru import logger
from torch import Tensor, nn

TARGET_SUFFIXES = (
    # full attention
    "q_proj", "k_proj", "v_proj", "o_proj",
    # linear-attention / GatedDeltaNet
    "in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b", "out_proj",
    # MLP
    "up_proj", "gate_proj", "down_proj",
)


def is_target(name: str) -> bool:
    return name.split(".")[-1] in TARGET_SUFFIXES


def _lora2r_hook(layer: nn.Linear, args: tuple, y: Tensor) -> Tensor:
    """Add the two-block delta to y, applying per-rollout block masks.

    Block masks (layer._lora2r_mask = (m, d), set by train.py per loss pass;
    None = unmasked for generation / gate pass / eval):
      m [G] quarantine on/off -- m=0: quarantine zero in forward AND backward,
            so the deployed block trains in its post-ablation state
      d [G] deployed detach   -- d=1: deployed kept in forward, zero grad
            (hack-gated rollouts update ONLY the quarantine block)
    Masks act on branch OUTPUTS so a detach blocks grads to BOTH the A rows and
    the B columns of that block.

    grad probe: c = ones[..., 2r] spliced as h*c. After backward
    c.grad = h ⊙ (Bᵀδ_y) = the per-sample WEIGHT grad of a virtual diagonal
    scale between A and B. Training no longer uses it (the routeA gate reads
    activations); kept for scripts/diag_pinning.py, whose grad-score baseline
    reads this space.
    """
    (x,) = args                                       # x: [..., d_in]
    A:  Float[Tensor, "two_r d_in"]  = layer._lora2r_A    # trainable
    B:  Float[Tensor, "d_out two_r"] = layer._lora2r_B    # trainable
    A0: Float[Tensor, "two_r d_in"]  = layer._lora2r_A0   # frozen init (subtracted: net delta 0 at init)
    B0: Float[Tensor, "d_out two_r"] = layer._lora2r_B0
    r = layer._lora2r_r
    h = F.linear(x, A.to(x.dtype))                    # [..., 2r]
    if layer._lora2r_grad_probe and torch.is_grad_enabled():
        c = torch.ones(h.shape[0], *([1] * (h.dim() - 2)), h.shape[-1],
                       device=h.device, dtype=h.dtype, requires_grad=True)
        layer._lora2r_gate = c
        h = h * c
    h0 = F.linear(x, A0.to(x.dtype))                  # [..., 2r] frozen init path
    dep = (F.linear(h[..., :r], B[:, :r].to(x.dtype))
           - F.linear(h0[..., :r], B0[:, :r].to(x.dtype)))
    quar = (F.linear(h[..., r:], B[:, r:].to(x.dtype))
            - F.linear(h0[..., r:], B0[:, r:].to(x.dtype)))
    if layer._lora2r_mask is not None:
        m, d = layer._lora2r_mask                                # [G] each
        G = m.shape[0]
        shape = dep.shape                                        # [G, s, d_out] or [G*s, d_out]
        dep = dep.reshape(G, -1, shape[-1])
        quar = quar.reshape(G, -1, shape[-1])
        d_ = d.view(G, 1, 1).to(dep.dtype)
        dep = ((1 - d_) * dep + d_ * dep.detach()).reshape(shape)
        quar = (m.view(G, 1, 1).to(quar.dtype) * quar).reshape(shape)
    return y + (dep + quar).to(y.dtype)


def wrap_model_with_lora2r(
    model: nn.Module,
    r: int = 32,
    init_seed: int = 0,
    grad_probe: bool = False,
) -> dict[str, dict]:
    """Attach a rank-2r Gaussian-init LoRA (A AND B trainable) to every target Linear.

    Init: A0 ~ N(0, 1/d_in) [2r, d_in], B0 ~ N(0, 1/2r) [d_out, 2r], seeded per
    module so runs reproduce; blocks are iid -> statistically matched. W stays
    untouched; the hook subtracts the frozen A0/B0 contribution. The quarantine's
    learned delta is (A[r:], B[:, r:]) minus init; deploy ablation resets that
    block to A0/B0 (eval.ablate_quarantine).

    Info dict per module: {layer, A, B, A0, B0, handle, r}; quarantine = block
    slices, no separate tensor.
    """
    targets = [(n, m) for n, m in model.named_modules()
               if isinstance(m, nn.Linear) and is_target(n)]
    logger.info(f"lora2r attach: {len(targets)} target Linear modules, "
                f"r={r}/block (2r={2 * r}), Gaussian init seed={init_seed}, A+B trainable")
    out: dict[str, dict] = {}
    for i, (name, linear) in enumerate(targets):
        d_out, d_in = linear.weight.shape
        dev = linear.weight.device
        gen = torch.Generator().manual_seed(init_seed * 100003 + i)
        A0 = (torch.randn(2 * r, d_in, generator=gen) / d_in ** 0.5).to(device=dev, dtype=torch.float32)
        B0 = (torch.randn(d_out, 2 * r, generator=gen) / (2 * r) ** 0.5).to(device=dev, dtype=torch.float32)
        linear.register_buffer("_lora2r_A0", A0, persistent=True)
        linear.register_buffer("_lora2r_B0", B0, persistent=True)
        A = nn.Parameter(A0.clone())
        B = nn.Parameter(B0.clone())
        linear.register_parameter("_lora2r_A", A)
        linear.register_parameter("_lora2r_B", B)
        linear._lora2r_r = r
        linear._lora2r_grad_probe = grad_probe
        linear._lora2r_gate = None
        linear._lora2r_mask = None
        out[name] = {"layer": linear, "A": A, "B": B, "A0": A0, "B0": B0,
                     "handle": linear.register_forward_hook(_lora2r_hook), "r": r}
    trainable = ("_lora2r_A", "_lora2r_B")
    for n, p in model.named_parameters():
        if not n.endswith(trainable):
            p.requires_grad_(False)
    return out