LLM fine-tuning

Deep dive

Decision table

SFT

1. Format every example with the base model's exact chat template (special tokens, roles, EOS at the end of each assistant turn).
2. Loss only on assistant tokens; mask system and user tokens (Llama 2-Chat: LR 2e-5, weight decay 0.1, batch 64, 4096-token sequences, 2 epochs, answer-only loss).
3. Pack several short examples per sequence with cross-example attention masked, or bucket by length; padding wastes compute.
4. 1-3 epochs; more overfits style and hurts general ability. Hold out 2-5% for eval loss.

LoRA and QLoRA

For a frozen weight W [k, n] (y = x·W), train A [k, r] and B [r, n]:

y = x·W + (α/r)·(x·A)·B        A ~ small random,  B = 0  →  ΔW = 0 at step 0

• Trainable params r·(k + n) per matrix instead of k·n. Apply to all attention and MLP projections (q, k, v, o, gate, up, down): QLoRA found this necessary to match full fine-tuning; the original paper's q, v only falls short.
• Rank 8-64; α = r to 2r. With α/r scaling, raising r shrinks each direction's update; rsLoRA uses α/√r so higher ranks keep learning.
• Merge for serving: W ← W + (α/r)·A·B. Unmerged adapters cost an extra small GEMM per layer but allow per-request adapter switching.
• QLoRA: base weights in 4-bit NormalFloat (NF4, 16 levels at normal-distribution quantiles) with one absmax scale per 64 weights. Double quantization stores those scales in 8 bits per block of 256, cutting overhead from 0.5 to 0.127 bits/param. Weights are dequantized to bf16 tile by tile inside each matmul; gradients flow through the frozen quantized weights into the bf16 adapters. Paged optimizer state spills to CPU on memory spikes. The paper's defaults: r = 64, α = 16, LR 2e-4 for 7-13B.

DPO

Given prompt x, chosen y_w, rejected y_l, policy π_θ and a frozen reference π_ref (the SFT model):

z = β·[(log π_θ(y_w|x) − log π_ref(y_w|x)) − (log π_θ(y_l|x) − log π_ref(y_l|x))]
L = −log σ(z)

log π(y|x) is the sum of token log-probs over the response only. β (typically 0.1) sets how far the policy may move from the reference: lower β allows more drift. Precompute reference log-probs once: saves a forward pass per step and the reference model's memory. Use LR around 5e-7 (Zephyr) to 5e-6, 1-3 epochs. Variants: IPO (squared loss, resists overfitting deterministic preferences), SimPO (length-normalized, no reference model), ORPO (SFT and preference in one stage).

Evaluation

• Held-out loss on assistant tokens for SFT; reward accuracy (fraction of pairs with z > 0) and chosen/rejected log-prob trends for DPO. If both log-probs fall, the policy is drifting away from the data distribution.
• Task metrics on your own held-out set first; then general suites before vs after (MMLU, GSM8K, HumanEval, IFEval) to catch forgetting.
• LLM-as-judge (MT-Bench, AlpacaEval 2 length-controlled) with swapped answer order, and length reported alongside: judges favour longer answers and the first position.
• Decontaminate: n-gram overlap between training data and eval prompts.

Code

LoRA layer forward/backward/merge, answer-only SFT cross-entropy and the DPO loss with its gradient. The checks are a finite-difference gradient test, merge equivalence, and DPO loss = ln 2 when policy equals reference.

// LoRA linear layer (forward, backward, merge), masked SFT loss and the DPO loss on plain slices.
// Row-major, y = x·W: x [t, k], W [k, n] frozen, A [k, r], B [r, n], scale s = alpha / r.
pub struct Lora { pub k: usize, pub n: usize, pub r: usize, pub s: f32, pub w: Vec<f32>, pub a: Vec<f32>, pub b: Vec<f32> }

