LoRA fine-tuning and its modern successors, implemented from scratch — adapt a frozen transformer by training a few percent of it, with the base provably untouched (and optionally int8-quantized to cut its memory ~4x).
No PyTorch, no PEFT library. LoRA, DoRA (Weight-Decomposed LoRA, 2024), and
QLoRA-style quantization are all hand-written — forward and backward — and
gradient-checked. The transformer being fine-tuned lives in forge.backbone —
the same from-scratch architecture as the companion Glassbox project, bundled
here so Forge is completely self-contained: clone it and demo it on its own,
nothing else to install. Adapter save/load/hot-swap and honest perplexity
evaluation are included.
Full fine-tuning updates every weight — expensive in memory and compute. LoRA (Hu et al. 2021) made adaptation cheap enough for one GPU and is now the default way to specialize open models. Forge proves I understand it from the inside: I derived the low-rank update and its gradients by hand, and I can show the base model is genuinely frozen while a tiny adapter does the adapting.
It's the third of a set (each stands alone — no shared imports, so any one can be demoed on its own):
- Glassbox — I understand what's inside a model (built one from scratch).
- Forge — I can adapt one efficiently (this project; it bundles its own
copy of that transformer as
forge.backbone). - Crucible — I can test what I build (honest eval infrastructure).
Pretrain a base model on one style, freeze it, attach LoRA, fine-tune on a different style — training ~5% of the parameters — and watch generation shift while the base weights stay byte-for-byte identical:
base model (trained on style A): 'the mat the cat sat on the mat and the dog r'
LoRA-tuned (adapted to style B): 'the mind to be on the that is the question...'
training 5120 adapter params out of 104064 (4.9% of the model)
base weights unchanged by fine-tuning? True
python -m forge.demo # LoRA fine-tune: pretrain -> adapt -> before/after
python -m forge.demo --method dora # same, but fine-tune with DoRA
python -m forge.compare_dora # LoRA vs DoRA on a controlled task (2024 technique)
python -m forge.web # dashboard: pretrain, fine-tune live, compare + perplexity
python -m pytest tests/ -q # the test suite (or: python tests/test_forge.py)Freeze a weight W (in,out); learn a low-rank correction A (in,r), B (r,out):
y = x @ W (frozen base, unchanged)
+ (x @ A @ B) * (alpha / r) (trainable low-rank update)
Points I can defend at a whiteboard:
- Why low-rank works: the change needed to adapt to a task has low intrinsic rank — it lives in a small subspace — even though the weights are full-rank. So a rank-8 adapter recovers most of a full fine-tune.
- Why B starts at zero: then
A@B = 0, so the adapted model is the base model at step 0 — fine-tuning begins as a no-op and moves away smoothly instead of jolting the model with a random perturbation. - The alpha/r scaling decouples the learning rate from the rank, so you can
change
rwithout re-tuning. - The backward pass: since
Wis frozen we never needdW— onlydA,dB, and thedxthat carries gradients to earlier layers. All three are gradient-checked to 1e-8. - Merging:
merged_weight() = W + scaling*A@Bfolds the adapter back into the base, so inference has zero overhead once merged (a real deployment step).
LoRAHead subclasses the backbone's attention head, reuses the frozen Q/K/V
weights,
and adds adapters to the query and value projections (the LoRA paper's common
recipe). apply_lora(model) swaps every head; lora_params_and_grads(model)
collects only the adapters, so the optimizer can never touch the base —
guaranteed, and asserted in the tests. count_params reports the honest
trainable/total ratio (every base parameter counted, not just attention).
Pretrain (streamed loss) → attach LoRA and fine-tune (streamed loss + live
parameter count) → generate from base vs adapted side by side, with a check that
the base stayed frozen. The "base" side works by temporarily zeroing the adapter
B matrices, which recovers the exact pre-LoRA output (verified).
LoRA changes a weight's magnitude and direction together. DoRA decomposes them:
it applies the low-rank update to the direction (unit-normalized columns),
re-normalizes, then rescales by a separately trained magnitude vector:
W' = m · (W + s·A@B) / ||W + s·A@B||_col. That decoupled magnitude knob is why
DoRA closes much of the gap to full fine-tuning. It's a first-class fine-tuning
option here: apply_lora(model, method="dora") swaps in weight-decomposed
adapters on attention Q/V, trained exactly like LoRA. On a controlled fit
(compare_dora) DoRA reaches ~half the error of LoRA at the same rank, for a
handful of extra parameters. Forward and the column-normalization backward are
hand-written and gradient-checked to 1e-8, both standalone and inside attention.
The frozen base is the memory cost of fine-tuning a big model. QLoRA stores it in low precision and keeps only the small adapters in full precision. We implement symmetric per-column int8 quantization (4x smaller; real QLoRA uses 4-bit NF4 for 8x, same principle). The key correctness point: the base is frozen, so we never backprop into it — we only need its value, which we dequantize on the fly. The adapters train identically over a quantized base (gradient-checked), and can even compensate for quantization error. Round-trip error is ~0.6%.
The real production pattern: keep ONE frozen base in memory, swap tiny adapters
(here ~6 KB) per task/customer. save_adapters/load_adapters write and restore
just the adapter arrays; a saved adapter round-trips exactly, and swapping it in
restores the adapted behavior precisely.
perplexity reports exp(mean cross-entropy) on held-out text (lower = better),
in the Crucible spirit of measuring generalization rather than asserting it. In
the demo the adapted model's perplexity on the target style drops dramatically
below the base's.
- Tiny, character-level, CPU — it demonstrates the mechanism, not scale. On a real model the trainable fraction would be ~1%; here it's a few percent because attention Q/V is a larger slice of a small model. The code reports the real measured number, not a marketing figure.
- Adapting only attention Q/V (not the MLPs) limits how far a small model can move; the adaptation is visible but not dramatic at this size.
- Fine-tuning "style B" reuses the shared character vocabulary so token ids line up cleanly between the two styles.
- LoRA/DoRA on the MLP projections too (more adaptation capacity — currently the adapters target attention Q/V).
- True 4-bit NF4 quantization with bit-packing (this uses int8 for legibility).
- A larger held-out eval with per-style perplexity tables.