LLMs and generative AI / Fine tuning and alignment
Techniques for Fine-Tuning Large Language Models
Contents
- Implementing LoRA from Scratch
- Matrix Decomposition in LoRA
- Delta LoRA Implementation
- Prefix Tuning Architecture
- LoRA Feature Augmentation
- VERA Framework Implementation
- Performance Metrics Implementation
- Real-world Application - Text Classification
- Real-world Application Results
- Advanced LoRA Training Strategies
- Optimization Techniques in LoRA
- Memory-Efficient LoRA Implementation
- Additional Resources
Implementing LoRA from Scratch
Low-Rank Adaptation (LoRA) reduces the number of trainable parameters by decomposing weight matrices into low-rank approximations. This implementation demonstrates the core mathematics behind LoRA using NumPy, focusing on the fundamental matrix operations and rank decomposition.
import numpy as np
class LoRA:
def __init__(self, input_dim, output_dim, rank=4):
# Initialize matrices A and B for low-rank decomposition
self.A = np.random.randn(input_dim, rank) / np.sqrt(rank)
self.B = np.random.randn(rank, output_dim) / np.sqrt(rank)
self.alpha = 8 # Scaling factor
def forward(self, x):
# Compute LoRA transformation: x → x + α(BA)x
delta = self.alpha * (x @ self.A @ self.B)
return x + delta
def get_effective_weights(self):
return self.alpha * (self.A @ self.B)
# Example usage
lora = LoRA(input_dim=768, output_dim=768, rank=8)
x = np.random.randn(1, 768)
output = lora.forward(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Effective weights shape: {lora.get_effective_weights().shape}")
Matrix Decomposition in LoRA
The effectiveness of LoRA lies in its ability to approximate large weight matrices through low-rank decomposition. This implementation showcases the mathematical foundation using SVD decomposition and demonstrates parameter reduction calculations.
import numpy as np
from scipy.linalg import svd
def compute_rank_reduction(original_dim, rank):
"""Calculate parameter reduction from LoRA"""
original_params = original_dim * original_dim
lora_params = original_dim * rank * 2
reduction = (1 - lora_params/original_params) * 100
return original_params, lora_params, reduction
def low_rank_approximation(W, rank):
"""Compute low-rank approximation using SVD"""
U, S, Vt = svd(W, full_matrices=False)
# Keep only top-k singular values
U_k = U[:, :rank]
S_k = np.diag(S[:rank])
Vt_k = Vt[:rank, :]
# Reconstruct approximated matrix
W_approx = U_k @ S_k @ Vt_k
error = np.linalg.norm(W - W_approx, 'fro')
return W_approx, error
# Example usage
dim = 768
rank = 8
W = np.random.randn(dim, dim)
W_approx, error = low_rank_approximation(W, rank)
orig, lora, reduction = compute_rank_reduction(dim, rank)
print(f"Original parameters: {orig:,}")
print(f"LoRA parameters: {lora:,}")
print(f"Parameter reduction: {reduction:.2f}%")
print(f"Approximation error: {error:.4f}")
Delta LoRA Implementation
Delta LoRA enhances the standard LoRA by incorporating dropout mechanisms and adaptive scaling. This implementation demonstrates the key modifications that improve model robustness and generalization capabilities.
import numpy as np
class DeltaLoRA:
def __init__(self, input_dim, output_dim, rank=4, dropout_rate=0.1):
self.A = np.random.randn(input_dim, rank) / np.sqrt(rank)
self.B = np.random.randn(rank, output_dim) / np.sqrt(rank)
self.dropout_rate = dropout_rate
self.alpha = 8
self.scaling_vector = np.ones(rank)
def dropout_mask(self, shape):
return np.random.binomial(1, 1-self.dropout_rate, shape)
def forward(self, x, training=True):
if training:
mask = self.dropout_mask(self.A.shape[1])
A_dropped = self.A * mask[:, np.newaxis]
delta = self.alpha * (x @ (A_dropped * self.scaling_vector) @ self.B)
else:
delta = self.alpha * (x @ (self.A * self.scaling_vector) @ self.B)
return x + delta
# Example usage
delta_lora = DeltaLoRA(input_dim=768, output_dim=768, rank=8)
x = np.random.randn(1, 768)
train_output = delta_lora.forward(x, training=True)
eval_output = delta_lora.forward(x, training=False)
print(f"Training output shape: {train_output.shape}")
print(f"Evaluation output shape: {eval_output.shape}")
Prefix Tuning Architecture
Prefix tuning optimizes a small set of continuous task-specific vectors while keeping the main model parameters frozen. This implementation showcases the prefix embedding generation and integration with the main model's hidden states.
import numpy as np
import torch
import torch.nn as nn
class PrefixTuningModel(nn.Module):
def __init__(self, hidden_size, prefix_length=10, prefix_projection_size=512):
super().__init__()
self.prefix_length = prefix_length
# Prefix embedding layer
self.prefix_embedding = nn.Sequential(
nn.Linear(hidden_size, prefix_projection_size),
nn.Tanh(),
nn.Linear(prefix_projection_size, hidden_size * 2) # 2x for key and value
)
# Initialize prefix tokens
self.prefix_tokens = nn.Parameter(
torch.randn(prefix_length, hidden_size)
)
def forward(self, hidden_states):
batch_size = hidden_states.shape[0]
# Generate prefix
prefix = self.prefix_embedding(self.prefix_tokens)
prefix = prefix.view(self.prefix_length, 2, -1, hidden_size)
prefix_keys, prefix_values = prefix.split(1, dim=1)
# Reshape for batch processing
prefix_keys = prefix_keys.squeeze(1).expand(batch_size, -1, -1)
prefix_values = prefix_values.squeeze(1).expand(batch_size, -1, -1)
# Concatenate with input hidden states
output = torch.cat([prefix_keys, hidden_states], dim=1)
return output
# Example usage
hidden_size = 768
model = PrefixTuningModel(hidden_size)
sample_hidden = torch.randn(2, 20, hidden_size) # Batch x Seq x Hidden
output = model(sample_hidden)
print(f"Input shape: {sample_hidden.shape}")
print(f"Output shape: {output.shape}")
LoRA Feature Augmentation
LoRA-FA extends traditional LoRA by incorporating feature-based decomposition and directional scaling. This implementation demonstrates the enhanced adaptation mechanism with feature-aware weight updates.
import numpy as np
import torch
import torch.nn as nn
class LoRAFeatureAugmentation(nn.Module):
def __init__(self, input_dim, output_dim, rank=8, num_features=4):
super().__init__()
self.rank = rank
self.num_features = num_features
# Feature-aware decomposition matrices
self.feature_projector = nn.Linear(input_dim, num_features)
self.A = nn.Parameter(torch.randn(num_features, input_dim, rank) / np.sqrt(rank))
self.B = nn.Parameter(torch.randn(num_features, rank, output_dim) / np.sqrt(rank))
# Directional scaling factors
self.feature_scales = nn.Parameter(torch.ones(num_features))
self.alpha = 8
def forward(self, x):
# Extract feature weights
feature_weights = torch.softmax(self.feature_projector(x), dim=-1)
# Compute feature-weighted adaptation
delta = torch.zeros_like(x)
for i in range(self.num_features):
weight = feature_weights[:, i].unsqueeze(-1)
feature_delta = x @ self.A[i] @ self.B[i]
delta += weight * self.feature_scales[i] * feature_delta
return x + self.alpha * delta
# Example usage
model = LoRAFeatureAugmentation(768, 768)
x = torch.randn(2, 768)
output = model(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Feature scales: {model.feature_scales.data}")
VERA Framework Implementation
VERA (Verification-Enhanced Relevance Adaptation) focuses on maintaining output coherence while fine-tuning. This implementation showcases the verification mechanism and relevance scoring components that ensure high-quality adaptations.
import torch
import torch.nn as nn
import torch.nn.functional as F
class VERAModule(nn.Module):
def __init__(self, hidden_size, num_heads=8, dropout_rate=0.1):
super().__init__()
self.hidden_size = hidden_size
# Relevance scoring components
self.relevance_scorer = nn.MultiheadAttention(
hidden_size, num_heads, dropout=dropout_rate
)
# Verification gate
self.verification_gate = nn.Sequential(
nn.Linear(hidden_size * 2, hidden_size),
nn.LayerNorm(hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, 1),
nn.Sigmoid()
)
def forward(self, x, context):
# Compute relevance scores
relevance_output, attention_weights = self.relevance_scorer(
x, context, context
)
# Verification gate computation
combined = torch.cat([x, relevance_output], dim=-1)
gate_values = self.verification_gate(combined)
# Apply gated adaptation
output = gate_values * relevance_output + (1 - gate_values) * x
return output, attention_weights, gate_values
# Example usage
batch_size, seq_length = 2, 10
hidden_size = 768
model = VERAModule(hidden_size)
x = torch.randn(seq_length, batch_size, hidden_size)
context = torch.randn(seq_length, batch_size, hidden_size)
output, attn_weights, gates = model(x, context)
print(f"Output shape: {output.shape}")
print(f"Attention weights shape: {attn_weights.shape}")
print(f"Gate values shape: {gates.shape}")
Performance Metrics Implementation
This implementation provides comprehensive evaluation metrics for fine-tuning techniques, including parameter efficiency, adaptation quality, and computational overhead measurements.
import numpy as np
from typing import Dict, List
import time
class FineTuningMetrics:
def __init__(self, base_params: int):
self.base_params = base_params
self.metrics_history: Dict[str, List[float]] = {
'parameter_efficiency': [],
'adaptation_quality': [],
'compute_overhead': []
}
def compute_parameter_efficiency(self, adapted_params: int) -> float:
"""Calculate parameter efficiency ratio"""
efficiency = 1 - (adapted_params / self.base_params)
self.metrics_history['parameter_efficiency'].append(efficiency)
return efficiency
def measure_adaptation_quality(self,
original_output: np.ndarray,
adapted_output: np.ndarray) -> float:
"""Measure output similarity using cosine similarity"""
orig_flat = original_output.reshape(original_output.shape[0], -1)
adapt_flat = adapted_output.reshape(adapted_output.shape[0], -1)
similarities = []
for orig, adapt in zip(orig_flat, adapt_flat):
sim = np.dot(orig, adapt) / (np.linalg.norm(orig) * np.linalg.norm(adapt))
similarities.append(sim)
quality = np.mean(similarities)
self.metrics_history['adaptation_quality'].append(quality)
return quality
def compute_overhead(self, func, *args) -> float:
"""Measure computational overhead"""
start_time = time.time()
func(*args)
overhead = time.time() - start_time
self.metrics_history['compute_overhead'].append(overhead)
return overhead
def get_summary(self) -> Dict[str, float]:
"""Get summary statistics of all metrics"""
return {
metric: {
'mean': np.mean(values),
'std': np.std(values),
'min': np.min(values),
'max': np.max(values)
}
for metric, values in self.metrics_history.items()
}
# Example usage
metrics = FineTuningMetrics(base_params=100_000_000)
# Simulate measurements
adapted_params = 1_000_000
orig_output = np.random.randn(32, 768)
adapted_output = np.random.randn(32, 768)
efficiency = metrics.compute_parameter_efficiency(adapted_params)
quality = metrics.measure_adaptation_quality(orig_output, adapted_output)
overhead = metrics.compute_overhead(np.dot, orig_output, adapted_output.T)
summary = metrics.get_summary()
print("Metrics Summary:")
print(f"Parameter Efficiency: {efficiency:.4f}")
print(f"Adaptation Quality: {quality:.4f}")
print(f"Computational Overhead: {overhead:.4f}s")
Real-world Application - Text Classification
This implementation demonstrates the application of LoRA for fine-tuning a pre-trained model on a text classification task, including data preprocessing, model adaptation, and evaluation metrics tracking.
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from torch.utils.data import Dataset, DataLoader
class CustomDataset(Dataset):
def __init__(self, texts, labels, tokenizer, max_length=128):
self.encodings = tokenizer(texts, truncation=True, padding=True,
max_length=max_length, return_tensors='pt')
self.labels = torch.tensor(labels)
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
return {key: val[idx] for key, val in self.encodings.items()}, self.labels[idx]
class TextClassifierWithLoRA(nn.Module):
def __init__(self, model_name, num_classes, rank=4):
super().__init__()
self.base_model = AutoModel.from_pretrained(model_name)
hidden_size = self.base_model.config.hidden_size
# LoRA layers
self.lora_A = nn.Parameter(torch.randn(hidden_size, rank) / np.sqrt(rank))
self.lora_B = nn.Parameter(torch.randn(rank, hidden_size) / np.sqrt(rank))
# Classification head
self.classifier = nn.Linear(hidden_size, num_classes)
def forward(self, input_ids, attention_mask):
# Get base model outputs
outputs = self.base_model(input_ids=input_ids,
attention_mask=attention_mask)
hidden_states = outputs.last_hidden_state[:, 0, :] # [CLS] token
# Apply LoRA adaptation
lora_output = hidden_states @ self.lora_A @ self.lora_B
adapted_hidden = hidden_states + lora_output
# Classification
logits = self.classifier(adapted_hidden)
return logits
# Example usage
def train_epoch(model, dataloader, optimizer, criterion):
model.train()
total_loss = 0
correct = 0
total = 0
for batch in dataloader:
inputs, labels = batch
optimizer.zero_grad()
outputs = model(**inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
return total_loss / len(dataloader), 100. * correct / total
# Initialize model and training components
model_name = "bert-base-uncased"
num_classes = 2
model = TextClassifierWithLoRA(model_name, num_classes)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Training setup
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
criterion = nn.CrossEntropyLoss()
Real-world Application Results
# Sample training results
training_metrics = {
'epoch': [1, 2, 3, 4, 5],
'train_loss': [0.6932, 0.3845, 0.2756, 0.2134, 0.1876],
'train_acc': [50.23, 82.45, 89.67, 92.34, 93.56],
'val_loss': [0.6821, 0.4012, 0.3123, 0.2897, 0.2765],
'val_acc': [51.12, 80.34, 86.78, 88.92, 89.45]
}
print("Training Results:")
for epoch in range(len(training_metrics['epoch'])):
print(f"Epoch {training_metrics['epoch'][epoch]}:")
print(f" Train Loss: {training_metrics['train_loss'][epoch]:.4f}")
print(f" Train Accuracy: {training_metrics['train_acc'][epoch]:.2f}%")
print(f" Validation Loss: {training_metrics['val_loss'][epoch]:.4f}")
print(f" Validation Accuracy: {training_metrics['val_acc'][epoch]:.2f}%")
print("-" * 50)
# Parameter efficiency analysis
base_params = sum(p.numel() for p in model.base_model.parameters())
lora_params = sum(p.numel() for p in [model.lora_A, model.lora_B])
reduction = (1 - lora_params/base_params) * 100
print(f"\nParameter Efficiency Analysis:")
print(f"Base Model Parameters: {base_params:,}")
print(f"LoRA Parameters: {lora_params:,}")
print(f"Parameter Reduction: {reduction:.2f}%")
Advanced LoRA Training Strategies
This implementation demonstrates sophisticated training techniques for LoRA, including gradient accumulation, mixed-precision training, and dynamic rank adjustment based on task complexity.
import torch
from torch.cuda.amp import autocast, GradScaler
import numpy as np
class AdvancedLoRATrainer:
def __init__(self, model, rank_min=2, rank_max=16,
accumulation_steps=4):
self.model = model
self.rank_bounds = (rank_min, rank_max)
self.accumulation_steps = accumulation_steps
self.scaler = GradScaler()
def dynamic_rank_adjustment(self, validation_loss_history):
"""Adjust LoRA rank based on validation performance"""
if len(validation_loss_history) < 2:
return
loss_delta = validation_loss_history[-1] - validation_loss_history[-2]
current_rank = self.model.lora_A.shape[1]
if loss_delta > 0: # Validation loss increasing
new_rank = min(current_rank * 2, self.rank_bounds[1])
else:
new_rank = max(current_rank // 2, self.rank_bounds[0])
if new_rank != current_rank:
self._resize_lora_matrices(new_rank)
def _resize_lora_matrices(self, new_rank):
"""Resize LoRA matrices while preserving learned patterns"""
with torch.no_grad():
old_A = self.model.lora_A.data
old_B = self.model.lora_B.data
# SVD for rank reduction/expansion
U, S, V = torch.svd(old_A @ old_B)
# Initialize new matrices
new_A = torch.randn(old_A.shape[0], new_rank) / np.sqrt(new_rank)
new_B = torch.randn(new_rank, old_B.shape[1]) / np.sqrt(new_rank)
# Copy existing patterns
min_rank = min(new_rank, old_A.shape[1])
new_A[:, :min_rank] = U[:, :min_rank] @ torch.diag(
torch.sqrt(S[:min_rank]))
new_B[:min_rank, :] = torch.diag(
torch.sqrt(S[:min_rank])) @ V[:min_rank, :]
# Update model parameters
self.model.lora_A = torch.nn.Parameter(new_A)
self.model.lora_B = torch.nn.Parameter(new_B)
def train_step(self, batch, optimizer, criterion):
"""Training step with mixed precision and gradient accumulation"""
optimizer.zero_grad()
accumulated_loss = 0
for micro_batch in torch.split(batch, batch.shape[0] //
self.accumulation_steps):
with autocast():
outputs = self.model(micro_batch)
loss = criterion(outputs) / self.accumulation_steps
self.scaler.scale(loss).backward()
accumulated_loss += loss.item()
self.scaler.step(optimizer)
self.scaler.update()
return accumulated_loss
# Example usage
model = TextClassifierWithLoRA("bert-base-uncased", num_classes=2)
trainer = AdvancedLoRATrainer(model)
# Training loop with metrics
validation_losses = []
batch = torch.randn(32, 512) # Example batch
optimizer = torch.optim.AdamW(model.parameters())
criterion = torch.nn.CrossEntropyLoss()
loss = trainer.train_step(batch, optimizer, criterion)
validation_losses.append(loss)
trainer.dynamic_rank_adjustment(validation_losses)
Optimization Techniques in LoRA
This implementation focuses on advanced optimization strategies for LoRA, including adaptive learning rates, weight decay scheduling, and gradient clipping mechanisms.
import torch
import math
from torch.optim import Optimizer
class LoRAOptimizer(Optimizer):
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999),
weight_decay=0.01, warmup_steps=1000):
defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
super().__init__(params, defaults)
self.warmup_steps = warmup_steps
self.current_step = 0
def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
# Compute warmup factor
warmup_factor = min(1.0, self.current_step / self.warmup_steps)
# Parameter-specific learning rate
if hasattr(p, 'is_lora'):
effective_lr = group['lr'] * math.sqrt(warmup_factor)
else:
effective_lr = group['lr'] * warmup_factor
# Gradient clipping
torch.nn.utils.clip_grad_norm_(p, max_norm=1.0)
# Weight decay with scheduling
if group['weight_decay'] > 0:
p.data.mul_(1 - effective_lr * group['weight_decay'])
# Update parameters
p.data.add_(p.grad, alpha=-effective_lr)
self.current_step += 1
return loss
# Example usage
def configure_optimizer(model, lr=2e-5):
# Separate LoRA and non-LoRA parameters
lora_params = []
base_params = []
for name, param in model.named_parameters():
if 'lora' in name:
param.is_lora = True
lora_params.append(param)
else:
base_params.append(param)
optimizer = LoRAOptimizer([
{'params': base_params},
{'params': lora_params, 'lr': lr * 10} # Higher LR for LoRA
], lr=lr)
return optimizer
# Initialize and use optimizer
model = TextClassifierWithLoRA("bert-base-uncased", num_classes=2)
optimizer = configure_optimizer(model)
Memory-Efficient LoRA Implementation
This implementation showcases advanced memory optimization techniques for LoRA, utilizing gradient checkpointing and efficient memory management strategies for handling large models with limited resources.
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
import gc
class MemoryEfficientLoRA(nn.Module):
def __init__(self, base_model, rank=8, chunk_size=1024):
super().__init__()
self.base_model = base_model
self.hidden_size = base_model.config.hidden_size
self.chunk_size = chunk_size
# Initialize LoRA matrices with memory-efficient storage
self.lora_A = nn.Parameter(
torch.empty(self.hidden_size, rank, dtype=torch.float16)
)
self.lora_B = nn.Parameter(
torch.empty(rank, self.hidden_size, dtype=torch.float16)
)
self.reset_parameters()
def reset_parameters(self):
nn.init.normal_(self.lora_A, std=0.02)
nn.init.normal_(self.lora_B, std=0.02)
def chunked_forward(self, hidden_states):
batch_size = hidden_states.shape[0]
output = torch.zeros_like(hidden_states)
for i in range(0, batch_size, self.chunk_size):
chunk = hidden_states[i:i + self.chunk_size]
def custom_forward(x):
return x @ self.lora_A @ self.lora_B
output[i:i + self.chunk_size] = checkpoint(
custom_forward, chunk, preserve_rng_state=False
)
# Explicit memory cleanup
if i % (self.chunk_size * 4) == 0:
torch.cuda.empty_cache()
gc.collect()
return output
def forward(self, input_ids, attention_mask=None):
base_output = self.base_model(
input_ids=input_ids,
attention_mask=attention_mask
)
hidden_states = base_output.last_hidden_state
# Apply LoRA with memory optimization
lora_output = self.chunked_forward(hidden_states)
return base_output.last_hidden_state + lora_output
# Example usage with memory monitoring
def train_with_memory_tracking(model, dataloader, epochs=1):
model = model.cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for epoch in range(epochs):
for batch in dataloader:
# Track memory usage
torch.cuda.reset_peak_memory_stats()
initial_mem = torch.cuda.memory_allocated()
inputs = {k: v.cuda() for k, v in batch.items()}
outputs = model(**inputs)
loss = outputs.mean() # Example loss
loss.backward()
optimizer.step()
optimizer.zero_grad()
final_mem = torch.cuda.memory_allocated()
peak_mem = torch.cuda.max_memory_allocated()
print(f"Memory Usage:")
print(f"Initial: {initial_mem/1e6:.2f}MB")
print(f"Final: {final_mem/1e6:.2f}MB")
print(f"Peak: {peak_mem/1e6:.2f}MB")
# Create sample data and model
sample_data = [
{"input_ids": torch.randint(0, 1000, (16, 512)),
"attention_mask": torch.ones(16, 512)}
]
sample_dataloader = torch.utils.data.DataLoader(sample_data)
base_model = AutoModel.from_pretrained("bert-base-uncased")
efficient_model = MemoryEfficientLoRA(base_model)
# Train with memory tracking
train_with_memory_tracking(efficient_model, sample_dataloader)
Additional Resources
- ArXiv Papers and References:
- LoRA: Low-Rank Adaptation of Large Language Models https://arxiv.org/abs/2106.09685
- Efficient Fine-Tuning of Language Models https://arxiv.org/abs/2312.15685
- Memory-Efficient Fine-Tuning of Large Language Models https://arxiv.org/abs/2303.16742
- Optimal Transport for Parameter-Efficient Fine-Tuning https://arxiv.org/abs/2310.12962
- A Survey of Deep Learning Approaches for Parameter-Efficient Transfer Learning https://arxiv.org/abs/2311.11017
For further research and implementation details, consider searching:
- "Parameter-efficient fine-tuning techniques"
- "Memory optimization in language models"
- "Low-rank adaptation methods"
- "Efficient transfer learning approaches"