hyper_lora.py

"""
hyper_lora.py — AR-conditioned hypernetwork for the nanoMoE v11 hybrid.

What this does
--------------
Given a single conditioning vector z (an NLA *Activation Reconstructor* output,
e.g. Qwen2.5-7B L20 -> 3584-d), a small hypernetwork emits a full adapter for a
FROZEN model_hybrid.GPT and injects it. The base never trains; only the
hypernetwork does. This is the "more expressive" target you picked:

  spectral layers : LoRA on in_proj, c_proj, and the implicit-kernel generator
                    MLP (kernel_mlp.0/2/4)  +  FiLM (per-channel additive delta)
                    on the kernel envelope (log_alpha) and gain (scale)
  attention layers: LoRA on q_proj, k_proj, v_proj, c_proj
  FFN             : optional LoRA on c_fc/c_proj (off by default — it adds
                    capacity but dilutes the Fourier-native story; flip
                    adapt_ffn=True if the spectral path underfits)

Because the kernel is *generated* by kernel_mlp, adapting it needs no dense
spatial dW and no inverse DFT: we perturb the generator's params + the envelope,
and causal_fft_conv runs exactly as before. That is the architecture-native win.

Pipeline at train time
-----------------------
    z      = AR(task_descriptor_text)          # your AR, computed upstream
    deltas = hyper(z)                           # this module
    model.set_deltas(deltas)                    # distribute into the wrapped base
    logits, loss = model(x, y)                  # x,y are examples from THAT task
    loss.backward()                             # grads -> hyper only

One z corresponds to one task; the (x, y) batch are examples from that task that
share the adapter. To train, you need (z, oracle-adapter) or (z, task-data)
pairs — see build_hyper_dataset.py for generating the descriptor/example side.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

import torch
import torch.nn as nn
from torch.nn import functional as F

from model_hybrid import GPT, ImplicitKernel


# ---------------------------------------------------------------------------
# Injection primitives
# ---------------------------------------------------------------------------

class LoRALinear(nn.Module):
    """Wraps a frozen nn.Linear. The low-rank factors are NOT parameters — they
    are written each step from the hypernetwork output (so autograd flows back
    into the hypernetwork, not into A/B). With no delta set, == the base."""

    def __init__(self, base: nn.Linear, rank: int, alpha: float | None = None):
        super().__init__()
        self.in_features = base.in_features
        self.out_features = base.out_features
        self.rank = rank
        self.scaling = (alpha if alpha is not None else rank) / rank
        # store base weights as buffers -> frozen, not in the optimizer
        self.register_buffer("weight", base.weight.detach().clone())
        if base.bias is not None:
            self.register_buffer("bias", base.bias.detach().clone())
        else:
            self.bias = None
        self._A = None        # (rank, in)   set externally
        self._B = None        # (out, rank)  set externally

    def set_delta(self, A: torch.Tensor | None, B: torch.Tensor | None):
        self._A, self._B = A, B

    def forward(self, x):
        y = F.linear(x, self.weight, self.bias)
        if self._A is not None:
            # x @ A^T -> (..., rank);  (...) @ B^T -> (..., out)
            y = y + self.scaling * F.linear(F.linear(x, self._A), self._B)
        return y


class FiLMKernel(nn.Module):
    """Re-implements ImplicitKernel.forward verbatim but adds two per-channel
    deltas: dlog_alpha (envelope / memory horizon) and dscale (gain). Reuses the
    *original* kernel's buffers, frozen params, and (LoRA-wrapped) kernel_mlp."""

    def __init__(self, orig: ImplicitKernel):
        super().__init__()
        self.orig = orig                  # holds frozen params + (wrapped) kernel_mlp
        self._dlog_alpha = None           # (dim,)
        self._dscale = None               # (dim,)

    def set_delta(self, dlog_alpha, dscale):
        self._dlog_alpha, self._dscale = dlog_alpha, dscale

    def forward(self, T, device):
        k = self.orig
        t = torch.arange(T, device=device, dtype=torch.float32) / k.block_size
        ang = 2 * math.pi * t[:, None] * k.freqs[None, :].to(device)
        feats = torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1)
        h = k.kernel_mlp(feats)                                  # (T, D), LoRA applies here
        log_alpha = k.kernel_log_alpha
        if self._dlog_alpha is not None:
            log_alpha = log_alpha + self._dlog_alpha
        env = torch.exp(-torch.exp(log_alpha)[None, :] * t[:, None])
        h = h * env
        scale = k.kernel_scale
        if self._dscale is not None:
            scale = scale + self._dscale
        h = h / (h.norm(dim=0, keepdim=True) + 1e-6) * scale[None, :]
        return h.t()                                             # (D, T)


