Low-Rank Adaptation (LoRA) Implementation in PyTorch

This implementation builds Low-Rank Adaptation (LoRA) from scratch. It wraps a frozen nn.Linear layer with low-rank A/B matrices whose update is scaled by alpha / rank, dramatically reducing the number of trainable parameters for fine-tuning large models.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.linalg  
import math

class LoRALayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4, alpha=1):
        super().__init__()
        
        self.alpha = alpha
        self.rank = rank
        
        # Initialize LoRA matrices A and B
        self.lora_A = nn.Parameter(torch.zeros(in_features, rank))
        self.lora_B = nn.Parameter(torch.zeros(rank, out_features))
        
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)  

    def forward(self, x):
        # Compute: (x @ A) @ B * (alpha/rank)
        return (x @ self.lora_A @ self.lora_B) * (self.alpha / self.rank)

class LinearWithLoRA(nn.Module):
    def __init__(self, linear_layer, rank=4, alpha=1, freeze_weights=True):
        super().__init__()
        
        self.linear = linear_layer
        if freeze_weights:
            for param in self.linear.parameters():
                param.requires_grad = False
                
        # Add LoRA layer
        self.lora = LoRALayer(
            linear_layer.in_features,
            linear_layer.out_features,
            rank=rank,
            alpha=alpha
        )

    def forward(self, x):
        # linear + LoRA adjustment
        return self.linear(x) + self.lora(x)
    
# Test cases
def test_lora():
    in_features, out_features = 10, 5
    batch_size = 3
    rank = 2
    
    original_layer = nn.Linear(in_features, out_features)
    
    x = torch.randn(batch_size, in_features)
    
    # Wrap with LoRA
    lora_layer = LinearWithLoRA(original_layer, rank=rank)

    # Test forward pass
    output = lora_layer(x)
    assert output.shape == (batch_size, out_features), f"Shape mismatch: {output.shape}"
    print("(Shape test) passed!")

    # Freezing test
    original_params = sum(p.requires_grad for p in original_layer.parameters())
    assert original_params == 0, "Original weights not frozen"
    lora_params = sum(p.requires_grad for p in lora_layer.lora.parameters())
    assert lora_params == 2, "LoRA parameters not trainable"
    print("(Parameter freezing test) passed!")

    # Low-rank property test
    lora_contribution = lora_layer.lora(x)
    if len(lora_contribution.shape) > 2:
        lora_contribution = lora_contribution.reshape(-1, lora_contribution.shape[-1])
    rank_effective = torch.linalg.matrix_rank(lora_contribution)
    assert rank_effective <= rank, f"Rank constraint violated: {rank_effective} > {rank}"
    print("(Rank test) passed!")

    return "All tests passed successfully!"

if __name__ == "__main__":
    print(test_lora())

    input_size = 128
    output_size = 64
    batch_size = 32
    
    x = torch.randn(batch_size, input_size)
    
    original_layer = nn.Linear(input_size, output_size)
    
    # Add LoRA
    lora_model = LinearWithLoRA(original_layer, rank=8, alpha=16)
    
    # Forward pass
    output = lora_model(x)
    print(f"\nExample output shape: {output.shape}")