#!/usr/bin/env python3 """Shared building blocks for the Memory-LoRA hypernetwork. Forked from Code2LoRA's ``hypernetwork/code2lora_core.py`` (Hotsko et al., "Code2LoRA: Hypernetwork-Generated Adapters for Code Language Models under Software Evolution", MIT-licensed code release). Same core trick, different target model and conditioning input: * Code2LoRA: repository embedding -> LoRA adapter for Qwen2.5-Coder-1.5B * Memory-LoRA: document embedding -> LoRA adapter for google/gemma-4-E2B The ``MemoryLoRAHead`` (renamed from ``Code2LoRAHead``, architecture unchanged) outputs ONE (A, B) pair per LoRA module *type* (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj), shared across every target transformer layer -- not per-layer. This keeps the head's parameter count tractable for local (MPS) training. Gemma-4-E2B specifics (verified against the model's actual safetensors header, not guessed): * Decoder is nested at ``model.language_model.layers.{i}.*`` -- NOT ``model.layers.{i}.*`` like Qwen2.5-Coder. ``get_module_specs`` below matches on ``language_model\\.layers\\.(\\d+)\\.``, not ``model\\.layers``. * Module type names are identical to Code2LoRA's defaults: q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj. * Only ``model.language_model.*`` is ever touched. ``vision_tower`` and ``audio_tower`` are left completely alone -- irrelevant to text recall and risky to perturb. """ from __future__ import annotations import math import re from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import numpy as np import pyarrow as pa import pyarrow.compute as pc import pyarrow.dataset as pads import torch import torch.nn as nn import torch.nn.functional as F # --------------------------------------------------------------------------- # LoRA module + injection (unchanged from Code2LoRA -- architecture-agnostic) # --------------------------------------------------------------------------- class LoRA(nn.Module): """Wraps an ``nn.Linear`` with an additive low-rank update. Forward: ``y = base(x) + scaling * (x @ A^T) @ B^T``, where per-batch A: ``[rank, in_features]`` and B: ``[out_features, rank]`` come from an external hypernet via :meth:`set_lora_weights`. IMPORTANT autograd contract: A and B are kept as **plain attributes**, not buffers, and stored **without detaching**, so the LM loss's backward graph flows through them straight into the hypernet parameters that produced them. The base ``nn.Linear`` is frozen and its forward sees a detached copy of the input to avoid building an autograd graph through the (much larger) frozen LLM weights. """ def __init__(self, base: nn.Linear, in_features: int, out_features: int, rank: int, alpha: float): super().__init__() self.base = base for p in self.base.parameters(): p.requires_grad = False self.in_features = in_features self.out_features = out_features self.rank = rank self.scaling = float(alpha) / float(max(1, rank)) self.A: Optional[torch.Tensor] = None # [rank, in_features] self.B: Optional[torch.Tensor] = None # [out_features, rank] def set_lora_weights(self, A: torch.Tensor, B: torch.Tensor) -> None: self.A = A self.B = B def forward(self, x: torch.Tensor) -> torch.Tensor: y = self.base(x) if self.A is None or self.B is None: return y x_f32 = x.detach().to(torch.float32) A = self.A.to(torch.float32) B = self.B.to(torch.float32) delta = F.linear(F.linear(x_f32, A), B) * self.scaling return y + delta.to(dtype=y.dtype) @dataclass class ModuleSpec: full_name: str # e.g. 'model.language_model.layers.5.self_attn.q_proj' layer_idx: int type: str # e.g. 'q_proj' in_features: int out_features: int # Gemma-4-E2B nests its text decoder here (verified via safetensors header). DEFAULT_ROOT_PREFIX = "model.language_model." _LAYER_IDX_RE = re.compile(r"\blanguage_model\.layers\.(\d+)\.") def get_module_specs(model: nn.Module, target_module_types: List[str], root_prefix: str = DEFAULT_ROOT_PREFIX ) -> List[ModuleSpec]: """Discover every nn.Linear under ``root_prefix`` whose name contains one of ``target_module_types`` and return one :class:`ModuleSpec` per match, sorted by (layer_idx, full_name). Restricting to ``root_prefix`` is what keeps ``vision_tower`` / ``audio_tower`` untouched even though they also contain q_proj/k_proj/ v_proj/o_proj-named linears. Gemma-4-E2B is architecturally heterogeneous across layers (unlike Qwen2.5-Coder, which Code2LoRA was built for): every 5th layer is a wider "full_attention" layer (q_proj/o_proj 2x the width of the "sliding_attention" layers), and 20 of the 35 layers have NO k_proj/ v_proj at all -- they reuse an earlier layer's KV cache (``num_kv_shared_layers=20`` in the model config). A LoRA (A, B) pair can only be shared across modules of IDENTICAL shape, so ``.type`` here is ``"{module_name}_{in}x{out}"`` (shape-qualified), not just the raw module name -- e.g. ``"q_proj_1536x2048"`` vs ``"q_proj_1536x4096"`` end up as distinct hypernetwork output heads. Layers with no matching module (e.g. k_proj on a KV-shared layer) simply produce no spec for that layer, which is architecturally correct: there is nothing to adapt there since that layer never computes its own K/V. """ specs: List[ModuleSpec] = [] for name, m in model.named_modules(): if root_prefix and not name.startswith(root_prefix): continue match_type = next( (t for t in target_module_types if t in name), None ) if match_type is None: continue if not isinstance(m, nn.Linear): continue m_layer = _LAYER_IDX_RE.search(name) layer_idx = int(m_layer.group(1)) if m_layer else -1 shape_qualified_type = f"{match_type}_{m.in_features}x{m.out_features}" specs.append(ModuleSpec( full_name=name, layer_idx=layer_idx, type=shape_qualified_type, in_features=int(m.in_features), out_features=int(m.out_features), )) specs.sort(key=lambda s: (s.layer_idx, s.full_name)) return specs def replace_with_lora(model: nn.Module, specs: List[ModuleSpec], rank: int, alpha: float) -> None: """Replace each target ``nn.Linear`` in ``model`` with a :class:`LoRA` wrapper. Idempotent.""" named = dict(model.named_modules()) device = next(model.parameters()).device dtype = next(model.parameters()).dtype for sp in specs: parent_name, attr = sp.full_name.rsplit(".", 1) orig = getattr(named[parent_name], attr) if isinstance(orig, LoRA): continue assert isinstance(orig, nn.Linear), \ f"{sp.full_name} is not nn.Linear (got {type(orig)})" wrapped = LoRA(orig, sp.in_features, sp.out_features, rank, alpha).to(device=device, dtype=dtype) setattr(named[parent_name], attr, wrapped) def inject_lora_weights(model: nn.Module, specs: List[ModuleSpec], head_out: Dict[str, Dict[str, torch.Tensor]], batch_index: int = 0) -> None: """Push ``head_out["A"][type]`` and ``head_out["B"][type]`` into the wrapper :class:`LoRA` modules for every spec sharing that type.""" A_by_type = head_out["A"] B_by_type = head_out["B"] named = dict(model.named_modules()) for sp in specs: named[sp.full_name].set_lora_weights( A_by_type[sp.type][batch_index], B_by_type[sp.type][batch_index], ) def discover_module_types_and_dims(specs: List[ModuleSpec] ) -> Dict[str, Tuple[int, int]]: """Return {type_name: (in_features, out_features)} -- one entry per target module type. Assumes all instances of the same type share dims.""" type_dims: Dict[str, Tuple[int, int]] = {} for sp in specs: if sp.type in type_dims: assert type_dims[sp.type] == (sp.in_features, sp.out_features), \ f"type {sp.type} appears with inconsistent dims" continue type_dims[sp.type] = (sp.in_features, sp.out_features) return type_dims # --------------------------------------------------------------------------- # Shared LoRA generation head (= Code2LoRAHead, renamed; architecture unchanged) # --------------------------------------------------------------------------- class MemoryLoRAHead(nn.Module): """Maps a document-context embedding to a LoRA adapter in one forward pass. Input : ctx ``[B, input_dim]`` -- a single document embedding. Output : ``{"A": {type: [B, rank, in_f]}, "B": {type: [B, out_f, rank]}}``, one (A, B) pair per LoRA module *type*, shared across all target transformer layers. Args: input_dim : Context-vector dim (2048, matches the Qwen3-Embedding weighted-mean + max-pool concat from ``encoder.py``). type_dims : ``{type: (in_features, out_features)}`` for each LoRA module type (q_proj, v_proj, gate_proj, ...). hidden_dim : Trunk hidden dimension. Default 128 -- deliberately small: with only ~165 training documents (~3K QA pairs), a 745M-param head (hidden_dim=512, the original default) overfits within ~2 epochs (train loss -> 0.4 while held-out cr_val/cr_test loss rises from ~1.9 to ~2.7). hidden_dim=128 cuts head size roughly 4x; combine with --head-dropout and higher weight decay for further regularization. rank : LoRA rank ``r``. init_log_scale : Initial log-scale for tanh squashing. -3.5 gives output magnitudes ~0.03 at init -> tiny LoRA delta. dropout : Dropout applied after each trunk GELU. 0.0 (paper's original setting) had no regularization at all; nonzero here specifically to counter the overfitting observed on this project's much smaller corpus. """ def __init__( self, input_dim: int, type_dims: Dict[str, Tuple[int, int]], hidden_dim: int = 128, rank: int = 16, init_log_scale: float = -3.5, dropout: float = 0.1, ): super().__init__() self.input_dim = input_dim self.hidden_dim = hidden_dim self.rank = rank self.dropout = dropout self.type_dims = dict(type_dims) self.types = sorted(type_dims.keys()) # Input standardization statistics (set via :meth:`fit_input_stats`). # # WHY THIS EXISTS: the 6-view repo embedding is ~64% a constant vector # shared by every repo (the frozen encoder's mean response to "source # code"), so raw embeddings have mean pairwise cosine ~0.73 -- and the # trunk, dominated by that DC component, mapped them to cosine ~0.98, # i.e. it emitted essentially the SAME adapter for every repository # (measured: emitted-delta cosine 0.96 across 24 unrelated repos). # A near-constant, non-trivially-sized delta is pure damage: it scored # WORSE than random noise of matched scale on held-out repo text. # Centering removes that DC term (pairwise cosine ~0.00), which is what # makes the conditioning signal actually reach the output heads. # Buffers (not parameters) so they persist in the checkpoint and are # applied identically at training and inference time. self.register_buffer("input_mean", torch.zeros(input_dim)) self.register_buffer("input_std", torch.ones(input_dim)) self.trunk = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Dropout(dropout), ) self.heads_A = nn.ModuleDict({ t: nn.Linear(hidden_dim, rank * type_dims[t][0]) for t in self.types }) self.heads_B = nn.ModuleDict({ t: nn.Linear(hidden_dim, type_dims[t][1] * rank) for t in self.types }) self.log_scale_A = nn.ParameterDict({ t: nn.Parameter(torch.tensor(init_log_scale)) for t in self.types }) self.log_scale_B = nn.ParameterDict({ t: nn.Parameter(torch.tensor(init_log_scale)) for t in self.types }) @torch.no_grad() def fit_input_stats(self, embeddings: torch.Tensor) -> None: """Set the standardization buffers from the TRAINING embeddings only. Must be called once before training and never refit afterwards -- inference has to apply exactly the same transform, which is why the stats travel inside the checkpoint. """ e = embeddings.float() self.input_mean.copy_(e.mean(0)) # Guard against near-constant dimensions blowing up when divided. self.input_std.copy_(e.std(0).clamp_min(1e-3)) def forward(self, ctx: torch.Tensor) -> Dict[str, Dict[str, torch.Tensor]]: if ctx.dim() == 3: ctx = torch.max(ctx, dim=1).values ctx = (ctx.float() - self.input_mean) / self.input_std h = self.trunk(ctx.float()) h = F.normalize(h, p=2, dim=-1) * math.sqrt(self.hidden_dim) A_out: Dict[str, torch.Tensor] = {} B_out: Dict[str, torch.Tensor] = {} for t in self.types: in_f, out_f = self.type_dims[t] A_raw = self.heads_A[t](h).view(-1, self.rank, in_f) B_raw = self.heads_B[t](h).view(-1, out_f, self.rank) scale_A = torch.exp(self.log_scale_A[t]).clamp(1e-5, 0.3) scale_B = torch.exp(self.log_scale_B[t]).clamp(1e-5, 0.3) A_out[t] = torch.tanh(A_raw) * scale_A B_out[t] = torch.tanh(B_raw) * scale_B return {"A": A_out, "B": B_out} def config_dict(self) -> Dict[str, Any]: return { "input_dim": self.input_dim, "hidden_dim": self.hidden_dim, "rank": self.rank, "dropout": self.dropout, "types": self.types, "type_dims": {t: list(v) for t, v in self.type_dims.items()}, } # --------------------------------------------------------------------------- # Parquet loaders -- documents + recall QA pairs # --------------------------------------------------------------------------- def _list_to_f32_array(col, dim: int) -> np.ndarray: """Vectorized fixed-width-list -> ndarray. The naive per-row Python loop (``for i, v in enumerate(col.to_pylist()): out[i] = v``) does dim * n_rows individual scalar assignments in pure Python -- fine at ~200 rows, but at real-corpus scale (74K rows x 2048 dims = 151M scalar ops) it single-handedly took 5+ minutes just to load embeddings before training could even start. pyarrow's own flatten() + numpy reshape does the same conversion in C. """ if len(col) == 0: return np.zeros((0, dim), dtype=np.float32) flat = col.combine_chunks().flatten() if hasattr(col, "combine_chunks") else col.flatten() arr = flat.to_numpy(zero_copy_only=False).astype(np.float32, copy=False) return arr.reshape(len(col), dim) @dataclass class DocRow: doc_id: str doc_version: str # constant "v1" for static (non-evolving) docs split: str # "train" | "cr_val" | "cr_test" (cross-corpus) doc_embedding: np.ndarray # fp32 [2048] @dataclass class QnaRow: doc_id: str doc_version: str split: str # cross-corpus split, inherited from DocRow qna_split: str # "train" | "held_out" (in-corpus split) question: str prefix: str target: str def load_doc_rows(parquet_path: Path, splits: Optional[List[str]] = None, embedding_col: str = "doc_embedding", ) -> List[DocRow]: needed = ["doc_id", "doc_version", "split", embedding_col] ds = pads.dataset(str(parquet_path), format="parquet") flt = None if splits: flt = pc.is_in(pads.field("split"), value_set=pa.array(splits, type=pa.string())) table = ds.to_table(columns=needed, filter=flt) n = table.num_rows if n == 0: return [] dim = len(table.column(embedding_col)[0].as_py()) embs = _list_to_f32_array(table.column(embedding_col), dim) doc_col = table.column("doc_id").to_pylist() ver_col = table.column("doc_version").to_pylist() split_col = table.column("split").to_pylist() rows: List[DocRow] = [] for i in range(n): rows.append(DocRow( doc_id=doc_col[i], doc_version=ver_col[i], split=split_col[i] or "", doc_embedding=embs[i], )) return rows def load_qna_rows(jsonl_path: Path, splits: Optional[List[str]] = None, qna_splits: Optional[List[str]] = None, doc_ids: Optional[List[str]] = None, ) -> List[QnaRow]: """QnA pairs are written as JSONL by ``generate_synthetic_dataset.py`` (one row per line, cheap to append incrementally during generation) -- unlike doc embeddings, which are batch-written parquet. Filters are applied in Python; at the scale of this project (low thousands of rows) that's simpler and fast enough.""" import json as _json splits_set = set(splits) if splits else None qna_splits_set = set(qna_splits) if qna_splits else None doc_ids_set = set(doc_ids) if doc_ids else None rows: List[QnaRow] = [] with open(jsonl_path) as f: for line in f: line = line.strip() if not line: continue d = _json.loads(line) if splits_set and d.get("split") not in splits_set: continue if qna_splits_set and d.get("qna_split") not in qna_splits_set: continue if doc_ids_set and d.get("doc_id") not in doc_ids_set: continue rows.append(QnaRow( doc_id=d.get("doc_id", ""), doc_version=d.get("doc_version", "v1"), split=d.get("split", ""), qna_split=d.get("qna_split", ""), question=d.get("question", ""), prefix=d.get("prefix", ""), target=d.get("target", ""), )) return rows __all__ = [ "LoRA", "ModuleSpec", "DEFAULT_ROOT_PREFIX", "get_module_specs", "replace_with_lora", "inject_lora_weights", "discover_module_types_and_dims", "MemoryLoRAHead", "DocRow", "QnaRow", "load_doc_rows", "load_qna_rows", ]