LoRA Implementation › Integrating LoRA into Transformer Models
Language AI HandbookAll chapters LoRA Implementation › Integrating LoRA into Transformer ModelsOne section · Sections in this chapter4 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Real transformer models contain many linear layers distributed throughout their architecture: the query, key, and value projections in each attention head, the output projection that recombines attention results, and the feed-forward networks that process each position independently. A typical LoRA configuration targets the attention projections, as these capture the most task-relevant transformations by controlling what information the model attends to and how it combines that information.
Think of the decision about which layers to target as similar to deciding which parts of a recipe to modify when adapting a dish for a new cuisine. The spices and seasonings, analogous to the attention query and value projections, have the most influence on flavor and are the natural place to start. The structural components like the cooking method and base ingredients, analogous to the feed-forward networks and embeddings, matter too, but changing them has broader effects that are harder to control. LoRA practitioners typically start with attention projections and add more target layers only when the initial adaptation falls short.
Applying LoRA to a full model requires identifying and replacing layers without manually editing the architecture. We need an automated approach that can traverse the model's module hierarchy, find layers matching our target criteria, and wrap them with LoRA adapters while preserving the model's structure. The traversal must handle the hierarchical naming scheme that PyTorch uses for nested modules, where a query projection might live at a path like layers.0.attention.q_proj.
Let's create a utility function that recursively replaces specified linear layers with LoRA versions:
def inject_lora(
model: nn.Module,
target_modules: list[str],
rank: int = 4,
alpha: float = 1.0,
dropout: float = 0.0,
) -> nn.Module:
"""
Inject LoRA adapters into specified modules of a model.
Args:
model: The base model to modify
target_modules: List of module name patterns to target
rank: LoRA rank
alpha: LoRA scaling factor
dropout: Dropout probability on LoRA path
Returns:
Modified model with LoRA layers injected
"""
# Collect modules to replace first to avoid modifying graph during traversal
modules_to_replace = []
for name, module in model.named_modules():
if any(target in name for target in target_modules) and isinstance(
module, nn.Linear
):
modules_to_replace.append((name, module))
for name, module in modules_to_replace:
# Create LoRA wrapper
lora_layer = LoRALinear(module, rank=rank, alpha=alpha, dropout=dropout)
# Replace the module in the parent
# Navigate to parent module
parts = name.split(".")
parent = model
for part in parts[:-1]:
if part.isdigit():
parent = parent[int(part)]
else:
parent = getattr(parent, part)
# Replace the layer
if parts[-1].isdigit():
parent[int(parts[-1])] = lora_layer
else:
setattr(parent, parts[-1], lora_layer)
return modelThis function implements a recursive traversal strategy that examines every module in the model hierarchy. For each module, it checks whether the module name contains any of the target patterns, allowing flexible specification of which layers to adapt. The function then verifies that the module is a linear layer (since LoRA applies specifically to linear transformations), creates a LoRA wrapper around it, and performs the replacement by locating the parent module and updating the attribute. This pattern-based approach means you can target layers by their role (like q_proj for query projections) rather than their exact position in the model structure.
One subtle but important design choice is that we collect all modules to replace before performing any replacements. If we replaced modules during traversal, the iterator could become confused by changes to the module graph mid-traversal, leading to missed replacements or double-replacements. Separating the collection phase from the replacement phase avoids this pitfall entirely. This two-pass approach is a common pattern in code that modifies a data structure while iterating over it.
Let's show this with a simplified transformer block that contains the needed components found in production models:
class SimpleAttention(nn.Module):
"""Simplified attention for demonstration."""
def __init__(self, hidden_size: int, num_heads: int):
super().__init__()
self.hidden_size = hidden_size
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.q_proj = nn.Linear(hidden_size, hidden_size)
self.k_proj = nn.Linear(hidden_size, hidden_size)
self.v_proj = nn.Linear(hidden_size, hidden_size)
self.o_proj = nn.Linear(hidden_size, hidden_size)
def forward(self, x):
batch, seq_len, _ = x.shape
q = self.q_proj(x)
k = self.k_proj(x)
v = self.v_proj(x)
# Reshape for multi-head attention
q = q.view(batch, seq_len, self.num_heads, self.head_dim).transpose(
1, 2
)
k = k.view(batch, seq_len, self.num_heads, self.head_dim).transpose(
1, 2
)
v = v.view(batch, seq_len, self.num_heads, self.head_dim).transpose(
1, 2
)
# Scaled dot-product attention
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
attn_weights = F.softmax(scores, dim=-1)
attn_output = torch.matmul(attn_weights, v)
# Reshape and project
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(batch, seq_len, self.hidden_size)
return self.o_proj(attn_output)
class SimpleFeedForward(nn.Module):
"""Simple FFN for demonstration."""
def __init__(self, hidden_size: int, intermediate_size: int):
super().__init__()
self.up_proj = nn.Linear(hidden_size, intermediate_size)
self.down_proj = nn.Linear(intermediate_size, hidden_size)
self.act = nn.GELU()
def forward(self, x):
return self.down_proj(self.act(self.up_proj(x)))
class SimpleTransformerBlock(nn.Module):
"""A basic transformer block."""
def __init__(self, hidden_size: int = 256, num_heads: int = 4):
super().__init__()
self.attention = SimpleAttention(hidden_size, num_heads)
self.ffn = SimpleFeedForward(hidden_size, hidden_size * 4)
self.norm1 = nn.LayerNorm(hidden_size)
self.norm2 = nn.LayerNorm(hidden_size)
def forward(self, x):
x = x + self.attention(self.norm1(x))
x = x + self.ffn(self.norm2(x))
return xNow we can inject LoRA into the attention projections, showing how the injection process affects parameter counts and training behavior:
# Create a transformer block
transformer = SimpleTransformerBlock(hidden_size=256, num_heads=4)
# Count parameters before LoRA
total_before = sum(p.numel() for p in transformer.parameters())
trainable_before = sum(
p.numel() for p in transformer.parameters() if p.requires_grad
)
# Freeze the complete base model before adding trainable adapters.
for parameter in transformer.parameters():
parameter.requires_grad = False
# Inject LoRA into query and value projections
transformer = inject_lora(
transformer, target_modules=["q_proj", "v_proj"], rank=8, alpha=8.0
)
# Count parameters after LoRA
total_after = sum(p.numel() for p in transformer.parameters())
trainable_after = sum(
p.numel() for p in transformer.parameters() if p.requires_grad
)Parameter counts: Before LoRA - Total: 789,760, Trainable: 789,760 After LoRA - Total: 797,952, Trainable: 8,192 Trainable parameter reduction: 99.0%
The model now has the same representational capacity as before, since all original weights are preserved and can contribute to the forward pass. However, only the 8,192 LoRA parameters are trainable, about 1.0% of the adapted block's 797,952 total parameters. This sharply reduces the parameters that require gradient computation and optimizer state storage. It shows LoRA's core value proposition: the ability to fine-tune large models on limited hardware by putting all learning into a small, efficient parameter set.
Why Target Only Query and Value Projections?
The choice to target query () and value () projections rather than key () or output () projections is supported by empirical evidence from the original LoRA paper and subsequent studies. The query projection determines what each token is "looking for" in the sequence, so adapting it changes the attention patterns the model forms. The value projection determines what information is extracted from each token once attention weights are computed, so adapting it changes what content the model reads out. Together, adapting and gives the model the most control over the information flow through the attention mechanism.
Adapting the key projection () is less common because the key and query projections jointly determine attention weights: changes what each token "offers" while changes what each token "seeks." Modifying both offers diminishing returns because the attention weight between tokens and is determined by the inner product ; adapting only one side of this product often suffices to shift attention patterns. The output projection () recombines the attended values, and adapting it overlaps functionally with adapting . In practice, targeting and reaches a good balance between adaptation capacity and parameter count, though targeting all four projections or including feed-forward layers may be warranted for tasks that require substantial behavioral changes.