fn mm(x: &[f32], w: &[f32], t: usize, k: usize, n: usize) -> Vec<f32> { // [t,k]·[k,n]
    let mut y = vec![0.0; t * n];
    for i in 0..t { for p in 0..k { let a = x[i * k + p]; for j in 0..n { y[i * n + j] += a * w[p * n + j]; } } }
    y
}
fn mm_tn(x: &[f32], y: &[f32], t: usize, k: usize, n: usize) -> Vec<f32> { // xᵀ·y: [k,t]·[t,n]
    let mut g = vec![0.0; k * n];
    for i in 0..t { for p in 0..k { let a = x[i * k + p]; for j in 0..n { g[p * n + j] += a * y[i * n + j]; } } }
    g
}
fn mm_nt(y: &[f32], w: &[f32], t: usize, n: usize, k: usize) -> Vec<f32> { // y·wᵀ: [t,n]·[n,k]
    let mut g = vec![0.0; t * k];
    for i in 0..t { for p in 0..k { g[i * k + p] = (0..n).map(|j| y[i * n + j] * w[p * n + j]).sum(); } }
    g
}

impl Lora {
    /// Returns y and u = x·A (kept for backward). B starts at zero, so y = x·W at step 0.
    pub fn forward(&self, x: &[f32], t: usize) -> (Vec<f32>, Vec<f32>) {
        let mut y = mm(x, &self.w, t, self.k, self.n);
        let u = mm(x, &self.a, t, self.k, self.r);
        for (yv, d) in y.iter_mut().zip(mm(&u, &self.b, t, self.r, self.n)) { *yv += self.s * d; }
        (y, u)
    }
    /// Given dL/dy, returns (dL/dx, dL/dA, dL/dB). W gets no gradient and no optimizer state.
    pub fn backward(&self, x: &[f32], u: &[f32], dy: &[f32], t: usize) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
        let mut db = mm_tn(u, dy, t, self.r, self.n);
        db.iter_mut().for_each(|g| *g *= self.s);
        let mut du = mm_nt(dy, &self.b, t, self.n, self.r);
        du.iter_mut().for_each(|g| *g *= self.s);
        let da = mm_tn(x, &du, t, self.k, self.r);
        let mut dx = mm_nt(dy, &self.w, t, self.n, self.k);
        for (d, e) in dx.iter_mut().zip(mm_nt(&du, &self.a, t, self.r, self.k)) { *d += e; }
        (dx, da, db)
    }
    /// Fold the adapter into W for serving: W += s·A·B. Zero inference overhead afterwards.
    pub fn merge(&mut self) {
        let ab = mm(&self.a, &self.b, self.k, self.r, self.n);
        for (w, d) in self.w.iter_mut().zip(ab) { *w += self.s * d; }
        self.b.iter_mut().for_each(|b| *b = 0.0);
    }
}

/// SFT cross-entropy over response tokens only (mask = true); prompt tokens give no gradient.
pub fn masked_ce(logits: &[f32], vocab: usize, targets: &[usize], mask: &[bool]) -> f32 {
    let (mut sum, mut cnt) = (0.0f64, 0usize);
    for (i, row) in logits.chunks(vocab).enumerate().filter(|(i, _)| mask[*i]) {
        let m = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let lse = m as f64 + row.iter().map(|v| ((v - m) as f64).exp()).sum::<f64>().ln();
        sum += lse - row[targets[i]] as f64;
        cnt += 1;
    }
    (sum / cnt.max(1) as f64) as f32
}

/// DPO on one pair. Inputs are summed response log-probs under the policy and frozen reference.
/// Returns (loss, dL/dlogp_chosen, dL/dlogp_rejected).
pub fn dpo(pc: f32, pr: f32, rc: f32, rr: f32, beta: f32) -> (f32, f32, f32) {
    let z = beta * ((pc - rc) - (pr - rr));
    let loss = if z > 0.0 { (-z).exp().ln_1p() } else { -z + z.exp().ln_1p() }; // softplus(-z), stable
    let g = 1.0 / (1.0 + z.exp()); // sigmoid(-z)
    (loss, -beta * g, beta * g)
}

