Fine Tuning And Aligning Llms
Parameter-Efficient Fine-Tuning: LoRA & QLoRA
DL Notes, Module 7, Chapter 4, Section 4 previewed LoRA briefly. This chapter covers it — and its quantized extension, QLoRA — in full depth, since it's essenti
JrCodex·7 min read
Jr Codex NLP & LLM Notes
Level: Advanced Prerequisites: Chapter 1: Instruction Tuning Time to complete: ~25 minutes
Table of Contents
- The Memory Problem at LLM Scale
- LoRA's Core Insight
- LoRA, Mathematically
- Implementing LoRA
- QLoRA: Adding Quantization
- A Practical LoRA Fine-Tuning Example
- When to Use Full Fine-Tuning vs LoRA
- Summary & Next Steps
1. The Memory Problem at LLM Scale
DL Notes, Module 7, Chapter 4, Section 4 previewed LoRA briefly. This chapter covers it — and its quantized extension, QLoRA — in full depth, since it's essential for working with modern LLMs practically.
Why Full Fine-Tuning Becomes Impractical
─────────────────────────────────────────
Full fine-tuning (DL Notes, Module 7, Ch.2, Sec.3) requires
storing GRADIENTS and OPTIMIZER STATE (DL Notes, Module 9,
Ch.1, Sec.5 — e.g. Adam's momentum AND squared-gradient
terms) for EVERY SINGLE parameter. For a 7-BILLION parameter
model, this can require well over 100GB of GPU memory just
for training state — far beyond a single consumer, or even
many enterprise, GPUs.
─────────────────────────────────────────
2. LoRA's Core Insight
LoRA (Low-Rank Adaptation) freezes the ENTIRE pretrained model and instead learns a small number of additional parameters that approximate the needed weight update.
The Key Mathematical Insight
─────────────────────────────────────────
When fine-tuning changes a weight matrix W, the ACTUAL
change (ΔW) required tends to have a LOW "intrinsic rank" —
meaning it can be well-APPROXIMATED by the product of two
much SMALLER matrices, rather than needing to represent a
full, dense update of the SAME size as W itself.
─────────────────────────────────────────
LoRA's Decomposition
─────────────────────────────────────────
Instead of learning a FULL ΔW (same shape as W, e.g.
4096 x 4096 = ~16.7 million values):
ΔW ≈ B × A
A: a (r × 4096) matrix r is the "RANK" — a SMALL
B: a (4096 × r) matrix number, e.g. 8 or 16
Total NEW parameters: r×4096 + 4096×r = 2 × r × 4096
(e.g., r=8: ~65,000 parameters — over
250x FEWER than the full 16.7 million)
─────────────────────────────────────────
3. LoRA, Mathematically
import torch
import torch.nn as nn
class LoRALayer(nn.Module):
def __init__(self, original_layer, rank=8, alpha=16):
super().__init__()
self.original_layer = original_layer
for param in self.original_layer.parameters():
param.requires_grad = False # FREEZE the original weights entirely (DL Notes, Module 7, Ch.2)
in_features = original_layer.in_features
out_features = original_layer.out_features
self.lora_A = nn.Parameter(torch.randn(rank, in_features) * 0.01) # SMALL, random init
self.lora_B = nn.Parameter(torch.zeros(out_features, rank)) # ZERO init — CRITICAL (see below)
self.scaling = alpha / rank
def forward(self, x):
original_output = self.original_layer(x) # the FROZEN, pretrained computation
lora_output = (x @ self.lora_A.T) @ self.lora_B.T # the SMALL, LEARNED adjustment
return original_output + self.scaling * lora_output # combine BOTH (DL Notes, Module 4,
# Ch.2's residual connection idea,
# applied here to WEIGHTS)Why lora_B Is Initialized to ZERO
─────────────────────────────────────────
At the very START of fine-tuning, lora_B × lora_A = 0 (since
lora_B is all zeros) — meaning the LoRA-adapted model is
IDENTICAL to the original pretrained model on the FIRST
forward pass. Training then GRADUALLY introduces the learned
adaptation, rather than starting from a RANDOM, potentially
disruptive adjustment — a deliberate, careful initialization
choice (echoing DL Notes, Module 3, Chapter 1's broader
initialization principles).
─────────────────────────────────────────
4. Implementing LoRA
In practice, the HuggingFace peft (Parameter-Efficient Fine-Tuning) library handles this automatically, rather than requiring manual implementation as in Section 3.
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model, TaskType
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
lora_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
r=8, # rank (Section 2)
lora_alpha=16, # scaling factor (Section 3)
lora_dropout=0.1, # DL Notes, Module 3, Ch.4 — applied to the LoRA path specifically
target_modules=["c_attn"], # WHICH layers get LoRA adapters — typically attention layers
# (DL Notes, Module 6, Ch.2)
)
peft_model = get_peft_model(model, lora_config)
peft_model.print_trainable_parameters()
# e.g. "trainable params: 294,912 || all params: 124,734,720 || trainable%: 0.24%"5. QLoRA: Adding Quantization
QLoRA combines LoRA with quantization — storing the FROZEN base model's weights using far fewer bits than standard 32-bit floats, further reducing memory requirements.
Quantization, Briefly
─────────────────────────────────────────
Standard precision (DL Notes, 32-bit floats (FP32) or
Module 9, Ch.2): 16-bit floats (FP16) per
weight
QLoRA's 4-bit quantization: compresses the
FROZEN base
model's weights
down to just 4
bits each — an
8x memory
reduction versus
FP32, while the
LoRA adapters
themselves (the
part actually
being TRAINED)
remain in higher
precision
─────────────────────────────────────────
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16, # compute happens in FP16 even though STORAGE is 4-bit
bnb_4bit_quant_type="nf4", # a quantization scheme optimized for neural network
# weight distributions specifically
)
model = AutoModelForCausalLM.from_pretrained(
"gpt2-large",
quantization_config=quantization_config,
device_map="auto",
)
# THEN apply LoRA on top of the quantized model, exactly as in Section 4
peft_model = get_peft_model(model, lora_config)Why This Combination Is So Effective
─────────────────────────────────────────
QLoRA's TWO techniques address DIFFERENT parts of the
memory problem: quantization shrinks the FROZEN base model's
STORAGE footprint (Section 1's biggest cost), while LoRA
shrinks the NUMBER OF TRAINABLE parameters (and therefore
the gradient/optimizer memory, Section 1). Combined, QLoRA
can fine-tune models with tens of billions of parameters on
a SINGLE consumer-grade GPU — a genuinely dramatic
accessibility improvement.
─────────────────────────────────────────
6. A Practical LoRA Fine-Tuning Example
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model, TaskType
from datasets import Dataset
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(model_name)
lora_config = LoraConfig(task_type=TaskType.CAUSAL_LM, r=8, lora_alpha=16, target_modules=["c_attn"])
peft_model = get_peft_model(model, lora_config)
# Reusing Chapter 1's instruction dataset and preprocessing directly
dataset = Dataset.from_list(instruction_examples).map(preprocess)
training_args = TrainingArguments(
output_dir="./lora-instruction-tuned",
num_train_epochs=3,
per_device_train_batch_size=4, # can afford a LARGER batch — far less memory used per step
learning_rate=2e-4, # LoRA fine-tuning often uses a SLIGHTLY higher rate than
# full fine-tuning, since fewer parameters are being adjusted
)
trainer = Trainer(model=peft_model, args=training_args, train_dataset=dataset)
trainer.train()
peft_model.save_pretrained("./lora-adapter-only") # saves ONLY the small LoRA weights — MEGABYTES,
# not GIGABYTES, since the base model is UNCHANGED7. When to Use Full Fine-Tuning vs LoRA
A Practical Decision Guide
─────────────────────────────────────────
Small model, PLENTY of Full fine-tuning
GPU memory available: (DL Notes, Module 7)
remains a
reasonable, simpler
default
Large model (billions of LoRA/QLoRA — often the
parameters), LIMITED ONLY practically
hardware: feasible option
Need to MAINTAIN MULTIPLE LoRA — swap
task-specific "versions" of DIFFERENT
the SAME base model: small
adapters
on and
off the
SAME
frozen
base
model,
rather
than
storing
multiple
FULL
copies
Maximum possible performance, Full
compute/storage NOT a concern: fine-
tuning
can
sometimes
slightly
OUTPERFORM
LoRA,
though the
gap is
often small
─────────────────────────────────────────
8. Summary & Next Steps
Key Takeaways
- Full fine-tuning of large LLMs requires storing gradients and optimizer state for every parameter, quickly becoming impractical at billions of parameters.
- LoRA freezes the entire pretrained model and learns a low-rank approximation of the needed weight update, dramatically reducing trainable parameters while initializing to an identical starting point via zero-initialized B matrices.
- QLoRA combines LoRA with 4-bit quantization of the frozen base model, addressing both storage and trainable-parameter memory costs simultaneously.
- LoRA's structural independence from the base model also enables swapping different task-specific adapters on and off a single frozen model, rather than storing multiple full model copies.
Concept Check
- Why is lora_B initialized to all zeros rather than small random values, unlike lora_A?
- What two distinct memory problems do LoRA and quantization each solve, and why does combining them (QLoRA) address both?
- Name a scenario where maintaining multiple LoRA adapters is more practical than maintaining multiple fully fine-tuned model copies.
Next Chapter
Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to NLP & LLM Index