tests/test_lora.py

"""Tests for LoRA infrastructure — injection, removal, parameter management."""

import pytest
import torch
from model.lora import LoRALayer, LoRAInjector, LoRAConfig, compute_total_lora_params
from model.attention import MultiHeadAttention


class TestLoRALayer:
    """Test the LoRA layer itself."""

    def test_init_shapes(self):
        lora = LoRALayer(in_dim=1024, out_dim=1024, rank=16)
        assert lora.A.shape == (16, 1024)
        assert lora.B.shape == (1024, 16)

    def test_forward_shape(self):
        lora = LoRALayer(in_dim=1024, out_dim=1024, rank=16)
        x = torch.randn(2, 10, 1024)
        out = lora(x)
        assert out.shape == (2, 10, 1024)

    def test_starts_as_zero(self):
        """B is initialized to zeros, so LoRA output should start as zero."""
        lora = LoRALayer(in_dim=1024, out_dim=1024, rank=16)
        x = torch.randn(2, 10, 1024)
        out = lora(x)
        assert torch.allclose(out, torch.zeros_like(out), atol=1e-6)

    def test_from_flat_params(self):
        """Create LoRA from flat parameter vector."""
        r, d = 16, 1024
        flat = torch.randn(r * d + d * r)
        lora = LoRALayer.from_flat_params(flat, in_dim=d, out_dim=d, rank=r)
        assert lora.A.shape == (r, d)
        assert lora.B.shape == (d, r)

    def test_from_flat_params_wrong_size(self):
        """Should raise on wrong parameter count."""
        flat = torch.randn(100)
        with pytest.raises(AssertionError):
            LoRALayer.from_flat_params(flat, in_dim=1024, out_dim=1024, rank=16)

    def test_num_params(self):
        lora = LoRALayer(in_dim=1024, out_dim=1024, rank=16)
        assert lora.num_params == 16 * 1024 + 1024 * 16  # A + B


class TestLoRAInjector:
    """Test the injector that creates LoRA layers for all decoder blocks."""

    def test_total_params_calculation(self):
        config = LoRAConfig(rank=16, targets=("q", "v"))
        injector = LoRAInjector(config, num_blocks=12, embed_dim=1024)
        expected = compute_total_lora_params(12, 1024, 16, ("q", "v"))
        assert injector.total_params == expected

    def test_create_random_layers(self):
        config = LoRAConfig(rank=16, targets=("q", "v"))
        injector = LoRAInjector(config, num_blocks=6, embed_dim=512)
        layers = injector.create_lora_layers()
        assert len(layers) == 6
        for block_layers in layers:
            assert "q" in block_layers
            assert "v" in block_layers

    def test_create_from_flat(self):
        config = LoRAConfig(rank=8, targets=("q", "v"))
        injector = LoRAInjector(config, num_blocks=3, embed_dim=256)
        flat = torch.randn(injector.total_params)
        layers = injector.create_lora_layers(flat)
        assert len(layers) == 3


class TestMultiHeadAttentionLoRA:
    """Test LoRA hooks in MultiHeadAttention."""

    def test_no_lora_by_default(self):
        attn = MultiHeadAttention(embed_dim=256, num_heads=4)
        assert not attn.has_lora

    def test_set_and_clear_lora(self):
        attn = MultiHeadAttention(embed_dim=256, num_heads=4)
        lora_q = LoRALayer(256, 256, rank=4)
        lora_v = LoRALayer(256, 256, rank=4)

        attn.set_lora(lora_q, lora_v)
        assert attn.has_lora

        attn.clear_lora()
        assert not attn.has_lora

    def test_output_changes_with_lora(self):
        """Output should differ when LoRA is active (unless B is zero)."""
        torch.manual_seed(42)
        attn = MultiHeadAttention(embed_dim=256, num_heads=4)
        attn.eval()  # Disable dropout for deterministic comparison
        x = torch.randn(2, 8, 256)

        with torch.no_grad():
            # Output without LoRA
            out_base = attn(x).clone()

            # Create LoRA with non-zero B
            lora_q = LoRALayer(256, 256, rank=4)
            lora_q.B.data = torch.randn_like(lora_q.B) * 0.1
            attn.set_lora(lora_q=lora_q)
            out_lora = attn(x)

            # Outputs should differ
            assert not torch.allclose(out_base, out_lora, atol=1e-5)

            # After clearing, output should match base
            attn.clear_lora()
            out_restored = attn(x)
            assert torch.allclose(out_base, out_restored, atol=1e-5)


class TestComputeTotalLoRAParams:
    def test_basic(self):
        total = compute_total_lora_params(12, 1024, 16, ("q", "v"))
        # Per layer: 16*1024 + 1024*16 = 32,768
        # 12 blocks * 2 targets = 24 layers
        # Total: 24 * 32,768 = 786,432
        assert total == 786432

    def test_single_target(self):
        total = compute_total_lora_params(6, 512, 8, ("q",))
        per_layer = 8 * 512 + 512 * 8
        assert total == per_layer * 6 * 1