Lab 01 — LoRA / QLoRA Delta + Knowledge Distillation

Phase: 05 — PEFT: LoRA, QLoRA & Knowledge Distillation Difficulty: ⭐⭐⭐☆☆ (the linear algebra is small; the invariants are the lesson) Time: 3–4 hours

Adapting a frontier model used to mean copying all of its weights and paying for gradients and optimizer state on every one of them. PEFT broke that: freeze the giant base, learn a tiny correction. This lab builds the three mechanisms that make adaptation and serving affordable — LoRA (the low-rank delta ΔW=(α/r)·A·B, with the all-important zero-initialized B so the adapted model starts as the pretrained one, and merge() for zero inference overhead), QLoRA-style blockwise quantization (a 4-bit frozen base with a per-block absmax scale and a fixed symmetric grid, with a bounded reconstruction error), and knowledge distillation (training a small student on a big teacher's softened logits — the stable softmax_T, the T²·KL soft term, and the blended loss).

What you build

  • matmul / mat_add / transpose — the three matrix helpers everything sits on; each raises ValueError on a shape mismatch.
  • LoRALinear(W_base, r, alpha, seed) — a frozen base (d×k), a Gaussian-init A (d×r), and a zero-init B (r×k). .forward(x) = x·W_base + (α/r)·(x·A)·B; .delta() = (α/r)·A·B; .merge() = W_base + ΔW. At init the delta is exactly the zero matrix.
  • lora_param_count(d,k,r)=r·(d+k) and full_param_count(d,k)=d·k — the "0.4% trainable" headline number, exactly.
  • quantize_blockwise / dequantize_blockwise — NF4-style 4-bit storage: a per-block absmax scale + a fixed symmetric codebook; reconstruction error is bounded by the grid spacing and shrinks as you add levels.
  • softmax_T / kl_divergence / cross_entropy / distillation_loss — the distillation kit: temperature softmax (max-subtracted), stable KL (log-space, 0·log0 guarded), and the Hinton loss T²·KL(teacher‖student) + α·CE(student,target).

Key concepts

ConceptWhat to understand
ΔW = (α/r)·A·Badapt the update, not W; A·B is rank-r, so it trains r·(d+k) ≪ d·k params
B is zero-initso ΔW = 0 at step 0 ⇒ the adapted model is the pretrained model — the soul of LoRA
α/r scalingα sets the delta's magnitude independent of rank; doubling α doubles ΔW linearly
merge()fold the delta into the base ⇒ zero extra inference latency; the served model is one matrix
blockwise absmaxa tiny per-block scale tames outliers; a fixed grid stores the values in 4 bits
finer grid = less errormore levels ⇒ smaller quantization step ⇒ a strictly smaller error bound
softened logitsdividing by T>1 raises the non-top probabilities — the teacher's "dark knowledge"
the T² factorrescales the soft-term gradient (shrunk by ~1/T²) back onto the hard term's footing

Files

FilePurpose
lab.pyskeleton with # TODO markers — your implementation
solution.pycomplete reference; python solution.py runs a worked example
test_lab.pythe proof — run it red, make it green
requirements.txtpytest only (pure stdlib otherwise)

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example

Success criteria

  • All tests pass against your implementation (and LAB_MODULE=solution).
  • You can explain why test_lora_delta_is_zero_at_init (and ..._forward_equals_base) is the soul test — B zero-init means training starts at the pretrained model.
  • You can explain why x · merge() equals forward(x) and why that gives LoRA zero inference overhead once deployed.
  • You can state the LoRA trainable fraction for a 4096×4096 projection at r=8 (65,536 / 16,777,216 ≈ 0.39%) and why that is the entire economic case for PEFT.
  • You can explain why test_quant_finer_grid_has_smaller_error is monotonic and what the per-block absmax scale buys you over a single tensor-wide scale.
  • You can explain why distillation_loss(t, t, …) == α·CE (the KL term vanishes) and why higher T increases entropy (test_higher_temperature_softens).

How this maps to the real stack

The miniatureThe production mechanismWhere to verify it
LoRALinear + merge()peft's LoraConfig / get_peft_model; model.merge_and_unload() folds adapters for inferenceHF peft; the adapter's adapter_config.json (r, lora_alpha, target_modules)
A Gaussian, B zeroexactly PEFT's init: lora_A Kaiming/Gaussian, lora_B zeros, so the model is unchanged at step 0peft.tuners.lora source; inspect lora_B weights right after wrapping
lora_param_countthe "trainable params: 0.4%" line print_trainable_parameters() reportsmodel.print_trainable_parameters()
quantize_blockwise (+absmax)bitsandbytes NF4: blockwise 4-bit with a per-block absmax scale, plus double quantization of the scalesbitsandbytes Linear4bit; QLoRA's bnb_4bit_quant_type="nf4"
softmax_T / kl_divergencethe distillation objective in TinyBERT / DistilBERT / transformers-based KD trainersa KD training loop; torch.nn.KLDivLoss(reduction="batchmean") on log-softmax
distillation_loss (T²·KL + α·CE)Hinton's response distillation; the T² rescale and the soft/hard blend are standardHinton et al. (2015); DistilBERT's loss

Limits of the miniature (be honest in the interview): there is no autograd here — you set A/B by hand and verify the forward algebra, not a training loop (P03 is the autograd phase); our grid is a uniform symmetric codebook, while real NF4 uses 16 non-uniform levels fit to a normal distribution (and double-quantizes the scales); distillation here is response-based on one example, not feature/relation distillation over a dataset; and merge() ignores the small numeric drift real fp16 merging incurs.

Extensions (build these on real hardware)

  • Add a tiny SGD loop: a target weight W*, an MSE loss on forward(x) vs x·W*, and hand-rolled gradients for A, B; watch the rank-r delta fit the low-rank part.
  • Replace the uniform grid with the real NF4 16-level codebook (quantiles of a unit normal) and re-run test_quant_finer_grid_has_smaller_error against it.
  • Add double quantization: quantize the per-block scales themselves and measure the extra bytes saved (QLoRA's ~0.4 bits/param trick).
  • Add target_modules selection: count params if you adapt only attention q,v vs also the MLP, and reproduce the LoRA paper's "q,v is enough" result on param budget.
  • Add feature distillation: an extra MSE term between a student and teacher hidden layer, and show it helps when the soft-logit signal alone is weak.

Interview / resume

  • Talking points: "Why is LoRA's B zero-initialized, and what breaks if it isn't?" "How does QLoRA fine-tune a 65B model on one 48GB GPU?" "Walk me through the distillation loss and the role of T and T²." "When do you merge adapters, and what does it cost?"
  • Resume bullet: Implemented LoRA/QLoRA from scratch — the low-rank delta ΔW=(α/r)·A·B with zero-initialized B and a mergeable adapter, NF4-style blockwise 4-bit quantization with a bounded reconstruction error, and a temperature-scaled knowledge-distillation loss (T²·KL + α·CE) — with a deterministic, dependency-free test suite proving each invariant.