Fine-Tuning Small Language Models (SLMs) with LoRA and Unsloth

Small Language Models (SLMs) like Llama 3.2, Gemma 2, and Phi-3 have made domain-specific fine-tuning accessible without needing enterprise cluster infrastructure. By leveraging Low-Rank Adaptation (LoRA) and optimized memory frameworks like Unsloth, you can fine-tune 3B to 8B models on a single consumer GPU (e.g., RTX 3090 / 4090).

1. Mathematical Intuition Behind LoRA

Low-Rank Adaptation freezes the pre-trained model weights $W_0 \in \mathbb{R}^{d \times k}$ and injects trainable rank decomposition matrices.

For a dense linear layer with weight update matrix $\Delta W$, LoRA decomposes $\Delta W$ into two low-rank matrices $A$ and $B$:

\[\Delta W = B \cdot A\]

Where $B \in \mathbb{R}^{d \times r}$ and $A \in \mathbb{R}^{r \times k}$, and the rank $r \ll \min(d, k)$.

During forward pass computation, the output $h$ is given by:

\[h = W_0 x + \Delta W x = W_0 x + \frac{\alpha}{r} B A x\]

Where $\alpha$ is a constant hyperparameter that scales the learned weights.

2. Environment Setup

To get started with Unsloth and Hugging Face trl:

pip install --upgrade pip pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes

3. Loading the Base Model in 4-bit Precision

from unsloth import FastLanguageModel import torch max_seq_length = 2048 dtype = None # Auto-detection (bfloat16 for Ampere+) load_in_4bit = True # Use 4bit quantization to reduce VRAM model, tokenizer = FastLanguageModel.from_pretrained( model_name = "unsloth/Llama-3.2-3B-Instruct", max_seq_length = max_seq_length, dtype = dtype, load_in_4bit = load_in_4bit, )

4. Applying LoRA Target Modules

Next, configure the LoRA adapters across the attention projection layers (q_proj, k_proj, v_proj, o_proj):

model = FastLanguageModel.get_peft_model( model, r = 16, # Rank dimension target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], lora_alpha = 16, lora_dropout = 0, # Optimized 0 for Unsloth bias = "none", use_gradient_checkpointing = "unsloth", # 30% VRAM savings random_state = 3407, )

5. Training Loop with SFTTrainer

from trl import SFTTrainer from transformers import TrainingArguments trainer = SFTTrainer( model = model, tokenizer = tokenizer, train_dataset = dataset, dataset_text_field = "text", max_seq_length = max_seq_length, dataset_num_proc = 2, args = TrainingArguments( per_device_train_batch_size = 2, gradient_accumulation_steps = 4, warmup_steps = 10, max_steps = 60, learning_rate = 2e-4, fp16 = not torch.cuda.is_bf16_supported(), bf16 = torch.cuda.is_bf16_supported(), logging_steps = 1, optim = "adamw_8bit", weight_decay = 0.01, lr_scheduler_type = "linear", seed = 3407, output_dir = "outputs", ), ) trainer.train()

6. Key Learnings & MLOps Considerations

• Rank Selection ($r$): For instruction tuning, $r=16$ or $r=32$ provides an optimal balance between parameter capacity and speed.
• Merging Weights: LoRA adapters can be merged back into the base model weights $W_{new} = W_0 + \frac{\alpha}{r}BA$ for zero-latency inference serving via vLLM or Ollama.