LoRA Implementation › The LoRA Training Loop
Language AI HandbookAll chapters LoRA Implementation › The LoRA Training LoopOne section · Sections in this chapter5 / 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
Training a LoRA-equipped model follows the standard PyTorch training pattern, with one key simplification: the optimizer only needs to handle the LoRA parameters. This means smaller optimizer state memory, faster parameter updates, and reduced gradient computation since gradients don't flow through the frozen base weights.
When you pass only LoRA parameters to an Adam optimizer, the memory benefit is significant. Adam maintains two running statistics per parameter: the first moment (the exponential moving average of gradients) and the second moment (the exponential moving average of squared gradients). For a 7-billion parameter model, full fine-tuning with Adam requires roughly GB just for parameters and optimizer states in float32. With LoRA targeting attention projections at rank 16, the trainable parameters drop to roughly 8 million, requiring less than 100 MB for optimizer states. This difference in memory requirement is what makes it possible to fine-tune on a single 24GB GPU instead of a cluster of eight.
Let's build a complete training example to illustrate these concepts in action:
def get_lora_parameters(model: nn.Module):
"""Extract only the LoRA parameters from a model."""
lora_params = []
for name, param in model.named_parameters():
if "lora_" in name:
lora_params.append(param)
return lora_params
def count_trainable_parameters(model: nn.Module) -> tuple[int, int]:
"""Count total and trainable parameters."""
total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
return total, trainableThe get_lora_parameters function is a simple but important utility. Rather than relying on requires_grad flags (which would work equally well), it filters by name, selecting only parameters whose name contains lora_. This makes the selection criterion explicit and readable: you can see at a glance that the optimizer will handle exactly the LoRA matrices and nothing else. The name-based filter also means you can add additional trainable parameters (like a task-specific classification head) without accidentally including them in LoRA-specific operations.
Now let's implement a training loop for a simple sequence classification task, which shows how LoRA integrates into a realistic training workflow:
class SequenceClassifier(nn.Module):
"""Simple classifier using transformer blocks."""
def __init__(
self,
vocab_size: int,
hidden_size: int = 256,
num_heads: int = 4,
num_layers: int = 2,
num_classes: int = 2,
):
super().__init__()
self.embedding = nn.Embedding(vocab_size, hidden_size)
self.pos_embedding = nn.Embedding(512, hidden_size)
self.layers = nn.ModuleList(
[
SimpleTransformerBlock(hidden_size, num_heads)
for _ in range(num_layers)
]
)
self.classifier = nn.Linear(hidden_size, num_classes)
def forward(self, input_ids):
batch_size, seq_len = input_ids.shape
positions = torch.arange(seq_len, device=input_ids.device)
x = self.embedding(input_ids) + self.pos_embedding(positions)
for layer in self.layers:
x = layer(x)
# Pool: take mean of sequence
pooled = x.mean(dim=1)
return self.classifier(pooled)from torch.utils.data import DataLoader, TensorDataset
# Create synthetic training data
torch.manual_seed(42)
vocab_size = 1000
num_samples = 500
seq_length = 32
# Generate "sentences" with a learnable pattern:
# class 1 sequences have tokens drawn from the upper half of the vocab (500-999),
# class 0 sequences from the lower half (0-499), with 15% noise tokens to add
# realistic overlap so the loss curve shows real convergence.
train_inputs = torch.zeros(num_samples, seq_length, dtype=torch.long)
train_labels = torch.randint(0, 2, (num_samples,))
noise_mask = torch.rand(num_samples, seq_length) < 0.15
for i in range(num_samples):
if train_labels[i] == 1:
train_inputs[i] = torch.randint(500, vocab_size, (seq_length,))
else:
train_inputs[i] = torch.randint(0, 500, (seq_length,))
# Apply noise: flip ~15% of tokens to the opposite range
train_inputs[noise_mask] = torch.randint(
0, vocab_size, (noise_mask.sum().item(),)
)
# Create data loader
dataset = TensorDataset(train_inputs, train_labels)
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)# Initialize model
model = SequenceClassifier(vocab_size=vocab_size, num_classes=2)
# Calculate baseline parameter counts
total, trainable = count_trainable_parameters(model)Before LoRA: Total parameters: 1,967,106 Trainable parameters: 1,967,106
At this stage, all parameters in the model are trainable. This is the baseline for full fine-tuning, which would require updating every weight in the network. For large models, this baseline is a significant memory burden since the optimizer must store momentum and variance estimates (in the case of Adam-family optimizers) for every parameter.
# Inject LoRA into attention projections across all layers
model = inject_lora(
model, target_modules=["q_proj", "v_proj"], rank=4, alpha=4.0
)
# Verify parameter counts
total, trainable = count_trainable_parameters(model)After LoRA injection: Total parameters: 1,975,298 Trainable parameters: 1,712,130 Percentage trainable: 86.68%
After injecting LoRA, the number of trainable parameters drops precipitously. The large majority of the model, comprising the original weights from embeddings, attention layers, feed-forward networks, and the classifier, is now frozen. Only the small adapter matrices remain trainable. This transformation is what lets LoRA to fine-tune models that would otherwise exceed available memory.
# Setup optimizer with only LoRA parameters
lora_params = get_lora_parameters(model)
optimizer = torch.optim.AdamW(lora_params, lr=1e-3, weight_decay=0.01)
criterion = nn.CrossEntropyLoss()
# Training loop
num_epochs = 5
train_losses = []
model.train()
for epoch in range(num_epochs):
epoch_loss = 0.0
num_batches = 0
for batch_inputs, batch_labels in train_loader:
optimizer.zero_grad()
outputs = model(batch_inputs)
loss = criterion(outputs, batch_labels)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
num_batches += 1
avg_loss = epoch_loss / num_batches
train_losses.append(avg_loss)Training progress: Epoch 1: Loss = 0.7006 Epoch 2: Loss = 0.6610 Epoch 3: Loss = 0.6073 Epoch 4: Loss = 0.5122 Epoch 5: Loss = 0.3411
The training loop shows that LoRA integrates naturally with standard PyTorch workflows. The code structure is virtually identical to full fine-tuning, with only two modifications required. First, we inject LoRA into target modules using our utility function. Second, we pass only LoRA parameters to the optimizer. This keeps gradient updates only apply to the adaptation matrices. Everything else proceeds exactly as in conventional training: the forward pass computes the loss, backpropagation computes gradients, and the optimizer steps.
The backpropagation through a LoRA layer deserves a moment's explanation. When loss.backward() runs, PyTorch traverses the computation graph in reverse. The gradient flows through the LoRA path as follows: first through the scaling factor, then through the linear transform with , then through the linear transform with . Both and accumulate gradients normally because they are leaf tensors with requires_grad=True. The gradient also flows into the original layer's output, but since original_layer.weight.requires_grad=False, PyTorch stops the gradient computation there. This selective gradient flow saves memory: you allocate gradient buffers only for parameters that need them.
Let's visualize the training progress:
