Efficient Fine-Tuning with LoRA

Introduction to LoRA

Low-Rank Adaptation (LoRA) is a technique for fine-tuning large language models more efficiently. It allows for task-specific adaptations without modifying most of the original model parameters. This approach significantly reduces computational requirements, memory usage, and training time compared to traditional fine-tuning methods.

import torch
import torch.nn as nn

class LoRALayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4):
        super().__init__()
        self.A = nn.Parameter(torch.randn(in_features, rank))
        self.B = nn.Parameter(torch.zeros(rank, out_features))
        
    def forward(self, x):
        return x @ (self.A @ self.B)

# Example usage
lora_layer = LoRALayer(768, 768)  # Assuming 768-dimensional embeddings
input_tensor = torch.randn(1, 768)
output = lora_layer(input_tensor)

Efficiency of LoRA

LoRA achieves efficiency by freezing the original model weights and introducing small, low-rank matrices for task-specific adjustments. This approach eliminates the need to retrain the entire model from scratch, resulting in significant computational savings. The low-rank matrices act as adapters, allowing the model to learn task-specific information while preserving the general knowledge encoded in the pre-trained weights.

class LoRAModel(nn.Module):
    def __init__(self, base_model, rank=4):
        super().__init__()
        self.base_model = base_model
        self.lora_layers = nn.ModuleDict()
        
        for name, module in self.base_model.named_modules():
            if isinstance(module, nn.Linear):
                self.lora_layers[name] = LoRALayer(module.in_features, module.out_features, rank)
    
    def forward(self, x):
        for name, module in self.base_model.named_modules():
            if isinstance(module, nn.Linear):
                x = module(x) + self.lora_layers[name](x)
            else:
                x = module(x)
        return x

# Example usage
base_model = nn.Sequential(nn.Linear(768, 512), nn.ReLU(), nn.Linear(512, 256))
lora_model = LoRAModel(base_model)

Memory-Friendly Approach

LoRA significantly reduces GPU memory usage, making it suitable for resource-constrained environments. By focusing on low-rank adaptations, LoRA can achieve up to a 3x reduction in memory requirements compared to full fine-tuning. This memory efficiency allows researchers and developers to work with larger models on more modest hardware setups.

def calculate_memory_savings(base_model, lora_model, input_size):
    def count_parameters(model):
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    base_params = count_parameters(base_model)
    lora_params = count_parameters(lora_model)
    
    memory_savings = 1 - (lora_params / base_params)
    return memory_savings

# Example usage
base_model = nn.Sequential(nn.Linear(768, 512), nn.ReLU(), nn.Linear(512, 256))
lora_model = LoRAModel(base_model)

savings = calculate_memory_savings(base_model, lora_model, 768)
print(f"Memory savings: {savings:.2%}")

LoRA+: Enhanced Learning Rates

LoRA+ improves upon the original LoRA by introducing different learning rates for adapters. This enhancement can speed up fine-tuning by up to 2x, making it particularly effective for complex tasks where standard LoRA may fall short. The adaptive learning rates allow for more nuanced parameter updates, leading to faster convergence and improved performance.

class LoRAPlusLayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4):
        super().__init__()
        self.A = nn.Parameter(torch.randn(in_features, rank))
        self.B = nn.Parameter(torch.zeros(rank, out_features))
        self.lr_A = nn.Parameter(torch.ones(1))
        self.lr_B = nn.Parameter(torch.ones(1))
        
    def forward(self, x):
        return x @ (self.lr_A * self.A @ (self.lr_B * self.B))

# Example usage
lora_plus_layer = LoRAPlusLayer(768, 768)
input_tensor = torch.randn(1, 768)
output = lora_plus_layer(input_tensor)

QLoRA: Quantization for Extreme Compression

QLoRA combines LoRA with quantization techniques to drastically reduce model size. This approach enables fine-tuning of extremely large models (up to 65 billion parameters) on consumer-grade GPUs with as little as 48GB of memory. QLoRA achieves this by quantizing the base model weights while keeping the LoRA adapters in full precision.

def quantize(tensor, bits=4):
    scale = tensor.abs().max() / (2**(bits-1) - 1)
    quantized = torch.round(tensor / scale).to(torch.int8)
    return quantized, scale

class QLoRALayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4, bits=4):
        super().__init__()
        self.quantized_weight, self.scale = quantize(torch.randn(in_features, out_features), bits)
        self.A = nn.Parameter(torch.randn(in_features, rank))
        self.B = nn.Parameter(torch.zeros(rank, out_features))
        
    def forward(self, x):
        base_output = x @ (self.quantized_weight.float() * self.scale)
        lora_output = x @ (self.A @ self.B)
        return base_output + lora_output

# Example usage
qlora_layer = QLoRALayer(768, 768)
input_tensor = torch.randn(1, 768)
output = qlora_layer(input_tensor)

DyLoRA: Dynamic Rank Adaptation

DyLoRA introduces dynamic adjustment of low-rank adapters during training. This adaptive approach eliminates the need for manual trial and error in finding the optimal rank, saving considerable time and effort. DyLoRA automatically tunes the rank of the adapters based on the task complexity and model performance.

class DyLoRALayer(nn.Module):
    def __init__(self, in_features, out_features, max_rank=16):
        super().__init__()
        self.A = nn.Parameter(torch.randn(in_features, max_rank))
        self.B = nn.Parameter(torch.zeros(max_rank, out_features))
        self.rank_importance = nn.Parameter(torch.ones(max_rank))
        
    def forward(self, x):
        effective_rank = torch.softmax(self.rank_importance, dim=0)
        A_weighted = self.A * effective_rank.unsqueeze(0)
        B_weighted = self.B * effective_rank.unsqueeze(1)
        return x @ (A_weighted @ B_weighted)

# Example usage
dylora_layer = DyLoRALayer(768, 768)
input_tensor = torch.randn(1, 768)
output = dylora_layer(input_tensor)

LoRA-FA: Further Memory Reduction

LoRA-FA (Frozen Adaptation) extends LoRA by freezing some weight matrices to reduce memory usage even further while maintaining performance. This technique is particularly useful when working with extremely large models or in severely resource-constrained environments.

class LoRAFALayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4, freeze_fraction=0.5):
        super().__init__()
        self.frozen_weight = nn.Parameter(torch.randn(in_features, out_features), requires_grad=False)
        unfrozen_size = int(out_features * (1 - freeze_fraction))
        self.A = nn.Parameter(torch.randn(in_features, rank))
        self.B = nn.Parameter(torch.zeros(rank, unfrozen_size))
        
    def forward(self, x):
        frozen_output = x @ self.frozen_weight
        lora_output = x @ (self.A @ self.B)
        return torch.cat([frozen_output[:, :self.B.size(1)], lora_output], dim=1)

# Example usage
lora_fa_layer = LoRAFALayer(768, 768)
input_tensor = torch.randn(1, 768)
output = lora_fa_layer(input_tensor)

DoRA: Direction and Magnitude Decomposition

DoRA (Decomposed Rank Adaptation) improves upon LoRA by decomposing weights into direction and magnitude components. This approach enhances performance while keeping changes to the original model minimal. DoRA allows for more nuanced adaptations, potentially leading to better task-specific fine-tuning.

class DoRALayer(nn.Module):
    def __init__(self, in_features, out_features, rank=4):
        super().__init__()
        self.direction = nn.Parameter(torch.randn(in_features, out_features))
        self.direction.data = torch.nn.functional.normalize(self.direction.data, dim=0)
        self.magnitude = nn.Parameter(torch.ones(out_features))
        self.A = nn.Parameter(torch.randn(in_features, rank))
        self.B = nn.Parameter(torch.zeros(rank, out_features))
        
    def forward(self, x):
        base_output = x @ (self.direction * self.magnitude)
        lora_output = x @ (self.A @ self.B)
        return base_output + lora_output

# Example usage
dora_layer = DoRALayer(768, 768)
input_tensor = torch.randn(1, 768)
output = dora_layer(input_tensor)

Real-Life Example: Sentiment Analysis

Let's explore how LoRA can be applied to fine-tune a pre-trained language model for sentiment analysis. We'll use a simplified example to demonstrate the process.

import torch
import torch.nn as nn

class SentimentClassifier(nn.Module):
    def __init__(self, base_model, num_classes=2):
        super().__init__()
        self.base_model = base_model
        self.classifier = nn.Linear(768, num_classes)
        
    def forward(self, x):
        features = self.base_model(x)
        return self.classifier(features)

# Simulated pre-trained model
base_model = nn.Sequential(nn.Embedding(10000, 768), nn.LSTM(768, 768, batch_first=True))
sentiment_model = SentimentClassifier(base_model)

# Apply LoRA
lora_sentiment_model = LoRAModel(sentiment_model)

# Example usage
input_ids = torch.randint(0, 10000, (1, 50))  # Batch of 1 sentence with 50 tokens
output = lora_sentiment_model(input_ids)
print(f"Sentiment logits: {output}")

Real-Life Example: Text Generation

Another practical application of LoRA is fine-tuning a language model for specific text generation tasks. Here's an example of how LoRA can be applied to a simple language model for generating text in a particular style or domain.

class LanguageModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)
        
    def forward(self, x):
        embedded = self.embedding(x)
        lstm_out, _ = self.lstm(embedded)
        return self.fc(lstm_out)

# Simulated pre-trained model
base_lm = LanguageModel(vocab_size=10000, embedding_dim=256, hidden_dim=512)

