Suppose a deployed model starts missing a task it needs to handle, but retraining it from scratch is too expensive. Low-rank adaptation (LoRA) offers a smaller update: train an adapter for the task while leaving the original model weights unchanged. The adapter can be saved separately, making task-specific versions easier to maintain. section 9.3.2 section 5.3.5
Ordinary weight updating can be pictured as W = W + ΔW. LoRA represents the update as the product of two narrower matrices: ΔW = Wa × Wb. For a 100 × 100 update with rank 2, those matrices contain 400 trainable elements rather than 10,000. This constrains the update; it does not make every possible 10,000-element update available. Rank is a practical trade-off: increasing it allows a more complex update but saves less memory. section 5.3.5
LoRA belongs to parameter-efficient fine-tuning (PEFT). It is especially useful when a base model already does useful work and needs a focused change for a new task or domain. The approach is not a guarantee of good output: training data still matters, and a worked model update produced a noticeable change without yet producing sufficiently coherent text. section 5.3.5 section 5.5 section 9.3.2
A small update alongside a linear layer
A PyTorch implementation creates two trainable matrices, A and B, with shapes input × rank and rank × output. Its forward pass adds a scaled update to the existing layer’s output: linear(x) + alpha × (x @ A @ B). The example initializes B to zero, so the adapter initially contributes no change. section 9.3.2
After training, keep the model and LoRA as separate state dictionaries. For inference, either load the adapter over the base model or merge its layers into the model. A separate adapter is useful when maintaining distinct task updates; merging produces a model that can be used without loading the adapter separately. section 9.3.2 section 5.3.5
Listing 9.6 Low-rank adaptation
class LoRALayer(nn.Module): #1
def __init__(self, in_dim, out_dim, rank, alpha):
super().__init__()
standard_deviation = 1 / torch.sqrt(torch.tensor(rank).float())
self.A = nn.Parameter(
torch.randn(in_dim, rank) * standard_deviation
)
self.B = nn.Parameter(torch.zeros(rank, out_dim))
self.alpha = alpha
def forward(self, x):
x = self.alpha * (x @ self.A @ self.B)
return x
class LinearWithLoRA(nn.Module):
def __init__(self, linear, rank, alpha):
super().__init__()
self.linear = linear
self.lora = LoRALayer(
linear.in_features, linear.out_features, rank, alpha
)
def forward(self, x):
return self.linear(x) + self.lora(x)
class LlamaBlock(nn.Module): #2
def __init__(self, config):
super().__init__()
self.config = config
self.rms = RMSNormalization(
(config["d_model"], config["d_model"])
).to(device)
self.attention = RoPEMaskedMultiheadAttention(config).to(device)
self.feedforward = nn.Sequential(
LinearWithLoRA(config["d_model"], config["d_model"]), #3
SwiGLU(config["d_model"]),
).to(device)
def forward(self, x):
x = self.rms(x)
x = x + self.attention(x)
x = self.rms(x)
x = x + self.feedforward(x)
return x
class SimpleLlama(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.embedding = nn.Embedding(
config["vocab_size"], config["d_model"]
)
self.llama_blocks = nn.Sequential(
OrderedDict(
[
(f"llama_{i}", LlamaBlock(config))
for i in range(config["n_layers"])
]
)
)
self.ffn = nn.Sequential(
LinearWithLoRA(config["d_model"], config["d_model"]), #4
SwiGLU(config["d_model"]),
LinearWithLoRA(config["d_model"], config["vocab_size"]), #5
)
print(
f"model params: {sum([m.numel() for m in self.parameters()])}"
)
def forward(self, idx, targets=None):
x = self.embedding(idx)
x = self.llama_blocks(x)
logits = self.ffn(x)
if targets is None:
return logits
else:
loss = F.cross_entropy(
logits.view(-1, self.config["vocab_size"]),
targets.view(-1),
ignore_index=tokenizer.pad_token_id,
reduction="sum",
)
return logits, loss
dataset = load_dataset( #6
"text",
data_files={
"train": ["../../data/Lima-train.csv"],
"val": ["../../data/Lima-test.csv"],
},
streaming=True,
)
encoded_dataset = dataset.map(
lambda examples: tokenizer(
examples["text"],
padding=True,
max_length=128,
truncation=True,
\]] return_tensors="pt",
),
batched=True,
)
train_data = iter(encoded_dataset["train"].shuffle())
val_data = iter(encoded_dataset["val"].shuffle())
train_data = cycle(train_data)
val_data = cycle(val_data)
llama.to("cpu") #7
add_lora(llama)
llama.to(device)
parameters = [{"params": list(get_lora_params(llama))}] #8
lora_optimizer = torch.optim.AdamW(parameters, lr=1e-3) #9
train( #10
llama,
lora_optimizer,
scheduler,
data=train_data,
config=MASTER_CONFIG,
lora=True,
print_logs=True,
)
state_dict = llama.state_dict() #11
lora_state_dict = {k: v for k, v in state_dict.items() if name_is_lora(k)}
torch.save(llama.state_dict(), "./llama.pth")
torch.save(lora_state_dict, "./lora.pth")
Inspect the PyTorch LoRA layer and wrapper to see exactly where the trainable update joins an existing linear layer. section 9.3.2
