Chapter 3: Implementing LoRA with PEFT

Hands-On LoRA Training

This chapter provides complete, production-ready code for implementing LoRA fine-tuning using the Hugging Face PEFT library.

Complete Training Script

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
)
from peft import (
    LoraConfig,
    get_peft_model,
    TaskType,
    prepare_model_for_kbit_training,
)
from trl import SFTTrainer
from datasets import load_dataset

# ============================================
# Configuration
# ============================================

MODEL_NAME = "meta-llama/Llama-3.1-8B"
OUTPUT_DIR = "./lora-output"
DATASET = "databricks/dolly-15k"

# LoRA hyperparameters
LORA_R = 16
LORA_ALPHA = 32
LORA_DROPOUT = 0.05
TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]

# Training hyperparameters
LEARNING_RATE = 2e-4
BATCH_SIZE = 4
GRADIENT_ACCUMULATION = 4
MAX_SEQ_LENGTH = 512
NUM_EPOCHS = 3

# ============================================
# Load Model and Tokenizer
# ============================================

print("Loading model and tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# ============================================
# Configure LoRA
# ============================================

print("Configuring LoRA...")

lora_config = LoraConfig(
    r=LORA_R,
    lora_alpha=LORA_ALPHA,
    lora_dropout=LORA_DROPOUT,
    target_modules=TARGET_MODULES,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)

# Apply LoRA
model = get_peft_model(model, lora_config)

# Print trainable parameters
model.print_trainable_parameters()

# ============================================
# Prepare Dataset
# ============================================

print("Loading dataset...")

dataset = load_dataset(DATASET, split="train")

def format_instruction(example):
    """Format data for instruction tuning."""
    if example.get("context"):
        text = f"""### Instruction:
{example['instruction']}

### Context:
{example['context']}

### Response:
{example['response']}"""
    else:
        text = f"""### Instruction:
{example['instruction']}

### Response:
{example['response']}"""
    return {"text": text}

dataset = dataset.map(format_instruction)

# ============================================
# Training Arguments
# ============================================

training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRADIENT_ACCUMULATION,
    learning_rate=LEARNING_RATE,
    weight_decay=0.01,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    logging_steps=10,
    save_strategy="epoch",
    bf16=True,
    gradient_checkpointing=True,
    optim="adamw_torch_fused",
    report_to="none",  # or "wandb"
    seed=42,
)

# ============================================
# Create Trainer
# ============================================

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    args=training_args,
    dataset_text_field="text",
    max_seq_length=MAX_SEQ_LENGTH,
    tokenizer=tokenizer,
    packing=False,
)

# ============================================
# Train
# ============================================

print("Starting training...")
trainer.train()

# ============================================
# Save
# ============================================

print("Saving model...")
model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

print(f"Training complete! Model saved to {OUTPUT_DIR}")

LoraConfig Parameters Explained

Core Parameters

LoraConfig(
    # Rank - capacity of adaptation
    r=16,  # Typical: 8, 16, 32, 64
    
    # Scaling factor
    lora_alpha=32,  # Typical: 1x to 2x of r
    
    # Regularization
    lora_dropout=0.05,  # Typical: 0.0 to 0.1
    
    # Which layers to adapt
    target_modules=["q_proj", "v_proj"],
    
    # Bias training
    bias="none",  # "none", "all", or "lora_only"
    
    # Task type
    task_type=TaskType.CAUSAL_LM,
)

Target Modules by Model

# Llama / Mistral
TARGET_MODULES_LLAMA = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj",
]

# GPT-2 / GPT-J
TARGET_MODULES_GPT = [
    "c_attn", "c_proj", "c_fc",
]

# BERT / RoBERTa
TARGET_MODULES_BERT = [
    "query", "key", "value",
    "dense",
]

Loading and Using a Trained Adapter

Basic Loading

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")

# Load adapter
model = PeftModel.from_pretrained(base_model, "./lora-output")