# Apply LoRA
lora_lm = LoRAModel(base_lm)

# Example usage
input_sequence = torch.randint(0, 10000, (1, 20))  # Batch of 1 sequence with 20 tokens
output_logits = lora_lm(input_sequence)
generated_tokens = torch.argmax(output_logits, dim=-1)
print(f"Generated token IDs: {generated_tokens}")

Implementing LoRA Training Loop

To effectively use LoRA, it's crucial to understand how to set up a training loop. This example demonstrates a basic training loop for a LoRA-adapted model, highlighting the key differences from standard fine-tuning.

def train_lora_model(model, train_dataloader, num_epochs, learning_rate):
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    criterion = nn.CrossEntropyLoss()
    
    for epoch in range(num_epochs):
        model.train()
        for batch in train_dataloader:
            inputs, labels = batch
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            
            # Update only LoRA parameters
            for name, param in model.named_parameters():
                if 'lora_layers' in name:
                    param.grad.data *= 1.0  # You can adjust scaling factor
                else:
                    param.grad.data *= 0.0  # Zero out gradients for non-LoRA params
            
            optimizer.step()
        
        print(f"Epoch {epoch+1}/{num_epochs} completed")

# Example usage (assuming you have a DataLoader set up)
# train_dataloader = ...
# lora_model = LoRAModel(base_model)
# train_lora_model(lora_model, train_dataloader, num_epochs=5, learning_rate=1e-4)

Evaluating LoRA Performance

After training a LoRA-adapted model, it's important to evaluate its performance and compare it to the base model. This slide demonstrates how to set up an evaluation loop and calculate relevant metrics.

def evaluate_model(model, eval_dataloader):
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for batch in eval_dataloader:
            inputs, labels = batch
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    accuracy = correct / total
    return accuracy

# Example usage
# eval_dataloader = ...
# base_model = ...
# lora_model = LoRAModel(base_model)

# base_accuracy = evaluate_model(base_model, eval_dataloader)
# lora_accuracy = evaluate_model(lora_model, eval_dataloader)

# print(f"Base model accuracy: {base_accuracy:.4f}")
# print(f"LoRA model accuracy: {lora_accuracy:.4f}")

Saving and Loading LoRA Models

One of the advantages of LoRA is the ability to save and distribute only the adapted parameters, which are much smaller than the full model. This approach allows for efficient storage and sharing of task-specific adaptations without the need to distribute the entire base model.

def save_lora_adaptations(model, path):
    lora_state_dict = {name: param for name, param in model.state_dict().items() if 'lora_layers' in name}
    torch.save(lora_state_dict, path)

def load_lora_adaptations(base_model, lora_path):
    lora_model = LoRAModel(base_model)
    lora_state_dict = torch.load(lora_path)
    lora_model.load_state_dict(lora_state_dict, strict=False)
    return lora_model

# Example usage
base_model = create_base_model()
lora_model = train_lora_model(base_model, train_data)

# Save LoRA adaptations
save_lora_adaptations(lora_model, "lora_adaptations.pth")

# Load LoRA adaptations
new_lora_model = load_lora_adaptations(base_model, "lora_adaptations.pth")

Comparing LoRA Variants

Different LoRA variants offer unique advantages for specific scenarios. This slide provides a simple comparison function to evaluate the performance of various LoRA implementations on a given task.

def compare_lora_variants(base_model, variants, dataset):
    results = {}
    for name, variant_class in variants.items():
        model = variant_class(base_model)
        accuracy = train_and_evaluate(model, dataset)
        results[name] = accuracy
    return results

# Example usage
variants = {
    "LoRA": LoRAModel,
    "QLoRA": QLoRAModel,
    "DyLoRA": DyLoRAModel,
    "DoRA": DoRAModel
}

dataset = load_dataset()
comparison_results = compare_lora_variants(base_model, variants, dataset)

for variant, accuracy in comparison_results.items():
    print(f"{variant} accuracy: {accuracy:.4f}")

Additional Resources

For those interested in diving deeper into LoRA and its variants, here are some valuable resources:

1. Original LoRA paper: "LoRA: Low-Rank Adaptation of Large Language Models" (https://arxiv.org/abs/2106.09685)
2. QLoRA: "QLoRA: Efficient Finetuning of Quantized LLMs" (https://arxiv.org/abs/2305.14314)
3. DoRA: "DoRA: Weight-Decomposed Low-Rank Adaptation" (https://arxiv.org/abs/2402.09353)
4. DyLoRA: "DyLoRA: Parameter-Efficient Tuning of Pretrained Models using Dynamic Search Strategy" (https://arxiv.org/abs/2210.07558)

These papers provide in-depth explanations of the techniques, their implementations, and experimental results. They serve as excellent starting points for understanding the theoretical foundations and practical applications of LoRA and its variants.