Deeply Learning 4: LoRA (Low Rank Adaptation) - Math, Intuition and Implementation from Scratch

!pip install transformers datasets accelerate torchvision -q

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
import copy
import time
from dataclasses import dataclass

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')

import logging
logging.getLogger("transformers").setLevel(logging.ERROR)

from transformers import GPT2Model

model = GPT2Model.from_pretrained("gpt2")
total_params = sum(p.numel() for p in model.parameters())

print(f"GPT-2 small total parameters: {total_params:,}")
print(f"Memory for full fine-tuning (FP32, with AdamW): {total_params * 4 * 4 / 1e9:.2f} GB")
print(f"Memory for 10 task-specific checkpoints: {total_params * 4 * 10 / 1e9:.2f} GB")
print()
print(f"Now scale that to 70B parameters (LLaMA 3.1 70B, mixed precision):")
print(f"  Fine-tuning memory: {70e9 * 12 / 1e9:.0f} GB  (params + grads in FP16, optimizer states in FP32)")
print(f"  10 task checkpoints: {70e9 * 2 * 10 / 1e9:.0f} GB")

del model

from torchvision import datasets, transforms

torch.manual_seed(42)

mnist_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
mnist_train = datasets.MNIST('./data', train=True, download=True, transform=mnist_transform)
mnist_loader = DataLoader(mnist_train, batch_size=128, shuffle=True)

d_in = 28 * 28  # flattened MNIST images
d_hidden = 256

class MNISTNet(nn.Module):
    def __init__(self, d_in=784, d_hidden=256, d_out=10):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)
        self.fc2 = nn.Linear(d_hidden, d_out)

    def forward(self, x: torch.Tensor):
        x = x.flatten(1)
        return self.fc2(F.relu(self.fc1(x)))

