# How LoRA is Implemented for Efficient Transformer Fine-Tuning in PyTorch

> Learn how LoRA fine-tuning efficiently adapts transformers in PyTorch. Freeze pre-trained weights and inject low-rank matrices drastically reducing parameters.

- Repository: [labml.ai/annotated_deep_learning_paper_implementations](https://github.com/labmlai/annotated_deep_learning_paper_implementations)
- Tags: how-to-guide
- Published: 2026-03-04

---

**LoRA (Low-Rank Adaptation) enables efficient transformer fine-tuning by freezing pre-trained weight matrices and injecting trainable low-rank matrices into linear projections, reducing trainable parameters by orders of magnitude while maintaining full model capacity.**

The `labmlai/annotated_deep_learning_paper_implementations` repository demonstrates exactly **how LoRA is implemented for efficient transformer fine-tuning in PyTorch** through a minimal, educational codebase. This implementation modifies standard `nn.Linear` and `nn.Embedding` layers to freeze original weights and learn low-rank decomposition matrices instead.

## Core LoRA Layer Mechanics

The foundation resides in [`labml_nn/lora/__init__.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/__init__.py), which reimplements PyTorch's linear and embedding layers with frozen base weights and trainable adapters.

### Freezing Pre-Trained Weights

In [`labml_nn/lora/__init__.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/__init__.py) lines 64-68 (for `Linear`) and lines 26-29 (for `Embedding`), the original weight parameters are initialized with `requires_grad=False`:

```python

# Linear layer implementation (simplified)

self.weight = nn.Parameter(torch.empty((out_features, in_features)), requires_grad=False)

```

This ensures gradient computation skips the massive pre-trained matrices during backpropagation.

### Injecting Low-Rank Adapters

Two small trainable matrices replace the full weight update. Lines 79-83 (Linear) and 33-35 (Embedding) initialize:

- `self.lora_a` with shape `(r, in_features)` — initialized with Gaussian noise
- `self.lora_b` with shape `(out_features, r)` — initialized with zeros

Where `r` is the LoRA rank (typically 4-64), dramatically smaller than the original dimensions.

### Scaling the Low-Rank Update

The implementation applies a scaling factor `alpha / r` to control adapter influence. Lines 60-63 (Linear) and 22-25 (Embedding) store:

```python
self.scaling = alpha / r

```

This keeps the magnitude of the adaptation comparable to the original weight updates, preventing instability during fine-tuning.

### The Forward Pass Computation

The critical logic appears in lines 90-96 (Linear) and 43-48 (Embedding), computing:

```python

# Original frozen path + low-rank adapter path

original = F.linear(x, self.weight, bias=self.bias)
adapter = F.linear(F.linear(x, self.lora_a), self.lora_b) * self.scaling
return original + adapter

```

Mathematically, this implements \(h = xW_0 + \frac{\alpha}{r}xA^{\top}B^{\top}\), where \(W_0\) remains frozen and \(A, B\) are trainable.

## Integrating LoRA into GPT-2 Transformers

The [`labml_nn/lora/gpt2.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/gpt2.py) file adapts the GPT-2 architecture by replacing every standard projection layer with LoRA-enabled equivalents. Within each transformer block:

- **Query, Key, Value projections**: use `Linear(d_model, d_model, r=lora_rank)`
- **Embedding layers**: use `Embedding(vocab_size, d_model, r=lora_rank)`
- **Layer normalization and attention mechanisms**: remain unchanged

This surgical replacement ensures only specific weight matrices receive gradient updates, while the bulk of the 124M+ parameter model stays frozen.

## Loading Pre-Trained Weights and Optimization Strategy

The [`labml_nn/lora/experiment.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/experiment.py) script handles weight loading and training configuration. Lines 64-68, 98-113, and 122-128 demonstrate loading Hugging Face GPT-2 checkpoints:

```python

# Load pre-trained state dict

hf_model = AutoModelForCausalLM.from_pretrained("gpt2")
state_dict = hf_model.state_dict()

# Copy weights with strict=False to allow missing LoRA parameters

model.load_state_dict(mapped_weights, strict=False)

```

The `strict=False` parameter is essential — it ignores the newly initialized `lora_a` and `lora_b` parameters that don't exist in the original checkpoint.

### Training Only the Adapters

Line 38 creates an optimizer receiving all parameters via `model.parameters()`, but only the LoRA matrices compute gradients:

```python
optimizer = Adam(model.parameters(), lr=learning_rate)

```

During the training loop (lines 48-64), `loss.backward()` propagates gradients exclusively through `lora_a` and `lora_b`, leaving the frozen base weights untouched. This reduces training memory requirements and checkpoint sizes from gigabytes to megabytes.

## Complete Fine-Tuning Workflow

Combine these components to fine-tune GPT-2 on a custom dataset:

```python
from labml_nn.lora.gpt2 import GPTModel
from torch.optim import Adam
import torch.nn.functional as F

# 1. Initialize LoRA-augmented model

model = GPTModel(
    d_model=768,
    n_heads=12,
    n_layers=12,
    vocab_size=50257,
    r=32,  # LoRA rank

).to(device)

# 2. Load pre-trained weights (base weights frozen, LoRA params initialized)

# Implementation details in labml_nn/lora/experiment.py lines 64-113

# 3. Setup optimizer — only LoRA parameters receive gradients

optimizer = Adam(model.parameters(), lr=1e-4)

# 4. Training loop

for epoch in range(num_epochs):
    for batch in dataloader:
        inputs = batch[0].to(device)
        logits = model(inputs[:, :-1])
        loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)),
            inputs[:, 1:].reshape(-1)
        )
        
        optimizer.zero_grad()
        loss.backward()  # Gradients flow only through lora_a and lora_b

        optimizer.step()

```

This workflow fine-tunes the model on new tasks while preserving the original pre-trained knowledge in the frozen weights.

## Summary

- **LoRA modifies linear projections** in [`labml_nn/lora/__init__.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/__init__.py) by freezing `self.weight` at lines 64-68 and injecting trainable `lora_a`/`lora_b` matrices at lines 79-83.
- **Scaling factor** `alpha/r` (lines 60-63) controls the magnitude of low-rank updates during the forward pass (lines 90-96).
- **GPT-2 integration** in [`labml_nn/lora/gpt2.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/gpt2.py) replaces standard layers while keeping attention mechanisms and layer norms unchanged.
- **Memory efficiency** comes from loading pre-trained weights with `strict=False` (lines 64-113 in [`experiment.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/experiment.py)) and optimizing only the small adapter matrices.
- **Training speed** improves because backpropagation skips the massive frozen base weights, computing gradients only for the low-rank decomposition matrices.

## Frequently Asked Questions

### What makes LoRA more parameter-efficient than full fine-tuning?

Full fine-tuning updates all weight matrices in a transformer (millions or billions of parameters), whereas LoRA freezes these matrices and trains only the low-rank decomposition matrices \(A\) and \(B\). As implemented in [`labml_nn/lora/__init__.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/__init__.py), if the original weight has shape `(768, 768)` and rank `r=32`, LoRA trains only `768×32 + 32×768 = 49,152` parameters instead of the full 589,824 — a 12x reduction for that single layer.

### Which layers in the transformer should use LoRA adapters?

According to the [`labml_nn/lora/gpt2.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/gpt2.py) implementation, LoRA adapters are applied to **query, key, value projections** and **embedding layers** — essentially every linear transformation that projects between the model dimension and itself. Layer normalization parameters and attention softmax operations remain frozen, as these typically contain less task-specific information.

### How does the scaling factor alpha/r affect model performance?

The scaling factor, defined in lines 60-63 of [`labml_nn/lora/__init__.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/__init__.py), multiplies the low-rank update by `alpha/r` (defaulting to `1.0` when `alpha=r`). This hyperparameter controls the learning rate of the adaptation relative to the frozen pre-trained weights. Higher `alpha` values increase the influence of the LoRA adapters during the forward pass (lines 90-96), allowing faster adaptation but potentially destabilizing training if set too high.

### Can this implementation be applied to architectures other than GPT-2?

Yes. The `Linear` and `Embedding` classes in [`labml_nn/lora/__init__.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/__init__.py) are drop-in replacements for standard PyTorch layers. Any transformer architecture using `nn.Linear` for projections — including BERT, T5, or Vision Transformers — can integrate these classes by replacing the standard layers, exactly as demonstrated in [`labml_nn/lora/gpt2.py`](https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/main/labml_nn/lora/gpt2.py) for the GPT-2 architecture.