# Generate
inputs = tokenizer("Write a haiku about coding:", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Merging Adapter

# Merge for deployment (no adapter overhead)
model = model.merge_and_unload()

# Save merged model
model.save_pretrained("./merged-model")
tokenizer.save_pretrained("./merged-model")

Multi-Adapter Training

Training Multiple Adapters

from peft import LoraConfig, get_peft_model

# Create base PEFT model
model = get_peft_model(base_model, lora_config)

# Train first adapter
trainer.train()
model.save_pretrained("./adapter-task-a")

# Add another adapter for a different task
model.add_adapter("task_b", lora_config)
model.set_adapter("task_b")

# Train second adapter
trainer.train()  # With different dataset
model.save_pretrained("./adapter-task-b", save_adapter="task_b")

Switching Adapters

# Load base with first adapter
model = PeftModel.from_pretrained(base_model, "./adapter-task-a")

# Load additional adapter
model.load_adapter("./adapter-task-b", adapter_name="task_b")

# Switch between adapters
model.set_adapter("default")  # Use task A
output_a = model.generate(...)

model.set_adapter("task_b")  # Use task B
output_b = model.generate(...)

Hyperparameter Tuning Guide

Rank Selection

# Start here
rank_experiments = [8, 16, 32]

for r in rank_experiments:
    config = LoraConfig(r=r, lora_alpha=r*2, ...)
    # Train and evaluate
    # Compare quality vs training time

Learning Rate

# LoRA typically uses higher LR than full FT
lr_options = [1e-4, 2e-4, 5e-4]

# Start with 2e-4, adjust based on:
# - Loss not decreasing → increase LR
# - Training unstable → decrease LR

Target Modules

# Minimal (fastest, lowest quality)
minimal = ["q_proj", "v_proj"]

# Standard (good balance)
standard = ["q_proj", "k_proj", "v_proj", "o_proj"]

# Full (highest quality, most params)
full = ["q_proj", "k_proj", "v_proj", "o_proj", 
        "gate_proj", "up_proj", "down_proj"]

Debugging Common Issues

Issue: Loss Not Decreasing

# Check 1: Is adapter trainable?
model.print_trainable_parameters()
# Should show non-zero trainable params

# Check 2: Learning rate too low?
# Try increasing: 1e-4 → 2e-4 → 5e-4

# Check 3: Data formatting correct?
print(dataset[0]["text"])  # Verify format

Issue: NaN Loss

# Solution 1: Lower learning rate
learning_rate = 1e-5

# Solution 2: Use gradient clipping
training_args = TrainingArguments(
    max_grad_norm=1.0,
    ...
)

# Solution 3: Check for inf/nan in data

Issue: Out of Memory

# Solution 1: Reduce batch size
per_device_train_batch_size = 1

# Solution 2: Increase gradient accumulation
gradient_accumulation_steps = 16

# Solution 3: Enable gradient checkpointing
gradient_checkpointing = True

# Solution 4: Use QLoRA instead (next module)

Production Checklist

□ Choose appropriate rank (start with 16)
□ Set alpha to 2x rank
□ Target all attention projections at minimum
□ Use appropriate learning rate (1e-4 to 5e-4)
□ Enable gradient checkpointing for memory
□ Set up evaluation metrics
□ Save checkpoints regularly
□ Merge adapter for deployment
□ Test merged model thoroughly

Key Takeaways

✅ PEFT library makes LoRA implementation straightforward

✅ Key config: r (rank), alpha (scaling), target_modules

✅ SFTTrainer simplifies the training loop

✅ Adapters can be saved/loaded/merged independently

✅ Multiple adapters enable task switching

✅ Start with standard settings, tune as needed

Next Steps

Test your knowledge with the Module 4 Quiz, then continue to Module 5: QLoRA and Quantization.

Quick Check

Before moving on, make sure you can answer:

What function creates a PEFT model from a base model?
How do you merge a LoRA adapter into the base model?
What is a typical learning rate for LoRA training?

Hands-On LoRA Training​

Complete Training Script​

LoraConfig Parameters Explained​

Core Parameters​

Target Modules by Model​

Loading and Using a Trained Adapter​

Basic Loading​

Merging Adapter​

Multi-Adapter Training​

Training Multiple Adapters​

Switching Adapters​

Hyperparameter Tuning Guide​

Rank Selection​

Learning Rate​

Target Modules​

Debugging Common Issues​

Issue: Loss Not Decreasing​

Issue: NaN Loss​

Issue: Out of Memory​

Production Checklist​

Key Takeaways​

Next Steps​

Quick Check​