# ---------------------------------------------------------------------------
# Site discovery: walk the real model and record what we will adapt
# ---------------------------------------------------------------------------

@dataclass
class SiteSpec:
    name: str          # unique key, e.g. "h.0.mixer.in_proj"
    kind: str          # "lora" | "film"
    in_dim: int = 0    # lora
    out_dim: int = 0   # lora
    dim: int = 0       # film
    rank: int = 0      # lora


@dataclass
class AdaptConfig:
    rank_kernel: int = 8       # LoRA rank on the kernel generator MLP
    rank_mixer: int = 8        # LoRA rank on in_proj / c_proj
    rank_attn: int = 8         # LoRA rank on q/k/v/o
    rank_ffn: int = 8          # LoRA rank on FFN (only if adapt_ffn)
    adapt_ffn: bool = False
    film: bool = True          # FiLM on kernel envelope + gain


def _swap(parent: nn.Module, child_name: str, new: nn.Module):
    setattr(parent, child_name, new)


def wrap_model(model: GPT, cfg: AdaptConfig):
    """In-place: replace target Linears with LoRALinear, kernels with FiLMKernel.
    Returns (sites, mods) where `mods` maps each site name to the live module —
    so injection never has to re-resolve a (post-wrap) dotted path."""
    sites: list[SiteSpec] = []
    mods: dict[str, nn.Module] = {}

    def wrap_linear(parent, attr, rank, prefix):
        base = getattr(parent, attr)
        assert isinstance(base, nn.Linear), f"{prefix}.{attr} is {type(base)}"
        lo = LoRALinear(base, rank)
        _swap(parent, attr, lo)
        name = f"{prefix}.{attr}"
        sites.append(SiteSpec(name, "lora",
                              in_dim=base.in_features, out_dim=base.out_features,
                              rank=rank))
        mods[name] = lo

    for i, block in enumerate(model.transformer.h):
        bp = f"h.{i}"
        mixer = block.mixer
        if block.is_attn:
            for attr in ("q_proj", "k_proj", "v_proj", "c_proj"):
                wrap_linear(mixer, attr, cfg.rank_attn, f"{bp}.mixer")
        else:
            wrap_linear(mixer, "in_proj", cfg.rank_mixer, f"{bp}.mixer")
            wrap_linear(mixer, "c_proj", cfg.rank_mixer, f"{bp}.mixer")
            # kernel generator MLP: indices 0,2,4 are Linears (wrap BEFORE FiLM)
            kmlp = mixer.kernel.kernel_mlp
            for idx in (0, 2, 4):
                wrap_linear(kmlp, str(idx), cfg.rank_kernel, f"{bp}.mixer.kernel.kernel_mlp")
            # FiLM on the envelope/gain — wrap the kernel itself
            orig_kernel = mixer.kernel
            film = FiLMKernel(orig_kernel)
            _swap(mixer, "kernel", film)
            if cfg.film:
                name = f"{bp}.mixer.kernel"
                sites.append(SiteSpec(name, "film", dim=orig_kernel.dim))
                mods[name] = film
        if cfg.adapt_ffn:
            wrap_linear(block.mlp, "c_fc", cfg.rank_ffn, f"{bp}.mlp")
            wrap_linear(block.mlp, "c_proj", cfg.rank_ffn, f"{bp}.mlp")

    # freeze everything in the base
    for p in model.parameters():
        p.requires_grad_(False)
    return sites, mods


# ---------------------------------------------------------------------------
# The hypernetwork: z -> {site -> delta}
# ---------------------------------------------------------------------------

