LoRA escrita a mano

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import TYPE_CHECKING

import torch
from pydantic import Field
from torch import Tensor, nn

from rukh.config import BaseConfig

if TYPE_CHECKING:
    from collections.abc import Iterator

__all__ = [
    "TARGETS",
    "AdapterLayout",
    "LoraConfig",
    "LoRALinear",
    "adapter_layout",
    "adapter_params",
    "apply_lora",
    "load_adapter",
    "lora_modules",
    "merge_lora",
    "merged_state_dict",
    "save_adapter",
    "stack_adapter",
]

ADAPTER_FILE = "adapter.safetensors"
ADAPTER_CONFIG = "adapter_config.json"

        def save(path: Path, step: int) -> Path:
            return write_checkpoint(
                path,
                step=step,
                model=model,
                # A LoRA run trains `A` and `B`; the optimizer state of two small matrices is not
                # worth 460 MB per checkpoint, and it cannot be resumed into anyway.
                optimizer=None if cfg.lora is not None else optimizer,
                cfg=cfg,
                model_cfg=model_cfg,
                vocab_hash=vocab_hash,
                manifest_sha=manifest_sha,
                best_val=best_val,
                run_id=this_run,
                state=merged_state_dict(model) if cfg.lora is not None else None,
            )

    *,
    step: int,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None,
    cfg: Mapping[str, Any],
    model_cfg: Mapping[str, Any],
    vocab_hash: str | None = None,
    data_manifest_sha: str | None = None,
    git_sha: str | None = None,
    best_val: float | None = None,
    run_id: str | None = None,
    state: Mapping[str, torch.Tensor] | None = None,
) -> Path:
    """Write one checkpoint atomically (temporary file plus replace) and return its path.

    ``run_id`` is the MLflow run that wrote it, so ``--resume`` can carry on logging into the
    same run and the loss curve stays one line instead of two.

    ``state`` replaces ``model.state_dict()``. A LoRA run passes the merged weights, so what
    lands on disk is an ordinary decoder checkpoint that every loader, exporter and publisher
    already reads, instead of a wrapped model only this repository could open.
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    source = state if state is not None else model.state_dict()
    payload: dict[str, Any] = {
        "step": int(step),
        "model_state": {k: v.detach().cpu() for k, v in source.items()},

def save_adapter(model: nn.Module, path: Path | str, cfg: LoraConfig) -> Path:
    """Write ``A`` and ``B`` to safetensors, keyed by the module they adapt, plus the config.

    The keys are the paths the adapter *would* have in the base model
    (``blocks.0.attn.qkv.lora_A.0``), not the paths the wrappers create, so loading does not
    depend on how the wrapping happened to be spelled.
    """
    from safetensors.torch import save_file

    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    tensors: dict[str, Tensor] = {}
    for name, adapter in lora_modules(model):
        for index, (a, b) in enumerate(zip(adapter.a, adapter.b, strict=True)):
            tensors[f"{name}.lora_A.{index}"] = a.detach().cpu().contiguous()
            tensors[f"{name}.lora_B.{index}"] = b.detach().cpu().contiguous()
    if not tensors:
        raise ValueError("model has no LoRA adapters to save")
    save_file(tensors, target.as_posix())
    config_path = target.with_name(ADAPTER_CONFIG)
    config_path.write_text(cfg.model_dump_json(indent=2) + "\n", encoding="utf-8")
    return target

def load_adapter(model: nn.Module, path: Path | str) -> LoraConfig:
    """Apply the adapter stored next to ``path`` to a base model and load its weights."""
    from safetensors.torch import load_file

    target = Path(path)
    cfg = LoraConfig.model_validate_json(
        target.with_name(ADAPTER_CONFIG).read_text(encoding="utf-8")
    )
    apply_lora(model, cfg)
    tensors = load_file(target.as_posix())
    state: dict[str, Tensor] = {}
    for name, adapter in lora_modules(model):
        for index in range(len(adapter.a)):
            state[f"{name}.a.{index}"] = tensors[f"{name}.lora_A.{index}"]
            state[f"{name}.b.{index}"] = tensors[f"{name}.lora_B.{index}"]
    missing, unexpected = model.load_state_dict(state, strict=False)
    if unexpected:
        raise ValueError(f"adapter carries tensors the model has no place for: {unexpected}")
    return cfg

class AdapterLayout(BaseConfig):
    """How a whole adapter is laid out as two stacked tensors, one per factor.

    ``save_adapter`` writes one ``A`` and one ``B`` per adapted slice, keyed by the module they
    belong to. That is the right shape to load into PyTorch and the wrong shape to feed a graph:
    an ONNX input is one tensor, not thirty-two. Stacking turns the whole adapter into
    ``A (layers, slices, r, in_features)`` and ``B (layers, slices, width, r)``, which is what
    the browser uploads once per move and what lets a 1.6 MB file change a 440 MB model without
    downloading a second one.

    The stacking needs the layers to agree: one adapted module per block, every slice the same
    width, the same input width everywhere. That covers the query, key and value ranges of the
    fused ``qkv`` -- which is what this project trains and publishes -- and it does not cover an
    adapter that also touches the MLP, whose ``fc`` is four times as wide. That case has a path
    of its own (``merge_lora`` and an ordinary export); it just cannot travel as one tensor.
    """

    module: str
    """Dotted path of the adapted matrix inside a block, e.g. ``attn.qkv``."""
    layers: int
    slices: tuple[tuple[int, int], ...]
    r: int
    scale: float
    in_features: int

    @property
    def width(self) -> int:
        return self.slices[0][1] - self.slices[0][0]

    @property
    def shape_a(self) -> tuple[int, int, int, int]:
        return (self.layers, len(self.slices), self.r, self.in_features)

    @property
    def shape_b(self) -> tuple[int, int, int, int]:
        return (self.layers, len(self.slices), self.width, self.r)

    @property
    def numel(self) -> int:
        return math.prod(self.shape_a) + math.prod(self.shape_b)

def adapter_layout(model: nn.Module) -> AdapterLayout:
    """Read the layout off a model that already has its adapters applied."""
    found = list(lora_modules(model))
    if not found:
        raise ValueError("model has no LoRA adapters")
    names = [name for name, _ in found]
    suffixes = {name.split(".", 2)[-1] for name in names}
    if len(suffixes) != 1:
        raise ValueError(
            f"adapters on more than one matrix per block ({sorted(suffixes)}) cannot be stacked "
            "into one tensor; merge them instead"
        )
    first = found[0][1]
    for name, adapter in found:
        if adapter.slices != first.slices or adapter.base.in_features != first.base.in_features:
            raise ValueError(f"{name} does not have the same shape as {names[0]}")
    widths = {stop - start for start, stop in first.slices}
    if len(widths) != 1:
        raise ValueError(f"slices of different widths ({sorted(widths)}) cannot be stacked")
    return AdapterLayout(
        module=suffixes.pop(),
        layers=len(found),
        slices=first.slices,
        r=first.r,
        scale=first.scale,
        in_features=first.base.in_features,
    )

@torch.no_grad()
def stack_adapter(model: nn.Module) -> tuple[Tensor, Tensor]:
    """The whole adapter as ``(A, B)``, stacked layer by layer in the order the blocks run."""
    layout = adapter_layout(model)
    a = torch.stack([torch.stack([p.detach() for p in ad.a]) for _, ad in lora_modules(model)])
    b = torch.stack([torch.stack([p.detach() for p in ad.b]) for _, ad in lora_modules(model)])
    if a.shape != layout.shape_a or b.shape != layout.shape_b:  # pragma: no cover - guarded above
        raise ValueError(f"stacked {tuple(a.shape)}/{tuple(b.shape)} against {layout}")
    return a.cpu().float().contiguous(), b.cpu().float().contiguous()


def adapter_size(model: nn.Module) -> tuple[int, int]:
    """``(adapter parameters, base parameters)``: the ratio the whole method exists for."""
    adapter = sum(p.numel() for p in adapter_params(model))
    total = sum(p.numel() for p in model.parameters())
    return adapter, total - adapter

def describe(model: nn.Module, cfg: LoraConfig) -> str:
    """One line for the run log: rank, targets and how little of the model is being trained."""
    adapter, base = adapter_size(model)
    share = 100.0 * adapter / (adapter + base)
    targets = ", ".join(cfg.targets)
    return (
        f"LoRA r={cfg.r} alpha={cfg.alpha} on [{targets}]: "
        f"{adapter:,} trainable of {adapter + base:,} ({share:.2f} %)"
    )


def adapter_json(cfg: LoraConfig, base_repo: str, extra: dict[str, object] | None = None) -> str:
    """The JSON that travels with a published adapter: what it is and what it mounts on."""
    payload: dict[str, object] = {"lora": cfg.model_dump(mode="json"), "base_model": base_repo}
    if extra:
        payload.update(extra)
    return json.dumps(payload, indent=2) + "\n"

    lora: LoraConfig | None = None
    """Train low-rank adapters instead of the weights.

    With it on, everything outside the adapters is frozen, the optimizer only ever sees ``A`` and
    ``B``, and the checkpoints hold the **merged** weights so the rest of the toolchain never
    learns that LoRA happened. The adapter itself is written next to them as its own small file,
    which is the artefact worth publishing."""

    if cfg.lora is not None:
        if resume is not None:
            raise ValueError(
                "--resume cannot continue a LoRA run: its checkpoints hold merged weights, "
                "which no longer say where the adapter ended. Start it again with init_from."
            )
        apply_lora(model, cfg.lora)
        log.info("%s", describe(model, cfg.lora))

        if cfg.lora is not None:
            # The artefact worth publishing: a couple of megabytes that mount on weights the
            # reader already has, next to the merged checkpoint the rest of the toolchain reads.
            adapter = save_adapter(model, out_dir / ADAPTER_FILE, cfg.lora)
            log.info("adapter written to %s (%.1f MB)", adapter, adapter.stat().st_size / 1e6)

"""Tests for rukh.models.lora: the hand-written adapters and what makes them LoRA.

