models/lora.py

#!/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,
    }