class HyperLoRA(nn.Module):
    """T2L-style and SMALL: a shared trunk over z, a per-site embedding, and ONE
    bottlenecked head per *shape group* (e.g. all q_proj across layers share a
    head). The per-site embedding is what specialises a shared head to each site,
    so the head count is ~#module-types, not #sites. This is the difference
    between a hypernetwork larger than the base and one a fraction of its size."""

    def __init__(self, sites: list[SiteSpec], d_z: int,
                 d_trunk: int = 512, d_site: int = 64, d_mid: int = 128,
                 dropout: float = 0.0):
        super().__init__()
        self.sites = sites
        self.trunk = nn.Sequential(
            nn.Linear(d_z, d_trunk), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(d_trunk, d_trunk), nn.GELU(),
        )
        self.site_emb = nn.ParameterDict()
        self.heads = nn.ModuleDict()
        self.group_of: dict[str, str] = {}
        self.out_of: dict[str, int] = {}

        def group_key(s: SiteSpec) -> str:
            if s.kind == "lora":
                return f"lora_{s.in_dim}_{s.out_dim}_{s.rank}"
            return f"film_{s.dim}"

        for s in sites:
            key = s.name.replace(".", "__")
            self.site_emb[key] = nn.Parameter(torch.randn(d_site) * 0.02)
            g = group_key(s)
            self.group_of[s.name] = g
            out = s.rank * (s.in_dim + s.out_dim) if s.kind == "lora" else 2 * s.dim
            self.out_of[s.name] = out
            if g not in self.heads:
                head = nn.Sequential(
                    nn.Linear(d_trunk + d_site, d_mid), nn.GELU(),
                    nn.Linear(d_mid, out),
                )
                final = head[-1]
                if s.kind == "lora":
                    # standard LoRA init translated to a hypernet: the A-slice of
                    # the output starts nonzero, the B-slice starts zero. Then
                    # dW = B@A = 0 at init (identity preserved) BUT dL/dB != 0
                    # because A != 0 -> B moves immediately, then A follows. If
                    # BOTH start at zero their gradients deadlock and the LoRA
                    # sites never train.
                    a = s.rank * s.in_dim
                    nn.init.normal_(final.weight[:a], std=0.02)
                    nn.init.zeros_(final.weight[a:])
                    nn.init.zeros_(final.bias)
                else:                                   # film: additive, zero is fine
                    nn.init.zeros_(final.weight); nn.init.zeros_(final.bias)
                self.heads[g] = head

    def forward(self, z: torch.Tensor) -> dict:
        """z: (d_z,) for one task. Returns {site_name: payload}."""
        if z.dim() == 2:
            assert z.shape[0] == 1, "one task (one z) per call; pass z[0] for a batch of 1"
            z = z[0]
        h = self.trunk(z)
        out = {}
        for s in self.sites:
            key = s.name.replace(".", "__")
            feat = torch.cat([h, self.site_emb[key]], dim=-1)
            raw = self.heads[self.group_of[s.name]](feat)
            if s.kind == "lora":
                a = s.rank * s.in_dim
                A = raw[:a].view(s.rank, s.in_dim)
                B = raw[a:].view(s.out_dim, s.rank)
                out[s.name] = ("lora", A, B)
            else:
                out[s.name] = ("film", raw[: s.dim], raw[s.dim:])
        return out


# ---------------------------------------------------------------------------
# Glue: distribute hyper output into the wrapped base
# ---------------------------------------------------------------------------

class AdaptedModel(nn.Module):
    """Wraps a frozen, already-wrapped GPT. set_deltas() pushes one task's
    adapter in; forward() then runs the specialised model."""

    def __init__(self, model: GPT, mods: dict[str, nn.Module]):
        super().__init__()
        self.model = model
        self._index = mods

    def set_deltas(self, deltas: dict | None):
        for name, mod in self._index.items():
            if deltas is None:
                mod.set_delta(None, None)
                continue
            payload = deltas[name]
            mod.set_delta(payload[1], payload[2])     # ("lora"/"film", a, b)

    def forward(self, idx, targets=None):
        return self.model(idx, targets)


def build(model: GPT, d_z: int, cfg: AdaptConfig | None = None,
          d_trunk: int = 512, dropout: float = 0.0):
    """Convenience: wrap `model` in place, build the hypernetwork for it."""
    cfg = cfg or AdaptConfig()
    sites, mods = wrap_model(model, cfg)
    hyper = HyperLoRA(sites, d_z=d_z, d_trunk=d_trunk, dropout=dropout)
    adapted = AdaptedModel(model, mods)
    return adapted, hyper, sites