The point of writing LoRA instead of importing it is that every claim the method makes can be
checked here: it starts as the identity, it trains only the factors, folding it back gives the
same model, and the file it produces is three orders of magnitude smaller than the weights.
"""

from __future__ import annotations

from pathlib import Path

import pytest
import torch

from rukh.models import DecoderConfig, MoveDecoder, preset
from rukh.models.lora import (
    LoraConfig,
    LoRALinear,
    adapter_params,
    apply_lora,
    load_adapter,
    lora_modules,
    merge_lora,
    save_adapter,
)

pytestmark = pytest.mark.unit


def _model(seed: int = 0) -> MoveDecoder:
    torch.manual_seed(seed)
    return MoveDecoder(DecoderConfig(n_layer=2, n_head=2, d_model=16, block=8, dropout=0.0)).eval()


def _idx() -> torch.Tensor:
    torch.manual_seed(1234)
    return torch.randint(1, 100, (2, 8))

def test_an_untrained_adapter_is_the_identity() -> None:
    """`B` starts at zero, so applying LoRA changes nothing until a gradient arrives.

    Any other initialisation would move the model before it learned anything, and every
    comparison against the base would start from a model that is already different.
    """
    idx = _idx()
    model = _model()
    with torch.no_grad():
        before, _ = model(idx)
    apply_lora(model, LoraConfig(r=4))
    with torch.no_grad():
        after, _ = model(idx)
    torch.testing.assert_close(before, after)


def test_only_the_factors_are_trainable() -> None:
    model = _model()
    trainable = apply_lora(model, LoraConfig(r=4, targets=("q", "v")))
    by_grad = [p for p in model.parameters() if p.requires_grad]
    assert trainable == sum(p.numel() for p in by_grad)
    assert {id(p) for p in by_grad} == {id(p) for p in adapter_params(model)}

def test_trainable_count_is_the_formula_on_the_card() -> None:
    """`2 * r * d_model` per adapted `d_model`-wide matrix, per layer, and nothing else."""
    cfg = LoraConfig(r=4, targets=("q", "v"))
    model = _model()
    trainable = apply_lora(model, cfg)
    n_layer, d_model = model.cfg.n_layer, model.cfg.d_model
    assert trainable == n_layer * len(cfg.targets) * 2 * cfg.r * d_model


def test_the_fused_projection_gives_each_target_its_own_factors() -> None:
    """Query and value are two slices of one `qkv` matrix, not one adapter over both.

    Sharing a single `A` between them would be cheaper and would not be LoRA: each adapted
    matrix is supposed to get its own low-rank subspace.
    """
    model = _model()
    apply_lora(model, LoraConfig(r=4, targets=("q", "v")))
    adapters = dict(lora_modules(model))
    assert set(adapters) == {"blocks.0.attn.qkv", "blocks.1.attn.qkv"}
    qkv = adapters["blocks.0.attn.qkv"]
    d = model.cfg.d_model
    assert qkv.slices == ((0, d), (2 * d, 3 * d))  # query and value, key untouched
    assert len(qkv.a) == len(qkv.b) == 2
    assert not torch.equal(qkv.a[0], qkv.a[1])


def test_the_key_slice_never_moves_when_only_q_and_v_are_adapted() -> None:
    model = _model()
    apply_lora(model, LoraConfig(r=4, targets=("q", "v")))
    d = model.cfg.d_model
    for _, adapter in lora_modules(model):
        with torch.no_grad():
            for b in adapter.b:
                b.normal_()
        weight = adapter.merged_weight()
        torch.testing.assert_close(weight[d : 2 * d], adapter.base.weight[d : 2 * d])
        assert not torch.allclose(weight[:d], adapter.base.weight[:d])

def test_merging_reproduces_the_adapted_model_exactly() -> None:
    """`W + BA` folded into the weight has to give the same logits as the wrapper computing it.

    This is the property that lets an adapter be published as a small file and still be
    exported, quantised and served as an ordinary model.
    """
    idx = _idx()
    model = _model()
    apply_lora(model, LoraConfig(r=4, alpha=8, targets=("q", "v")))
    torch.manual_seed(7)
    with torch.no_grad():
        for _, adapter in lora_modules(model):
            for b in adapter.b:
                b.normal_(std=0.05)
        adapted, _ = model(idx)
    assert merge_lora(model) == model.cfg.n_layer
    assert not list(lora_modules(model))
    with torch.no_grad():
        merged, _ = model(idx)
    torch.testing.assert_close(adapted, merged, rtol=1e-5, atol=1e-6)


def test_a_merged_model_has_the_state_dict_of_a_plain_decoder() -> None:
    model = _model()
    keys_before = set(model.state_dict())
    apply_lora(model, LoraConfig(r=4))
    assert "blocks.0.attn.qkv.base.weight" in model.state_dict()
    merge_lora(model)
    assert set(model.state_dict()) == keys_before

def test_adapter_round_trips_through_safetensors(tmp_path: Path) -> None:
    idx = _idx()
    cfg = LoraConfig(r=4, alpha=8, targets=("q", "v"))
    trained = _model()
    apply_lora(trained, cfg)
    torch.manual_seed(11)
    with torch.no_grad():
        for _, adapter in lora_modules(trained):
            for b in adapter.b:
                b.normal_(std=0.05)
        expected, _ = trained(idx)

    path = save_adapter(trained, tmp_path / "adapter.safetensors", cfg)
    assert path.with_name("adapter_config.json").is_file()

    restored = _model()
    loaded = load_adapter(restored, path)
    assert loaded == cfg
    with torch.no_grad():
        got, _ = restored(idx)
    torch.testing.assert_close(expected, got)

def test_the_adapter_is_a_rounding_error_next_to_the_published_model() -> None:
    """The number the whole method exists for, on the model that is actually published.

    The toy decoder of these tests is 39 k parameters and two thirds of that is the embedding
    table, so its ratio says nothing. `medium` is the size the claim is about: r=8 on query and
    value across 16 layers of `d_model` 768 is 393 216 numbers, 0.34 % of 115 M -- 1 572 864 bytes,
    the 1.6 MB file a reader downloads instead of a second copy of the weights.
    """
    cfg = LoraConfig(r=8, alpha=16, targets=("q", "v"))
    medium = preset("medium")
    adapter = medium.n_layer * len(cfg.targets) * 2 * cfg.r * medium.d_model
    published_params = 115_120_128  # `medium`, as every card and report of the project says
    assert adapter == 393_216
    assert adapter / published_params < 0.004
    assert adapter * 4 == 1_572_864  # bytes in float32: 1.6 MB, the size of the published file


def test_gradients_reach_the_factors_and_nothing_else() -> None:
    model = _model()
    apply_lora(model, LoraConfig(r=4, targets=("q", "v")))
    idx = _idx()
    _, loss = model(idx, idx)
    assert loss is not None
    loss.backward()
    assert all(p.grad is not None for p in adapter_params(model))
    assert model.tokens.weight.grad is None
    for _, adapter in lora_modules(model):
        assert adapter.base.weight.grad is None

def test_targets_are_validated() -> None:
    with pytest.raises(ValueError, match="unknown LoRA targets"):
        apply_lora(_model(), LoraConfig(targets=("qq",)))
    with pytest.raises(ValueError, match="at least one"):
        apply_lora(_model(), LoraConfig(targets=()))


def test_mlp_targets_use_the_feed_forward_width() -> None:
    """`mlp.fc` is `ff` wide, not `d_model`; a slice computed in the wrong unit would truncate."""
    model = _model()
    apply_lora(model, LoraConfig(r=2, targets=("mlp_in", "mlp_out")))
    adapters = dict(lora_modules(model))
    assert adapters["blocks.0.mlp.fc"].slices == ((0, model.cfg.ff),)
    assert adapters["blocks.0.mlp.proj"].slices == ((0, model.cfg.d_model),)


def test_scale_follows_alpha_over_r() -> None:
    base = torch.nn.Linear(4, 4)
    adapter = LoRALinear(base, r=2, alpha=8, dropout=0.0, slices=((0, 4),))
    assert adapter.scale == 4.0


def test_the_factors_are_built_where_the_weight_they_correct_lives() -> None:
    """`apply_lora` runs after the model has been moved to its device.

    The training loop loads the checkpoint first -- loading needs the module names of a plain
    decoder -- so by the time the adapters go on, the weights are already on the GPU. A factor
    created on the CPU by default meets a CUDA activation on the first forward and the run dies
    there. Every other test in this file runs on the CPU, where the bug cannot be seen; this one
    checks the property that makes it impossible instead of the device it happens to run on.
    """
    model = _model()
    base = model.blocks[0].attn.qkv.weight
    apply_lora(model, LoraConfig(r=4, targets=("q", "v")))
    for _, adapter in lora_modules(model):
        for factor in [*adapter.a, *adapter.b]:
            assert factor.device == base.device
            assert factor.dtype == base.dtype