Kohya-ss-sd-scripts/tests/library/test_network_utils_urae.py

import pytest
import torch

from library.network_utils import convert_urae_to_standard_lora


class TestConvertURAEToStandardLoRA:
    @pytest.fixture
    def sample_matrices(self):
        """Create sample matrices for testing"""
        # Original up matrix (4x2)
        orig_up = torch.tensor([
            [1.0, 2.0],
            [3.0, 4.0],
            [5.0, 6.0],
            [7.0, 8.0]
        ])
        
        # Original down matrix (2x6)
        orig_down = torch.tensor([
            [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
            [0.7, 0.8, 0.9, 1.0, 1.1, 1.2]
        ])
        
        # Trained up matrix (4x2) - same shape as orig_up but with changed values
        trained_up = torch.tensor([
            [1.1, 2.1],
            [3.1, 4.1],
            [5.1, 6.1],
            [7.1, 8.1]
        ])
        
        # Trained down matrix (2x6) - same shape as orig_down but with changed values
        trained_down = torch.tensor([
            [0.15, 0.25, 0.35, 0.45, 0.55, 0.65],
            [0.75, 0.85, 0.95, 1.05, 1.15, 1.25]
        ])
        
        return {
            'orig_up': orig_up,
            'orig_down': orig_down,
            'trained_up': trained_up,
            'trained_down': trained_down
        }

    def test_basic_conversion(self, sample_matrices):
        """Test the basic functionality of convert_urae_to_standard_lora"""
        lora_up, lora_down, alpha = convert_urae_to_standard_lora(
            sample_matrices['trained_up'],
            sample_matrices['trained_down'],
            sample_matrices['orig_up'],
            sample_matrices['orig_down']
        )
        
        # Check shapes
        assert lora_up.shape[0] == sample_matrices['trained_up'].shape[0]  # Same number of rows as trained_up
        assert lora_up.shape[1] == sample_matrices['trained_up'].shape[1]  # Same rank as trained_up
        assert lora_down.shape[0] == sample_matrices['trained_up'].shape[1]  # Same rank as trained_up
        assert lora_down.shape[1] == sample_matrices['trained_down'].shape[1]  # Same number of columns as trained_down
        
        # Check alpha is a reasonable value
        assert 0.1 <= alpha <= 1024.0
        
        # Check that lora_up @ lora_down approximates the weight delta
        delta = (sample_matrices['trained_up'] @ sample_matrices['trained_down']) - (sample_matrices['orig_up'] @ sample_matrices['orig_down'])
        
        # The approximation should be close in Frobenius norm after scaling
        lora_effect = lora_up @ lora_down
        delta_norm = torch.norm(delta, p="fro").item()
        lora_norm = torch.norm(lora_effect, p="fro").item()
        
        # Either they are close, or the alpha scaling brings them close
        scaled_lora_effect = (alpha / sample_matrices['trained_up'].shape[1]) * lora_effect
        scaled_lora_norm = torch.norm(scaled_lora_effect, p="fro").item()
        
        # At least one of these should be true
        assert abs(delta_norm - lora_norm) < 1e-4 or abs(delta_norm - scaled_lora_norm) < 1e-4

    def test_specified_rank(self, sample_matrices):
        """Test conversion with a specified rank"""
        new_rank = 1  # Lower than trained_up's rank of 2
        
        lora_up, lora_down, alpha = convert_urae_to_standard_lora(
            sample_matrices['trained_up'],
            sample_matrices['trained_down'],
            sample_matrices['orig_up'],
            sample_matrices['orig_down'],
            rank=new_rank
        )
        
        # Check that the new rank is used
        assert lora_up.shape[1] == new_rank
        assert lora_down.shape[0] == new_rank
        
        # Should still produce a reasonable alpha
        assert 0.1 <= alpha <= 1024.0

    def test_with_initial_alpha(self, sample_matrices):
        """Test conversion with a specified initial alpha"""
        initial_alpha = 16.0
        
        lora_up, lora_down, alpha = convert_urae_to_standard_lora(
            sample_matrices['trained_up'],
            sample_matrices['trained_down'],
            sample_matrices['orig_up'],
            sample_matrices['orig_down'],
            initial_alpha=initial_alpha
        )
        
        # Alpha should be influenced by initial_alpha but may be adjusted
        # Since we're using same rank, should be reasonably close to initial_alpha
        assert 0.1 <= alpha <= 1024.0
        # Should at least preserve the order of magnitude in typical cases
        assert abs(alpha - initial_alpha) <= initial_alpha * 4.0

    def test_large_initial_alpha(self, sample_matrices):
        """Test conversion with a very large initial alpha that should be capped"""
        initial_alpha = 2000.0  # Larger than the 1024.0 cap
        
        lora_up, lora_down, alpha = convert_urae_to_standard_lora(
            sample_matrices['trained_up'],
            sample_matrices['trained_down'],
            sample_matrices['orig_up'],
            sample_matrices['orig_down'],
            initial_alpha=initial_alpha
        )
        
        # Alpha should be capped at 1024.0
        assert alpha <= 1024.0

    def test_very_small_initial_alpha(self, sample_matrices):
        """Test conversion with a very small initial alpha that should be floored"""
        initial_alpha = 0.01  # Smaller than the 0.1 floor
        
        lora_up, lora_down, alpha = convert_urae_to_standard_lora(
            sample_matrices['trained_up'],
            sample_matrices['trained_down'],
            sample_matrices['orig_up'],
            sample_matrices['orig_down'],
            initial_alpha=initial_alpha
        )
        
        # Alpha should be floored at 0.1
        assert alpha >= 0.1

    def test_change_rank_with_initial_alpha(self, sample_matrices):
        """Test conversion with both rank change and initial alpha"""
        initial_alpha = 16.0
        new_rank = 1  # Half of original rank 2
        
        lora_up, lora_down, alpha = convert_urae_to_standard_lora(
            sample_matrices['trained_up'],
            sample_matrices['trained_down'],
            sample_matrices['orig_up'],
            sample_matrices['orig_down'],
            initial_alpha=initial_alpha,
            rank=new_rank
        )
        
        # Check shapes
        assert lora_up.shape[1] == new_rank
        assert lora_down.shape[0] == new_rank
        
        # Alpha should be adjusted for the rank change (approx halved in this case)
        expected_alpha = initial_alpha * (new_rank / sample_matrices['trained_up'].shape[1])
        # Allow some tolerance for adjustments from norm-based capping
        assert abs(alpha - expected_alpha) <= expected_alpha * 4.0 or alpha >= 0.1

    def test_zero_delta(self):
        """Test conversion when delta is zero"""
        # Create matrices where the delta will be zero
        dim_in, rank, dim_out = 4, 2, 6
        
        # Create identical matrices for original and trained
        orig_up = torch.randn(dim_in, rank)
        orig_down = torch.randn(rank, dim_out)
        trained_up = orig_up.clone()
        trained_down = orig_down.clone()
        
        lora_up, lora_down, alpha = convert_urae_to_standard_lora(
            trained_up,
            trained_down,
            orig_up,
            orig_down
        )
        
        # Should still return matrices of correct shape
        assert lora_up.shape == (dim_in, rank)
        assert lora_down.shape == (rank, dim_out)
        
        # Alpha should be at least the minimum
        assert alpha >= 0.1

    def test_large_dimensions(self):
        """Test with larger matrix dimensions"""
        dim_in, rank, dim_out = 100, 8, 200
        
        orig_up = torch.randn(dim_in, rank)
        orig_down = torch.randn(rank, dim_out)
        trained_up = orig_up + 0.01 * torch.randn(dim_in, rank)  # Small perturbation
        trained_down = orig_down + 0.01 * torch.randn(rank, dim_out)  # Small perturbation
        
        lora_up, lora_down, alpha = convert_urae_to_standard_lora(
            trained_up,
            trained_down,
            orig_up,
            orig_down
        )
        
        # Check shapes
        assert lora_up.shape == (dim_in, rank)
        assert lora_down.shape == (rank, dim_out)
        
        # Should produce a reasonable alpha
        assert 0.1 <= alpha <= 1024.0

    def test_rank_exceeding_singular_values(self):
        """Test when requested rank exceeds available singular values"""
        # Small matrices with limited rank
        dim_in, rank, dim_out = 3, 2, 3
        
        orig_up = torch.randn(dim_in, rank)
        orig_down = torch.randn(rank, dim_out)
        trained_up = orig_up + 0.1 * torch.randn(dim_in, rank)
        trained_down = orig_down + 0.1 * torch.randn(rank, dim_out)
        
        # Request rank larger than possible
        too_large_rank = 10
        
        lora_up, lora_down, alpha = convert_urae_to_standard_lora(
            trained_up,
            trained_down,
            orig_up,
            orig_down,
            rank=too_large_rank
        )
        
        # Rank should be limited to min(dim_in, dim_out, S.size)
        max_possible_rank = min(dim_in, dim_out)
        assert lora_up.shape[1] <= max_possible_rank
        assert lora_down.shape[0] <= max_possible_rank