fn main() {
    let (t, k, n, r) = (3, 6, 5, 2);
    let f = |i: usize, m: f32| ((i * 7919 % 101) as f32 / 101.0 - 0.5) * m;
    let mut l = Lora { k, n, r, s: 16.0 / r as f32, w: (0..k * n).map(|i| f(i, 1.0)).collect(),
        a: (0..k * r).map(|i| f(i + 3, 0.2)).collect(), b: (0..r * n).map(|i| f(i + 9, 0.1)).collect() };
    let x: Vec<f32> = (0..t * k).map(|i| f(i + 17, 2.0)).collect();
    // Loss = sum(y): dL/dy = 1. Check dL/dA[0] against a finite difference.
    let (y, u) = l.forward(&x, t);
    let (_, da, _) = l.backward(&x, &u, &vec![1.0; t * n], t);
    l.a[0] += 1e-2;
    let num = (l.forward(&x, t).0.iter().sum::<f32>() - y.iter().sum::<f32>()) / 1e-2;
    l.a[0] -= 1e-2;
    assert!((num - da[0]).abs() < 1e-2 * num.abs().max(1.0));
    l.merge();
    assert!(l.forward(&x, t).0.iter().zip(&y).all(|(a, b)| (a - b).abs() < 1e-4));
    let (loss, gc, gr) = dpo(-10.0, -12.0, -10.0, -12.0, 0.1); // policy == reference: loss ln 2
    assert!((loss - std::f32::consts::LN_2).abs() < 1e-6 && gc < 0.0 && gr > 0.0);
    let ce = masked_ce(&[0.0, 0.0, 5.0, 0.0], 2, &[1, 0], &[false, true]);
    assert!((ce - (1.0 + (-5.0f32).exp()).ln()).abs() < 1e-6);
    println!("ok");
}

Efficiency

• Cost: about 16 B/param for full fine-tuning (112 GB for 7B), 2 B/param for LoRA on bf16, 0.5 B/param for QLoRA. Bound: GPU memory, then compute.
• Levers:
  • LoRA — 10,000x fewer trainable parameters and 3x less GPU memory than Adam on GPT-3 175B (LoRA paper) — many tasks
  • QLoRA 4-bit NF4 base — 65B model on one 48 GB GPU (QLoRA paper) — single-GPU runs
  • merge adapters into W — no inference overhead — single-task serving
  • precomputed DPO reference log-probs — saves a forward pass and the reference model — DPO
• Naive → standard → state of the art:
  • full fine-tuning, fp32 Adam · about 16 B/param
  • LoRA on a bf16 base · about 2 B/param
  • QLoRA · about 0.5 B/param

Pitfalls

• Chat template mismatch between training and inference (missing BOS/EOS, different role tokens): the model rambles or never stops.
• Loss on prompt tokens: the model learns to generate user turns.
• LoRA with B and A both random: ΔW ≠ 0 at start and the model is damaged before training. One factor must be zero.
• Forgetting α/r when merging, or merging into the 4-bit weights instead of a bf16 copy (re-quantizing loses the update).
• DPO on pairs generated by a different model than the SFT reference: off-distribution data, unstable training. Generate pairs with the SFT model when you can.
• DPO length exploitation: chosen answers longer on average → policy just gets longer. Balance lengths or use a length-normalized variant.
• Evaluating only the target task: regressions in general ability go unnoticed until users find them.

Implementations

Sources

• Hu et al., LoRA: Low-Rank Adaptation of Large Language Models, ICLR 2022. paper · code
• Dettmers et al., QLoRA: Efficient Finetuning of Quantized LLMs, NeurIPS 2023. paper · code
• Kalajdzievski, A Rank Stabilization Scaling Factor for Fine-Tuning with LoRA, 2023. paper
• Biderman et al., LoRA Learns Less and Forgets Less, TMLR 2024. paper
• Ouyang et al., Training language models to follow instructions with human feedback, NeurIPS 2022. paper
• Rafailov et al., Direct Preference Optimization: Your Language Model is Secretly a Reward Model, NeurIPS 2023. paper
• Azar et al., A General Theoretical Paradigm to Understand Learning from Human Preferences, 2023. paper
• Ethayarajh et al., KTO: Model Alignment as Prospect Theoretic Optimization, 2024. paper
• Meng, Xia, Chen, SimPO: Simple Preference Optimization with a Reference-Free Reward, 2024. paper
• Zhou et al., LIMA: Less Is More for Alignment, NeurIPS 2023. paper
• Tunstall et al., Zephyr: Direct Distillation of LM Alignment, 2023. paper
• Touvron et al., Llama 2: Open Foundation and Fine-Tuned Chat Models, 2023. paper
• Hugging Face, LoRA, diffusers training guide (text-to-image). docs