#!/usr/bin/env python3 """LoRA (Low-Rank Adaptation) implemented from scratch as a plain ``nn.Module``. No ``peft``. The whole point of this file is that the low-rank math is visible. THE PROBLEM ----------- Fine-tuning normally updates a pretrained weight matrix ``W`` directly. One DistilBERT attention projection is 768x768 = 589,824 parameters. Across 6 layers and 4 projections each, that is ~14M parameters just for attention, plus you must store an Adam optimizer state (2 extra floats) for every one of them. THE IDEA -------- Freeze ``W``. Learn a *low-rank correction* instead: W_effective = W + (alpha / r) * B @ A with ``A: (r, in_features)`` and ``B: (out_features, r)``. At ``r = 8`` that is ``8*768 + 768*8 = 12,288`` trainable parameters instead of 589,824 — about 2%. The bet being made is that the *change* a task needs is low-rank even though the pretrained weights are not. Empirically, for fine-tuning, that holds up well. TWO DETAILS THAT MATTER ----------------------- 1. The forward pass never builds ``W + BA``. It runs the two paths separately: y = xW^T + b + (alpha/r) * ((xA^T)B^T) Note the bracketing. ``(xA^T)B^T`` projects down to ``r`` dimensions first, so the intermediate tensor is tiny. ``x(BA)^T`` gives the same answer but materializes a full 768x768 matrix — exactly the thing we are avoiding. 2. ``B`` initializes to **zeros**, so ``BA = 0`` at step 0 and the adapted model is numerically identical to the pretrained one. Training departs smoothly from pretrained behaviour instead of starting from a random perturbation of it. (Both ``A`` and ``B`` zero would be useless — the gradient would be zero too, and they would never move. Exactly one of them must be random.) """ from __future__ import annotations import math from typing import Iterable, Sequence import torch import torch.nn as nn import torch.nn.functional as F __all__ = [ "LoRALinear", "inject_lora", "mark_only_lora_as_trainable", "count_parameters", "parameter_report", ] class LoRALinear(nn.Module): """A frozen ``nn.Linear`` plus a trainable low-rank update. Drop-in replacement: same input and output shapes as the layer it wraps, so it can be swapped into a pretrained model without touching anything else. Attributes: base: The original ``nn.Linear``, frozen (``requires_grad=False``). lora_A: ``(r, in_features)`` — the "down" projection. Randomly initialized. lora_B: ``(out_features, r)`` — the "up" projection. Zero-initialized. scaling: ``alpha / r``, applied to the low-rank path. """ def __init__( self, base_layer: nn.Linear, r: int = 8, alpha: int = 16, dropout: float = 0.0, ) -> None: """ Args: base_layer: The pretrained linear layer to adapt. Its weights are reused in place (not copied) and frozen. r: The rank. Higher = more capacity and more trainable parameters. 8 is the paper's default and works well; 1-4 is surprisingly competitive, 64+ starts to defeat the purpose. alpha: Scaling numerator. The low-rank path is multiplied by ``alpha / r``. Keeping ``alpha`` fixed while sweeping ``r`` means you do not have to retune the learning rate for each rank — that is the entire reason this parameter exists. dropout: Dropout applied to the *input* of the low-rank path only. Regularizes the adapter without disturbing the frozen path. """ super().__init__() if r <= 0: raise ValueError(f"LoRA rank must be positive, got r={r}") self.base = base_layer self.in_features = base_layer.in_features self.out_features = base_layer.out_features self.r = r self.alpha = alpha self.scaling = alpha / r # Freeze the pretrained weights. This is what makes the method # "parameter-efficient" — these tensors get no gradient and no optimizer # state, which is where most of the memory saving actually comes from. self.base.weight.requires_grad = False if self.base.bias is not None: self.base.bias.requires_grad = False # The two trainable matrices. nn.Parameter registers them with the module # so .parameters(), .to(device), and the optimizer all find them. self.lora_A = nn.Parameter(torch.empty(r, self.in_features)) self.lora_B = nn.Parameter(torch.zeros(self.out_features, r)) # Kaiming-uniform for A, matching how nn.Linear initializes its own # weight. B stays zero — see the module docstring for why. nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() def forward(self, x: torch.Tensor) -> torch.Tensor: """Run both paths and add them. Args: x: ``(..., in_features)``. In DistilBERT this is ``(batch, seq_len, 768)``. Returns: ``(..., out_features)`` — same shape the wrapped layer would produce. """ # Frozen pretrained path. No gradient flows into base.weight. base_out = self.base(x) # Low-rank path, evaluated right-to-left so the intermediate is (..., r): # x (..., 768) # @ A^T -> (..., r) <- the bottleneck; r is 8, not 768 # @ B^T -> (..., 768) dropped = self.lora_dropout(x) lora_out = F.linear(F.linear(dropped, self.lora_A), self.lora_B) return base_out + lora_out * self.scaling def extra_repr(self) -> str: """Make ``print(model)`` show the rank, so mis-set ranks are visible.""" return ( f"in_features={self.in_features}, out_features={self.out_features}, " f"r={self.r}, alpha={self.alpha}, scaling={self.scaling:.3f}" ) def inject_lora( model: nn.Module, target_names: Sequence[str] = ("q_lin", "v_lin"), r: int = 8, alpha: int = 16, dropout: float = 0.0, ) -> int: """Replace matching ``nn.Linear`` layers throughout ``model`` with ``LoRALinear``. Args: model: Any module tree — here, the DistilBERT encoder. target_names: Attribute names to adapt. DistilBERT's attention block names its projections ``q_lin``, ``k_lin``, ``v_lin``, ``out_lin``. The LoRA paper found query+value gives the best accuracy per parameter; adding keys costs more for little gain. r: Rank passed to each adapter. alpha: Scaling numerator passed to each adapter. dropout: Dropout passed to each adapter. Returns: How many layers were replaced. For DistilBERT with the default targets this is 12 (6 layers x {query, value}). **Check this number** — if it comes back 0, the layer names are wrong for your model and you have silently trained nothing. Raises: ValueError: If no layers matched, which is always a bug. """ # Collect first, mutate second. Calling setattr while iterating # named_modules() modifies the tree mid-walk, which skips entries. to_replace: list[tuple[nn.Module, str, nn.Linear]] = [] for parent in model.modules(): for child_name, child in parent.named_children(): if child_name in target_names and isinstance(child, nn.Linear): to_replace.append((parent, child_name, child)) for parent, child_name, child in to_replace: setattr(parent, child_name, LoRALinear(child, r=r, alpha=alpha, dropout=dropout)) if not to_replace: available = sorted( {name for _, name, _ in _iter_linear_children(model)} ) raise ValueError( f"inject_lora matched no layers for target_names={tuple(target_names)}. " f"Linear layer names present in this model: {available}" ) return len(to_replace) def _iter_linear_children(model: nn.Module) -> Iterable[tuple[nn.Module, str, nn.Linear]]: """Yield every ``(parent, attribute_name, layer)`` triple for ``nn.Linear`` children.""" for parent in model.modules(): for child_name, child in parent.named_children(): if isinstance(child, nn.Linear): yield parent, child_name, child def mark_only_lora_as_trainable( model: nn.Module, also_train: Sequence[str] = ("classifier",) ) -> None: """Freeze everything except the LoRA matrices and any named exceptions. ``LoRALinear.__init__`` already froze the base layers it wrapped, but the rest of the network — embeddings, feed-forward blocks, layer norms — is still trainable at this point. This freezes all of it in one pass. Args: model: The full classifier. also_train: Parameter-name prefixes to keep trainable. The classification head must be in here: it is randomly initialized, has no pretrained knowledge to preserve, and freezing random weights makes the task unlearnable. It is also tiny (768 -> 7), so it barely moves the headline parameter count. """ for name, param in model.named_parameters(): is_lora = "lora_A" in name or "lora_B" in name is_exception = any(name.startswith(prefix) for prefix in also_train) param.requires_grad = is_lora or is_exception def count_parameters(model: nn.Module) -> tuple[int, int]: """Return ``(trainable, total)`` parameter counts. ``requires_grad`` is the definition of "trainable" — it is exactly what the optimizer looks at, so this cannot drift from reality. """ total = sum(p.numel() for p in model.parameters()) trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) return trainable, total def parameter_report(model: nn.Module) -> dict[str, float | int]: """Summarize trainable vs. total parameters — the project's headline result.""" trainable, total = count_parameters(model) return { "trainable_params": trainable, "total_params": total, "trainable_pct": 100.0 * trainable / total if total else 0.0, "frozen_params": total - trainable, }