test_optimized_lokr.py

#!/usr/bin/env python3
"""
Performance test script for optimized LoKr implementation

This script compares the performance of the original LoKr implementation
with the optimized version on different matrix sizes and devices.
"""

import os
import sys
import time
import torch
import torch.nn as nn
from pathlib import Path
import argparse

# Add local directory to path so we can import modules
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from modules.module.LoRAModule import LoKrModule
from modules.util.enum.ModelType import PeftType

# Set environment variable to enable weight merging
os.environ["ONETRAINER_FORCE_LOKR_WEIGHT_MERGING"] = "1"

def create_test_module(matrix_size, rank, factor=4, device='cuda'):
    """Create a test LoKr module with specified dimensions."""
    linear = nn.Linear(matrix_size, matrix_size).to(device)
    prefix = f"test_lokr_{matrix_size}_{rank}"
    
    # Create a LoKr module
    lokr_module = LoKrModule(prefix, linear, rank, alpha=16.0, factor=factor)
    
    # Initialize weights
    lokr_module.lokr_w1_a.data.normal_(0, 0.1)
    lokr_module.lokr_w1_b.data.normal_(0, 0.1)
    lokr_module.lokr_w2_a.data.normal_(0, 0.1)
    lokr_module.lokr_w2_b.data.normal_(0, 0.1)
    
    return lokr_module

def benchmark_kronecker_product(module, iterations=100, warmup=10):
    """Benchmark the Kronecker product computation."""
    # Get matrices
    W1 = module.lokr_w1_b @ module.lokr_w1_a
    W2 = module.lokr_w2_b @ module.lokr_w2_a
    
    # Warm up
    for _ in range(warmup):
        _ = module.kronecker_product(W1, W2)
    
    # Benchmark
    start_time = time.time()
    for _ in range(iterations):
        result = module.kronecker_product(W1, W2)
    end_time = time.time()
    
    # Return average time per operation in milliseconds
    return (end_time - start_time) * 1000 / iterations, result

def benchmark_make_weight_kronecker(module, iterations=100, warmup=10):
    """Benchmark the make_weight_kronecker method."""
    # Warm up
    for _ in range(warmup):
        _ = module.make_weight_kronecker()
    
    # Clear cache to ensure fair test
    if hasattr(module, '_cached_weight'):
        del module._cached_weight
        module._cached_computation_key = None
    if hasattr(module, '_cached_W1'):
        del module._cached_W1
        module._cached_W1_key = None
    if hasattr(module, '_cached_W2'):
        del module._cached_W2
        module._cached_W2_key = None
    
    # Benchmark
    start_time = time.time()
    for _ in range(iterations):
        result = module.make_weight_kronecker()
    end_time = time.time()
    
    # Return average time per operation in milliseconds
    return (end_time - start_time) * 1000 / iterations, result

def benchmark_forward(module, batch_size=4, iterations=100, warmup=10):
    """Benchmark the forward pass."""
    # Get device from module
    device = module.lokr_w1_a.device
    
    # Create random input
    input_size = module.orig_module.in_features
    x = torch.randn(batch_size, input_size).to(device)
    
    # Hook to module
    module.hook_to_module()
    
    # Warm up
    for _ in range(warmup):
        _ = module.orig_module(x)
    
    # Benchmark
    start_time = time.time()
    for _ in range(iterations):
        result = module.orig_module(x)
    end_time = time.time()
    
    # Remove hook
    module.remove_hook_from_module()
    
    # Return average time per operation in milliseconds
    return (end_time - start_time) * 1000 / iterations, result

def benchmark_weight_merging(module, batch_size=4, iterations=100, warmup=10):
    """Benchmark performance with weight merging."""
    # Get device from module
    device = module.lokr_w1_a.device
    
    # Create random input
    input_size = module.orig_module.in_features
    x = torch.randn(batch_size, input_size).to(device)
    
    # Hook to module
    module.hook_to_module()
    
    # Apply weight merging
    module.apply_to_module()
    
    # Warm up
    for _ in range(warmup):
        _ = module.orig_module(x)
    
    # Benchmark
    start_time = time.time()
    for _ in range(iterations):
        result = module.orig_module(x)
    end_time = time.time()
    
    # Restore original weights and remove hook
    module.restore_original_weights()
    module.remove_hook_from_module()
    
    # Return average time per operation in milliseconds
    return (end_time - start_time) * 1000 / iterations, result

def run_benchmarks(device='cuda', iterations=100):
    """Run a complete set of benchmarks for different matrix sizes."""
    print(f"\n{'=' * 60}")
    print(f"Running LoKr optimization benchmarks on {device}")
    print(f"{'=' * 60}")
    
    # Matrix sizes to test
    sizes = [(512, 16), (1024, 32), (2048, 64)]
    
    results = []
    
    for size, rank in sizes:
        print(f"\nBenchmarking with matrix size {size}x{size}, rank {rank}")
        
        # Create a module for testing
        module = create_test_module(size, rank, device=device)
        
        # Test in eval mode (caching enabled)
        module.eval()
        
        # Benchmark Kronecker product
        kron_time, _ = benchmark_kronecker_product(module, iterations=iterations)
        print(f"Kronecker product: {kron_time:.3f} ms")
        
        # Benchmark make_weight_kronecker
        make_weight_time, _ = benchmark_make_weight_kronecker(module, iterations=iterations)
        print(f"make_weight_kronecker: {make_weight_time:.3f} ms")
        
        # Benchmark forward pass
        forward_time, _ = benchmark_forward(module, iterations=iterations)
        print(f"Forward pass: {forward_time:.3f} ms")
        
        # Benchmark with weight merging
        merged_time, _ = benchmark_weight_merging(module, iterations=iterations)
        print(f"Forward with weight merging: {merged_time:.3f} ms")
        
        # Speedup from weight merging
        speedup = forward_time / merged_time if merged_time > 0 else float('inf')
        print(f"Weight merging speedup: {speedup:.2f}x")
        
        results.append({
            'size': size,
            'rank': rank,
            'kronecker_time': kron_time,
            'make_weight_time': make_weight_time,
            'forward_time': forward_time,
            'merged_time': merged_time,
            'speedup': speedup
        })
    
    # Print summary
    print(f"\n{'=' * 60}")
    print(f"Summary of optimizations")
    print(f"{'=' * 60}")
    for result in results:
        print(f"Size {result['size']}x{result['size']}, Rank {result['rank']}:")
        print(f"  - Forward pass: {result['forward_time']:.3f} ms")
        print(f"  - With weight merging: {result['merged_time']:.3f} ms")
        print(f"  - Speedup: {result['speedup']:.2f}x")
    
    # Calculate average speedup
    avg_speedup = sum(r['speedup'] for r in results) / len(results)
    print(f"\nAverage speedup: {avg_speedup:.2f}x")
    return results

def main():
    parser = argparse.ArgumentParser(description="Benchmark LoKr optimization")
    parser.add_argument('--device', type=str, default='cuda', help='Device to run on (cuda or cpu)')
    parser.add_argument('--iterations', type=int, default=100, help='Number of iterations for benchmarking')
    args = parser.parse_args()
    
    if args.device == 'cuda' and not torch.cuda.is_available():
        print("CUDA is not available, falling back to CPU")
        args.device = 'cpu'
    
    run_benchmarks(device=args.device, iterations=args.iterations)

if __name__ == "__main__":
    main()