# ---------------------------------------------------------------------------
# Self-test against the REAL model_hybrid.GPT
# ---------------------------------------------------------------------------

def _self_test():
    from model_hybrid import GPTConfig
    torch.manual_seed(0)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"== hyper_lora self-test (device={dev}) ==")

    cfg_m = GPTConfig(block_size=64, vocab_size=256, n_layer=6, n_head=4,
                      n_embd=64, mixer="hybrid", attn_every=3, kv_heads=1,
                      kernel_d_hidden=32, kernel_n_fourier=8)
    base = GPT(cfg_m).to(dev)

    # snapshot the clean base output BEFORE wrapping, for the identity check
    x = torch.randint(0, cfg_m.vocab_size, (2, cfg_m.block_size), device=dev)
    y = torch.randint(0, cfg_m.vocab_size, (2, cfg_m.block_size), device=dev)
    base.eval()
    with torch.no_grad():
        clean_logits, _ = base(x, y)

    d_z = 128
    adapted, hyper, sites = build(base, d_z=d_z, cfg=AdaptConfig(adapt_ffn=True))
    adapted, hyper = adapted.to(dev), hyper.to(dev)
    n_lora = sum(s.kind == "lora" for s in sites)
    n_film = sum(s.kind == "film" for s in sites)
    print(f"[1/5] discovered {len(sites)} sites ({n_lora} LoRA, {n_film} FiLM)")
    hp = sum(p.numel() for p in hyper.parameters())
    print(f"      hypernetwork params: {hp/1e6:.3f}M")

    # 2. identity at init: zero-init heads -> deltas == 0 -> output == clean base
    z = torch.randn(d_z, device=dev)
    adapted.eval()
    with torch.no_grad():
        adapted.set_deltas(hyper(z))
        adapted_logits, _ = adapted(x, y)
    drift = (adapted_logits - clean_logits).abs().max().item()
    assert drift < 1e-4, f"identity-at-init broken: drift {drift}"
    print(f"[2/5] identity-at-init OK (max drift {drift:.2e} < 1e-4)")

    # 3. once heads are perturbed, the adapter actually changes the output
    with torch.no_grad():
        for p in hyper.heads.parameters():
            p.normal_(0, 0.02)
        adapted.set_deltas(hyper(z))
        moved_logits, _ = adapted(x, y)
    moved = (moved_logits - clean_logits).abs().max().item()
    assert moved > 1e-3, f"adapter had no effect: {moved}"
    print(f"[3/5] adapter changes output OK (max delta {moved:.3e})")

    # 4. gradient flow: hyper gets grads, base gets NONE
    adapted.train(); hyper.train()
    deltas = hyper(z)
    adapted.set_deltas(deltas)
    _, loss = adapted(x, y)
    loss.backward()
    hyper_with_grad = [n for n, p in hyper.named_parameters()
                       if p.grad is not None and torch.isfinite(p.grad).all()]
    hyper_total = sum(1 for _ in hyper.parameters())
    base_with_grad = [n for n, p in adapted.model.named_parameters()
                      if p.grad is not None]
    assert len(hyper_with_grad) == hyper_total, \
        f"only {len(hyper_with_grad)}/{hyper_total} hyper params got finite grad"
    assert not base_with_grad, f"FROZEN base leaked grad into: {base_with_grad[:3]}"
    print(f"[4/5] grad flow OK ({hyper_total} hyper params grad'd, base frozen)")

    # 5. two different z -> two different adapters
    hyper.zero_grad()
    with torch.no_grad():
        z2 = torch.randn(d_z, device=dev)
        adapted.set_deltas(hyper(z)); l_a, _ = adapted(x)
        adapted.set_deltas(hyper(z2)); l_b, _ = adapted(x)
    sep = (l_a - l_b).abs().max().item()
    assert sep > 1e-4, f"different z gave identical adapters: {sep}"
    print(f"[5/5] z-conditioning OK (two z differ by {sep:.3e})")

    print("ALL HYPER_LORA SELF-TESTS PASSED.")


if __name__ == "__main__":
    _self_test()