""" 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()