model/lora.py

"""
LoRA (Low-Rank Adaptation) — lightweight adapter injection for ArcisVLM v2.

Provides LoRA layers that can be dynamically injected into and removed from
the MoE decoder's attention projections at inference time. The HyperMother
generates LoRA weights per-camera/per-scene via the HyperNetwork, and this
module handles the injection mechanics.

Key design decisions:
  - Target the fused QKV projection, applying LoRA only to Q and V slices
    (NOT K, following LoRA best practices from Hu et al. 2021)
  - NOT applied to MoE FFN gating — the learned router was trained across
    Stages 2-3 and disrupting it with LoRA would break load-balanced routing
  - LoRA modifies what the model *attends to*; MoE controls how it *processes*

References:
  - LoRA: Low-Rank Adaptation (Hu et al., 2021)
  - HypeLoRA: Calibrated adapters with uncertainty (arXiv: 2603.19278)
  - Doc-to-LoRA / Sakana AI (arXiv: 2602.15902)
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn


@dataclass
class LoRAConfig:
    """Configuration for LoRA adapter injection."""
    rank: int = 16           # Low-rank dimension (r)
    alpha: float = 32.0      # Scaling factor (α)
    dropout: float = 0.05    # Dropout on LoRA path
    targets: tuple[str, ...] = ("q", "v")  # Which projections to adapt (subset of q, k, v)


class LoRALayer(nn.Module):
    """
    Low-Rank Adaptation layer: W' = W + (α/r) * B @ A

    A is initialized with Kaiming uniform, B with zeros, so the LoRA
    contribution starts as zero (identity behavior).

    Args:
        in_dim: Input dimension of the original projection
        out_dim: Output dimension of the original projection
        rank: Low-rank bottleneck dimension
        alpha: Scaling factor
        dropout: Dropout rate on the LoRA path
    """

    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        rank: int = 16,
        alpha: float = 32.0,
        dropout: float = 0.05,
    ):
        super().__init__()
        self.rank = rank
        self.scaling = alpha / rank

        # Down-projection: in_dim → rank
        self.A = nn.Parameter(torch.empty(rank, in_dim))
        # Up-projection: rank → out_dim
        self.B = nn.Parameter(torch.empty(out_dim, rank))

        self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()

        self.reset_parameters()

    def reset_parameters(self) -> None:
        """Initialize A with Kaiming uniform, B with zeros (start as identity)."""
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
        nn.init.zeros_(self.B)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Compute LoRA delta: (α/r) * x @ A^T @ B^T

        Args:
            x: [..., in_dim] — input tensor (any leading batch dims)
        Returns:
            [..., out_dim] — LoRA contribution to add to original output
        """
        return (self.dropout(x) @ self.A.T @ self.B.T) * self.scaling

    @property
    def num_params(self) -> int:
        """Total number of parameters in this LoRA layer."""
        return self.A.numel() + self.B.numel()

    @classmethod
    def from_flat_params(
        cls,
        flat_params: torch.Tensor,
        in_dim: int,
        out_dim: int,
        rank: int = 16,
        alpha: float = 32.0,
        dropout: float = 0.05,
    ) -> "LoRALayer":
        """
        Create a LoRALayer from a flat parameter vector (output of HyperNetwork).

        Args:
            flat_params: [rank * in_dim + out_dim * rank] — concatenated A and B
            in_dim, out_dim, rank, alpha, dropout: LoRA configuration
        Returns:
            LoRALayer with parameters set from flat_params
        """
        layer = cls(in_dim, out_dim, rank, alpha, dropout)
        a_size = rank * in_dim
        b_size = out_dim * rank
        assert flat_params.numel() == a_size + b_size, (
            f"Expected {a_size + b_size} params, got {flat_params.numel()}"
        )
        with torch.no_grad():
            layer.A.copy_(flat_params[:a_size].reshape(rank, in_dim))
            layer.B.copy_(flat_params[a_size:a_size + b_size].reshape(out_dim, rank))
        return layer


class LoRAInjector:
    """
    Manages LoRA injection into the MoE decoder's attention layers.

    Handles the mapping from a flat parameter vector (HyperNetwork output)
    to individual LoRA layers across all decoder transformer blocks,
    targeting Q and V projections in the fused QKV linear.

    The fused QKV layout is: [Q_params | K_params | V_params] where each
    segment is embed_dim wide. LoRA is applied as a residual on the Q and V
    slices of the output.
    """

    def __init__(self, config: LoRAConfig, num_blocks: int, embed_dim: int):
        self.config = config
        self.num_blocks = num_blocks
        self.embed_dim = embed_dim

        # Calculate total LoRA params needed
        self._params_per_layer = self._calc_params_per_layer()
        self._num_lora_layers = num_blocks * len(config.targets)
        self.total_params = self._params_per_layer * self._num_lora_layers

    def _calc_params_per_layer(self) -> int:
        """Params for one LoRA layer: rank * in_dim + out_dim * rank."""
        r = self.config.rank
        d = self.embed_dim
        return r * d + d * r  # A: [r, d] + B: [d, r]

    def create_lora_layers(
        self,
        flat_params: Optional[torch.Tensor] = None,
        device: Optional[torch.device] = None,
    ) -> list[dict[str, LoRALayer]]:
        """
        Create LoRA layers for all decoder blocks.

        Args:
            flat_params: Optional flat parameter vector from HyperNetwork.
                         If None, creates randomly initialized LoRA layers.
            device: Target device for the layers
        Returns:
            List of dicts, one per block, mapping target name → LoRALayer
            e.g., [{"q": LoRALayer, "v": LoRALayer}, ...]
        """
        layers = []
        offset = 0

        for block_idx in range(self.num_blocks):
            block_layers = {}
            for target in self.config.targets:
                if flat_params is not None:
                    chunk = flat_params[offset:offset + self._params_per_layer]
                    lora = LoRALayer.from_flat_params(
                        chunk,
                        in_dim=self.embed_dim,
                        out_dim=self.embed_dim,
                        rank=self.config.rank,
                        alpha=self.config.alpha,
                        dropout=self.config.dropout,
                    )
                    offset += self._params_per_layer
                else:
                    lora = LoRALayer(
                        in_dim=self.embed_dim,
                        out_dim=self.embed_dim,
                        rank=self.config.rank,
                        alpha=self.config.alpha,
                        dropout=self.config.dropout,
                    )
                if device is not None:
                    lora = lora.to(device)
                block_layers[target] = lora
            layers.append(block_layers)

        return layers


def compute_total_lora_params(
    num_blocks: int,
    embed_dim: int,
    rank: int = 16,
    targets: tuple[str, ...] = ("q", "v"),
) -> int:
    """
    Calculate the total number of LoRA parameters for the decoder.

    Useful for sizing the HyperNetwork output layer.

    Args:
        num_blocks: Number of transformer blocks in decoder
        embed_dim: Model dimension
        rank: LoRA rank
        targets: Which projections get LoRA (default: Q and V)
    Returns:
        Total parameter count
    """
    params_per_layer = rank * embed_dim + embed_dim * rank  # A + B
    return params_per_layer * num_blocks * len(targets)