LoRA: The Clever Trick That Makes Fine-Tuning Feasible

import torch
import torch.nn as nn
import math

class LoRALayer(nn.Module):
    """
    A Low-Rank Adaptation layer.
    
    This doesn't replace a linear layer. It sits *alongside* one,
    adding a small trainable update.
    """
    
    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 8,
        alpha: float = 16.0,
        dropout: float = 0.1
    ):
        super().__init__()
        self.rank = rank
        self.alpha = alpha
        
        # This scaling factor controls how much the LoRA update affects the output
        # alpha/rank is a common choice that keeps the update magnitude reasonable
        self.scaling = alpha / rank
        
        # The two low-rank matrices
        # A is initialized with random values (Kaiming initialization)
        # B is initialized to zeros (so initially, LoRA contributes nothing)
        self.lora_A = nn.Parameter(torch.zeros(rank, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
        
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)
        
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Compute the LoRA contribution: (B @ A) @ x * scaling
        
        x has shape: (batch_size, sequence_length, in_features)
        We're computing: x @ A^T @ B^T * scaling
        """
        lora_out = self.dropout(x) @ self.lora_A.T @ self.lora_B.T
        return lora_out * self.scaling


# Let's see the parameter savings in action
in_features, out_features = 4096, 4096
rank = 8

# Full fine-tuning would update this many parameters:
full_params = in_features * out_features

# LoRA only needs this many:
lora_params = (in_features * rank) + (out_features * rank)

print(f"Full fine-tuning: {full_params:,} parameters")
print(f"LoRA (rank={rank}): {lora_params:,} parameters")
print(f"Reduction: {full_params / lora_params:.1f}x fewer parameters")
print()
print(f"That's {100 * lora_params / full_params:.2f}% of the original size.")
print(f"(Or to put it another way: we're using {100 * (1 - lora_params / full_params):.1f}% less parameters!)")

class LoRALinear(nn.Module):
    """
    A linear layer with LoRA adaptation.
    
    Think of this as a wrapper around a pre-trained linear layer.
    The original layer is frozen, and we add a small LoRA "correction" on top.
    """
    
    def __init__(
        self,
        original_layer: nn.Linear,
        rank: int = 8,
        alpha: float = 16.0,
        dropout: float = 0.1
    ):
        super().__init__()
        
        # Store the original layer and freeze it
        # (We're not training these weights: they stay as pre-trained)
        self.original = original_layer
        for param in self.original.parameters():
            param.requires_grad = False
        
        # Add the LoRA adapter
        self.lora = LoRALayer(
            in_features=original_layer.in_features,
            out_features=original_layer.out_features,
            rank=rank,
            alpha=alpha,
            dropout=dropout
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass: original output + LoRA output.
        
        The original layer produces the "base" output.
        The LoRA layer produces a small "correction."
        We add them together.
        """
        return self.original(x) + self.lora(x)


# Let's test it with a realistic size (GPT-2's hidden dimension)
original = nn.Linear(768, 768)
lora_layer = LoRALinear(original, rank=8)

# Count trainable vs. frozen parameters
trainable = sum(p.numel() for p in lora_layer.parameters() if p.requires_grad)
total = sum(p.numel() for p in lora_layer.parameters())

print(f"Trainable parameters: {trainable:,}")
print(f"Total parameters: {total:,}")
print(f"Trainable: {100 * trainable / total:.2f}%")
print()
print("So we're only updating ~2% of the layer's parameters!")
print("(The other 98% stay frozen at their pre-trained values.)")

from transformers import AutoModelForCausalLM, AutoTokenizer

# Load GPT-2 (the small version, 124M parameters)
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

total_params = sum(p.numel() for p in model.parameters())
print(f"Base model: {model_name}")
print(f"Total parameters: {total_params:,}")
print()
print("Now let's add LoRA to this...")

from peft import LoraConfig, get_peft_model, TaskType

# Configure LoRA
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,  # We're doing causal language modeling
    r=8,                            # Rank: the size of the low-rank matrices
    lora_alpha=16,                  # Scaling factor (typically 2*r)
    lora_dropout=0.1,               # Dropout applied to LoRA layers
    target_modules=["c_attn", "c_proj"],  # Which layers to adapt
    fan_in_fan_out=True,            # Required for GPT-2's Conv1D layers
    bias="none"                     # Don't add LoRA to bias terms
)

# Apply LoRA to the model
# This wraps the specified layers with LoRA adapters
peft_model = get_peft_model(model, lora_config)

# Print the breakdown
peft_model.print_trainable_parameters()

print()
print("Wait, what are 'c_attn' and 'c_proj'?")
print()
print("Those are GPT-2's attention layers:")
print("  - c_attn: The combined Query/Key/Value projection")
print("  - c_proj: The output projection after attention")
print()
print("We're adding LoRA *only* to these layers, not the entire model.")
print("(Empirically, adapting attention layers gives the best results.)")

def merge_lora_weights(original_weight, lora_A, lora_B, scaling):
    """
    Merge LoRA weights into the original weight matrix.
    
    W_merged = W_original + (B @ A) * scaling
    
    After this, you can throw away the LoRA matrices and just use W_merged.
    """
    # Compute the low-rank update
    delta_W = (lora_B @ lora_A) * scaling
    
    # Add it to the original weights
    return original_weight + delta_W


# Example with realistic dimensions
d, k, r = 768, 768, 8
alpha = 16
scaling = alpha / r

# Simulate the weights
W_original = torch.randn(d, k)  # Original pre-trained weights
lora_A = torch.randn(r, k)      # LoRA matrix A (trained)
lora_B = torch.randn(d, r)      # LoRA matrix B (trained)

# Merge them
W_merged = merge_lora_weights(W_original, lora_A, lora_B, scaling)

print(f"Original W shape: {W_original.shape}")
print(f"LoRA A shape: {lora_A.shape}")
print(f"LoRA B shape: {lora_B.shape}")
print(f"Merged W shape: {W_merged.shape}")
print()
print("Before merging:")
print(f"  - Two separate matrix multiplies during inference")
print(f"  - Slightly slower, but LoRA weights can be swapped out")
print()
print("After merging:")
print(f"  - Single matrix multiply (same as original model)")
print(f"  - Zero inference overhead!")
print(f"  - Trade-off: can't easily swap LoRA adapters anymore")
print()
print("(In practice, you'd merge for production deployment.)")