net = MNISTNet(d_in, d_hidden, d_out=10).to(device)
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
for epoch in range(3):
    total_loss, correct, total = 0, 0, 0
    for images, labels in mnist_loader:
        images, labels = images.to(device), labels.to(device)
        logits = net(images)
        loss = F.cross_entropy(logits, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * labels.size(0)
        correct += (logits.detach().argmax(1) == labels).sum().item()
        total += labels.size(0)
    print(f"Pre-train epoch {epoch+1}/3 | Loss: {total_loss/total:.4f} | Acc: {correct/total:.2%}")

W_pretrained = {name: p.data.clone() for name, p in net.named_parameters()}

# Fine-tune: odd vs. even (binary task that reuses digit features)
optimizer = torch.optim.Adam(net.parameters(), lr=1e-4)
for epoch in range(2):
    total_loss, correct, total = 0, 0, 0
    for images, labels in mnist_loader:
        images, labels = images.to(device), (labels % 2).to(device)
        logits = net(images)[:, :2]
        loss = F.cross_entropy(logits, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * labels.size(0)
        correct += (logits.argmax(1) == labels).sum().item()
        total += labels.size(0)
    print(f"Fine-tune epoch {epoch+1}/2 | Loss: {total_loss/total:.4f} | Acc: {correct/total:.2%}")

# Compute delta_W and its singular values
delta_W = net.fc1.weight.data - W_pretrained['fc1.weight']
U, S, V = torch.linalg.svd(delta_W)

S_normalized = S / S.sum()
cumulative = torch.cumsum(S_normalized, dim=0)

print(f"\nShape of delta_W (fc1): {delta_W.shape}")
print(f"Full rank: {min(delta_W.shape)}")
print(f"\nSingular value energy concentration:")
for r in [1, 2, 4, 8, 16]:
    print(f"  Top {r:2d} singular values capture {cumulative[r-1]:.1%} of total energy")

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

ax1.bar(range(len(S)), S.cpu().numpy(), color='steelblue', alpha=0.8)
ax1.set_xlabel('Singular Value Index')
ax1.set_ylabel('Singular Value')
ax1.set_title('Singular Values of ΔW (fc1)')
ax1.set_xlim(-1, len(S))

ax2.plot(range(1, len(cumulative)+1), cumulative.cpu().numpy(),
         'o-', color='steelblue', markersize=3)
ax2.axhline(y=0.95, color='red', linestyle='--', alpha=0.7, label='95% energy')
ax2.set_xlabel('Rank r')
ax2.set_ylabel('Cumulative Energy')
ax2.set_title('Cumulative Energy vs. Rank')
ax2.legend()
ax2.set_xlim(0, len(S))

plt.tight_layout()
plt.show()

# Find rank needed for 95% energy
rank_95 = (cumulative >= 0.95).nonzero(as_tuple=True)[0][0].item() + 1
print(f"\nRank needed to capture 95% of ΔW energy: {rank_95} out of {min(delta_W.shape)}")
print(f"That's only {rank_95/min(delta_W.shape):.1%} of the full rank!")

# GPT-2 small architecture
d_model = 768
n_layers = 12
n_heads = 12
d_head = d_model // n_heads  # 64

# Each attention layer has Wq, Wk, Wv, Wo, each of shape (d_model, d_model)
# We'll apply LoRA to Wq and Wv

params_per_target_layer = d_model * d_model  # 768 * 768 = 589,824
full_ft_params = 2 * n_layers * params_per_target_layer  # Wq and Wv across all layers

print("=== Full Fine-Tuning (Wq + Wv only) ===")
print(f"Parameters per matrix: {params_per_target_layer:,}")
print(f"Total trainable (Wq + Wv, all layers): {full_ft_params:,}")

print("\n=== LoRA ===")
for r in [1, 2, 4, 8, 16]:
    lora_per_matrix = d_model * r + r * d_model  # B (d_model x r) + A (r x d_model)
    lora_total = 2 * n_layers * lora_per_matrix
    reduction = full_ft_params / lora_total
    print(f"  r={r:2d}: {lora_total:>8,} params ({reduction:5.1f}× reduction, "
          f"{lora_total/full_ft_params:.2%} of full)")

from transformers.pytorch_utils import Conv1D

class LoRALinear(nn.Module):
  """A Linear layer augmented with LoRA low-rank adaptation.

  Implements: h = W_0 @ x + (alpha / r) * B @ A @ x
  where W_0 is frozen and only A, B are trained.

  Supports both nn.Linear and HuggingFace's Conv1D (used by GPT-2).
  Conv1D stores weights as (in_features, out_features), transposed
  relative to nn.Linear's (out_features, in_features).
  """
  def __init__(self, original_layer: nn.Linear | Conv1D, r: int = 8, alpha: float = 16.0):
    super().__init__()
    self.original_layer = original_layer
    self.r = r
    self.alpha = alpha
    self.scaling = alpha / r
    self.is_conv1d = isinstance(original_layer, Conv1D)

    if self.is_conv1d:
        # Conv1D weight: (in_features, out_features)
        self.in_features = original_layer.weight.shape[0]
        self.out_features = original_layer.nf
    else:
        self.in_features = original_layer.in_features
        self.out_features = original_layer.out_features

    # Freeze the original weight
    self.original_layer.weight.requires_grad = False
    if self.original_layer.bias is not None:
        self.original_layer.bias.requires_grad = False

    # LoRA matrices
    # A: (r, in_features) - projects input down to rank r
    # B: (out_features, r) - projects back up to output dim
    self.lora_A = nn.Parameter(torch.empty(r, self.in_features))
    self.lora_B = nn.Parameter(torch.zeros(self.out_features, r))

    nn.init.kaiming_uniform_(self.lora_A, a=5**0.5)
    # B is initialized to zero, so delta_W = BA = 0 at init

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    h = self.original_layer(x)

    # LoRA branch: x @ A^T @ B^T
    # x: (..., in_features) -> (..., r) -> (..., out_features)
    lora_out = x @ self.lora_A.T @ self.lora_B.T

    return h + self.scaling * lora_out

  def merge_weights(self) -> nn.Module:
    """Merge LoRA weights back into the original layer.
    After merging, this layer behaves identically to the original
    with no additional computation cost.
    """
    with torch.no_grad():
        delta_W = self.lora_B @ self.lora_A  # (out_features, in_features)
        if self.is_conv1d:
            # Conv1D weight is (in_features, out_features), so transpose
            self.original_layer.weight.add_(self.scaling * delta_W.T)
        else:
            # nn.Linear weight is (out_features, in_features)
            self.original_layer.weight.add_(self.scaling * delta_W)
    return self.original_layer

  @property
  def lora_params(self) -> int:
    return self.lora_A.numel() + self.lora_B.numel()

  def __repr__(self) -> str:
    layer_type = "Conv1D" if self.is_conv1d else "Linear"
    return (f"LoRALinear({layer_type}, in={self.in_features}, "
            f"out={self.out_features}, "
            f"r={self.r}, alpha={self.alpha}, "
            f"lora_params={self.lora_params:,})")

# Verification: LoRA at init == original (test both nn.Linear and Conv1D)
torch.manual_seed(42)
x = torch.randn(2, 10, 768)  # (batch, seq_len, d_model)

for layer_type, original in [
    ("nn.Linear", nn.Linear(768, 768)),
    ("Conv1D", Conv1D(768, 768)),
]:
    lora_layer = LoRALinear(original, r=8, alpha=16)
    with torch.no_grad():
        diff = (original(x) - lora_layer(x)).abs().max().item()
    trainable = sum(p.numel() for p in lora_layer.parameters() if p.requires_grad)
    frozen = sum(p.numel() for p in lora_layer.parameters() if not p.requires_grad)
    print(f"{layer_type}: max diff at init = {diff:.2e}, "
          f"trainable = {trainable:,}, frozen = {frozen:,}")

def inject_lora(
  model: nn.Module,
  target_modules: list[str],
  r: int = 8,
  alpha: float = 16.0
  ) -> nn.Module:
  """Replace target nn.Linear layers with LoRALinear layers.

  Args:
      model: Pre-trained model to inject LoRA into
      target_modules: List of substrings to match against layer names.
                      e.g., ['c_attn'] for GPT-2's combined QKV projection,
                      or ['q_proj', 'v_proj'] for LLaMA-style models.
      r: LoRA rank
      alpha: LoRA scaling factor

  Returns:
      Model with LoRA layers injected (modified in-place)
  """
  for name, module in model.named_modules():
    if not isinstance(module, (nn.Linear, Conv1D)):
      continue
    if not any(target in name for target in target_modules):
      continue

    # Navigate to parent module and replace the child
    parts = name.split('.')
    parent = model
    for part in parts[:-1]:
      parent = getattr(parent, part)

    original_linear = getattr(parent, parts[-1])
    lora_linear = LoRALinear(original_linear, r=r, alpha=alpha)
    setattr(parent, parts[-1], lora_linear)

  # Freeze all non-LoRA parameters
  for name, param in model.named_parameters():
    if 'lora_' not in name:
        param.requires_grad = False

  return model

def print_trainable_params(model: nn.Module):
  """Print trainable vs total parameter counts."""
  trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
  total = sum(p.numel() for p in model.parameters())
  print(f"Trainable: {trainable:>10,} ({trainable/total:.2%})")
  print(f"Frozen:    {total - trainable:>10,} ({(total-trainable)/total:.2%})")
  print(f"Total:     {total:>10,}")
  return trainable, total

from datasets import load_dataset
from transformers import GPT2Tokenizer

# Load SST-2
dataset = load_dataset('glue', 'sst2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
tokenizer.pad_token = tokenizer.eos_token

print(f"Train size: {len(dataset['train']):,}")
print(f"Validation size: {len(dataset['validation']):,}")
print(f"\nExample: {dataset['train'][4]}")

MAX_LEN = 128

def tokenize_fn(examples):
  tokens = tokenizer(
      examples['sentence'],
      padding='max_length',
      truncation=True,
      max_length=MAX_LEN,
      return_tensors=None #return lists for datasets
  )

  tokens['labels'] = examples['label']
  return tokens

tokenized = dataset.map(tokenize_fn, batched=True, remove_columns=['sentence', 'idx'])
tokenized.set_format('torch')

```python
train_subset = tokenized['train']
val_subset = tokenized['validation']

train_loader = DataLoader(train_subset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_subset, batch_size=64)

from transformers import GPT2Model

class GPT2ForClassification(nn.Module):
  """GPT-2 with a classification head on top.
  Uses the last non-padding token's hidden state for classification.
  """
  def __init__(self, num_labels=2):
      super().__init__()
      self.gpt2 = GPT2Model.from_pretrained('gpt2')
      self.classifier = nn.Linear(self.gpt2.config.n_embd, num_labels)
      self.num_labels = num_labels

  def forward(
      self,
      input_ids: torch.Tensor,
      attention_mask: torch.Tensor,
      labels: torch.Tensor | None = None
  ) -> tuple[torch.Tensor | None, torch.Tensor]:
      outputs = self.gpt2(input_ids=input_ids, attention_mask=attention_mask)
      hidden_states = outputs.last_hidden_state  # (batch, seq_len, d_model)

      # Get the last non-padding token for each sequence
      # attention_mask: 1 for real tokens, 0 for padding
      # Sum along seq dim to get lengths, subtract 1 for 0-indexing
      seq_lengths = attention_mask.sum(dim=1) - 1  # (batch,)
      batch_idx = torch.arange(hidden_states.size(0), device=hidden_states.device)
      pooled = hidden_states[batch_idx, seq_lengths]  # (batch, d_model)

      logits = self.classifier(pooled)  # (batch, num_labels)

      loss = None
      if labels is not None:
          loss = F.cross_entropy(logits, labels)

      return loss, logits

@torch.no_grad()
def evaluate(
  model: nn.Module,
  dataloader: DataLoader,
  device: torch.device
) -> tuple[float, float]:
  model.eval()
  correct = 0
  total = 0
  total_loss = 0
  for batch in dataloader:
      input_ids = batch['input_ids'].to(device)
      attention_mask = batch['attention_mask'].to(device)
      labels = batch['labels'].to(device)

      loss, logits = model(input_ids, attention_mask, labels)
      preds = logits.argmax(dim=-1)
      correct += (preds == labels).sum().item()
      total += labels.size(0)
      total_loss += loss.item() * labels.size(0)

  return correct / total, total_loss / total


def train_model(
  model: nn.Module,
  train_loader: DataLoader,
  val_loader: DataLoader,
  epochs: int = 3,
  lr: float = 2e-5,
  device: str | torch.device = 'cuda'
) -> dict[str, list[float]]:
  model = model.to(device)
  optimizer = torch.optim.AdamW(
      filter(lambda p: p.requires_grad, model.parameters()),
      lr=lr, weight_decay=0.01
  )

  history: dict[str, list[float]] = {
      'train_loss': [], 'val_loss': [], 'val_acc': [], 'epoch_time': []
  }

  for epoch in range(epochs):
      model.train()
      total_loss = 0
      n_batches = 0
      start_time = time.time()

      for batch in train_loader:
          input_ids = batch['input_ids'].to(device)
          attention_mask = batch['attention_mask'].to(device)
          labels = batch['labels'].to(device)

          loss, _ = model(input_ids, attention_mask, labels)

          optimizer.zero_grad()
          loss.backward()
          torch.nn.utils.clip_grad_norm_(
              filter(lambda p: p.requires_grad, model.parameters()), 1.0
          )
          optimizer.step()

          total_loss += loss.item()
          n_batches += 1

      epoch_time = time.time() - start_time
      avg_train_loss = total_loss / n_batches
      val_acc, val_loss = evaluate(model, val_loader, device)

      history['train_loss'].append(avg_train_loss)
      history['val_loss'].append(val_loss)
      history['val_acc'].append(val_acc)
      history['epoch_time'].append(epoch_time)

      print(f"Epoch {epoch+1}/{epochs} | "
            f"Train Loss: {avg_train_loss:.4f} | "
            f"Val Loss: {val_loss:.4f} | "
            f"Val Acc: {val_acc:.4f} | "
            f"Time: {epoch_time:.1f}s")

  return history

print("=" * 60)
print("EXPERIMENT 1: Full Fine-Tuning")
print("=" * 60)

model_full = GPT2ForClassification(num_labels=2)
print("\nParameter counts:")
print_trainable_params(model_full)

history_full = train_model(model_full, train_loader, val_loader, epochs=3, lr=2e-5, device=device)

print("=" * 60)
print("EXPERIMENT 2: LoRA (r=8)")
print("=" * 60)

model_lora8 = GPT2ForClassification(num_labels=2)

# Let's first see the layer names to know what to target
# Inject LoRA into attention + MLP projections
# c_attn = combined QKV, attn.c_proj = attention output, mlp.c_proj = MLP output
print("\nAttention layer names in GPT-2:")
for name, module in model_lora8.named_modules():
  if isinstance(module, (nn.Linear, Conv1D)) and 'attn' in name:
    if isinstance(module, Conv1D):
        print(f"  {name} (Conv1D): {module.weight.shape[0]} -> {module.nf}")
    else:
        print(f"  {name} (Linear): {module.in_features} -> {module.out_features}")

# Inject LoRA into attention projections
# c_attn is the combined QKV projection, c_proj is the output projection
model_lora8 = inject_lora(model_lora8, target_modules=['c_attn', 'c_proj'], r=8, alpha=16)

# The classification head should also be trainable
for param in model_lora8.classifier.parameters():
  param.requires_grad = True

print("\nParameter counts after LoRA injection:")
print_trainable_params(model_lora8)

# Show the LoRA layers
print("\nLoRA layers:")
for name, module in model_lora8.named_modules():
  if isinstance(module, LoRALinear):
    print(f"  {name}: {module}")

history_lora8 = train_model(model_lora8, train_loader, val_loader, epochs=3, lr=3e-4, device=device)

print("=" * 60)
print("EXPERIMENT 3: LoRA (r=1)")
print("=" * 60)

model_lora1 = GPT2ForClassification(num_labels=2)
model_lora1 = inject_lora(model_lora1, target_modules=['c_attn', 'c_proj'], r=1, alpha=8)

for param in model_lora1.classifier.parameters():
    param.requires_grad = True

print("\nParameter counts after LoRA injection (r=1):")
print_trainable_params(model_lora1)

history_lora1 = train_model(model_lora1, train_loader, val_loader, epochs=3, lr=3e-4, device=device)

# Comparison table
results = {
    'Full Fine-Tune': {
        'val_acc': max(history_full['val_acc']),
        'params': sum(p.numel() for p in model_full.parameters() if p.requires_grad),
        'time_per_epoch': np.mean(history_full['epoch_time']),
    },
    'LoRA (r=8)': {
        'val_acc': max(history_lora8['val_acc']),
        'params': sum(p.numel() for p in model_lora8.parameters() if p.requires_grad),
        'time_per_epoch': np.mean(history_lora8['epoch_time']),
    },
    'LoRA (r=1)': {
        'val_acc': max(history_lora1['val_acc']),
        'params': sum(p.numel() for p in model_lora1.parameters() if p.requires_grad),
        'time_per_epoch': np.mean(history_lora1['epoch_time']),
    },
}

print(f"{'Method':<20} {'Val Acc':>10} {'Trainable':>12} {'% of Full':>10} {'Time/Epoch':>12}")
print("-" * 66)
full_params = results['Full Fine-Tune']['params']
for method, r in results.items():
    print(f"{method:<20} {r['val_acc']:>10.4f} {r['params']:>12,} "
          f"{r['params']/full_params:>10.2%} {r['time_per_epoch']:>10.1f}s")

# Demonstrate weight merging
# First, verify outputs match before and after merging

model_lora8.eval()
test_input_ids = torch.randint(0, 50257, (1, 32), device=device)
test_attention_mask = torch.ones(1, 32, device=device, dtype=torch.long)

with torch.no_grad():
    _, logits_before = model_lora8(test_input_ids, test_attention_mask)

# Merge all LoRA weights
merge_count = 0
for name, module in model_lora8.named_modules():
    if isinstance(module, LoRALinear):
        # Get parent and attribute name
        parts = name.split('.')
        parent = model_lora8
        for part in parts[:-1]:
            parent = getattr(parent, part)

        # Merge and replace with standard Linear
        merged_linear = module.merge_weights()
        setattr(parent, parts[-1], merged_linear)
        merge_count += 1

print(f"Merged {merge_count} LoRA layers back into base weights")

with torch.no_grad():
    _, logits_after = model_lora8(test_input_ids, test_attention_mask)

diff = (logits_before - logits_after).abs().max().item()
print(f"Max logit difference after merging: {diff:.2e}")

# After merging, there are no more LoRA layers
lora_layers = [n for n, m in model_lora8.named_modules() if isinstance(m, LoRALinear)]
print(f"\nRemaining LoRA layers: {len(lora_layers)} (should be 0)")
print("The model is now a standard GPT-2 with no extra inference cost!")

# We need a fresh LoRA model for analysis (since we merged the previous one)
# Let's train a quick one and analyze it

model_analysis = GPT2ForClassification(num_labels=2)
model_analysis = inject_lora(model_analysis, target_modules=['c_attn', 'c_proj'], r=8, alpha=16)
for param in model_analysis.classifier.parameters():
  param.requires_grad = True

# Quick train
print("Training for analysis...")
_ = train_model(model_analysis, train_loader, val_loader, epochs=3, lr=3e-4, device=device)

!pip install peft -q

from peft import LoraConfig, get_peft_model, TaskType
from transformers import GPT2ForSequenceClassification

# Using HuggingFace's built-in classification model
hf_model = GPT2ForSequenceClassification.from_pretrained('gpt2', num_labels=2)
hf_model.config.pad_token_id = tokenizer.pad_token_id

# Configure LoRA with PEFT
peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    r=8,
    lora_alpha=16,
    lora_dropout=0.0,
    target_modules=['c_attn', 'c_proj'],  # Same targets as our implementation
)

hf_model_peft = get_peft_model(hf_model, peft_config)

print("HuggingFace PEFT LoRA:")
hf_model_peft.print_trainable_parameters()

print("\nOur implementation:")
print_trainable_params(model_analysis)