Fine-Tune a Small Model with LoRA (TRL + PEFT)

LoRA fine-tuning trains a small set of low-rank adapter matrices while the base model stays frozen. It is the right tool for changing form and behaviour at volume: output format, tone, a fixed response structure. It is the wrong tool for teaching facts. Read Fine-tuning vs RAG vs prompting before you commit.

This runs a 0.5B instruct model through TRL’s SFTTrainer with a PEFT LoraConfig on a toy dataset. Install trl, peft, datasets and transformers first. The dataset uses the conversational prompt-completion shape, which SFT accepts and renders with the model’s chat template. Real data needs far more rows and cleaning: see Prepare a dataset for fine-tuning.

import torch
from datasets import Dataset
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer

base = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(base)
model = AutoModelForCausalLM.from_pretrained(base, dtype=torch.bfloat16)

# Toy behaviour to learn: every answer starts with "Short answer:".
pairs = [
    ("What is a GPU?", "Short answer: a chip built for parallel math."),
    ("What is an embedding?", "Short answer: a vector that encodes meaning."),
    ("What is LoRA?", "Short answer: small trainable matrices added to a frozen model."),
    ("What is a token?", "Short answer: a chunk of text the model reads as one unit."),
]
dataset = Dataset.from_list([
    {"prompt": [{"role": "user", "content": q}],
     "completion": [{"role": "assistant", "content": a}]}
    for q, a in pairs * 25
])

trainer = SFTTrainer(
    model=model,
    args=SFTConfig(
        output_dir="lora-out",
        learning_rate=1e-4,
        num_train_epochs=3,
        per_device_train_batch_size=4,
        logging_steps=5,
    ),
    train_dataset=dataset,
    processing_class=tokenizer,
    peft_config=LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules="all-linear",
        task_type="CAUSAL_LM",
    ),
)
trainer.train()
trainer.model.save_pretrained("lora-adapter")

target_modules="all-linear" targets every linear layer, so you don’t name modules per architecture. The adapter scale is lora_alpha / r, so r=16, lora_alpha=32 scales by 2. TRL’s docs suggest a higher learning rate (around 1e-4) when only adapter weights train; the SFTConfig default is 2e-5. SFTConfig turns bf16 on by default when fp16 is unset, so use bf16-capable hardware, or set bf16=False and drop the dtype argument.

The folder holds only the adapter, not the base model. Load it for inference:

from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
model = AutoPeftModelForCausalLM.from_pretrained("lora-adapter")
model.eval()

inputs = tokenizer.apply_chat_template(
    [{"role": "user", "content": "What is a vector database?"}],
    tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True,
).to(model.device)
out = model.generate(**inputs, max_new_tokens=40)
print(tokenizer.decode(out[0], skip_special_tokens=True))

What to watch

• LoRA doesn’t add knowledge. Expect it to learn style and structure, not new facts. Use retrieval for facts.
• An adapter is not a merged model. It needs the same base at load time. merge_and_unload() is not in-place: use the model it returns. To share the adapter, see Upload a model or dataset to Hugging Face.
• Change one knob at a time. Higher r means more trainable parameters; changing r or lora_alpha changes the scale. If target_modules is unset, PEFT picks modules by architecture and raises an error for an unknown one.
• Memory. A 0.5B model is small. For bigger bases, see QLoRA on a consumer GPU, and fix CUDA out of memory when a run fails.

Reference: TRL SFTTrainer and PEFT LoRA.

---

Next: QLoRA on a consumer GPU, Upload a model or dataset to Hugging Face, Prepare a dataset for fine-tuning