LoRA Fine-Tuning Engine

A from-scratch implementation of LoRA (Low-Rank Adaptation) fine-tuning for a small decoder-only transformer language model, built directly on PyTorch's tensor/autograd primitives — no PEFT, no transformers.Trainer. The transformer itself (embeddings, positional encoding, multi-head causal self-attention, feed-forward blocks, layer norm, residual connections) is also hand-written rather than assembled from nn.TransformerEncoderLayer / nn.MultiheadAttention.

This is the training counterpart to LLM-inference-server, a from-scratch Llama-family inference/serving engine (KV caching, batching, scheduling) in the same account. Together they cover both ends of the model lifecycle this account demonstrates: training-time adaptation here, serving-time efficiency there.

What LoRA does

Instead of updating a full weight matrix W during fine-tuning, LoRA freezes W and learns a low-rank correction:

W_effective = W_frozen + (alpha / r) * B @ A

where A is (r, in_features), B is (out_features, r), and r << min(in_features, out_features). Only A and B are trained. B is initialized to all zeros, so at the start of training the adapted layer is numerically identical to the frozen base layer — a standard LoRA init trick, verified directly in tests/test_lora.py with an exact hand-computed numeric check (not just a "doesn't crash" test).

apply_lora() freezes every parameter in the model, then walks its named children and replaces any Linear submodule whose attribute name matches a target list (q_proj, k_proj, v_proj, out_proj by default) with a LoRALinear wrapping the original weights. A dedicated test (test_frozen_base_vs_trainable_adapter_requires_grad) asserts the base weight/bias have requires_grad=False and only lora_A/lora_B are trainable — a real bug class in hand-rolled LoRA implementations is accidentally leaving the base trainable.

Honest scope

This sandbox has no GPU (torch.cuda.is_available() is False) and a modest, shared CPU. Real production-scale LLM fine-tuning needs a GPU and would not complete in reasonable time here, so the model and task are deliberately scoped down to make CPU experimentation honest and reproducible:

• Model: a tiny character-level transformer — 3 layers, d_model=128, 4 heads, ~615K total parameters. This is not a claim about fine-tuning a production language model; it's a small, from-scratch transformer sized for CPU iteration.
• Data: there is no bundled real-text corpus. Both corpora used below are generated from scratch by src/lora_ft/data.py from simple sentence templates with a fixed seed — a "generic" everyday-topic corpus standing in for pretraining data, and a "finance" business-jargon corpus (EBITDA, shareholders, quarterly forecasts, dollar amounts) standing in for a narrow domain-adaptation fine-tuning target. This is disclosed, not hidden, and lets the comparison below be run with zero network access.
• Tokenization: character-level over a fixed ~90-character vocabulary, not a subword/BPE tokenizer — the right scope for a model this small.

Real measured results

Full run: python -m lora_ft.cli (defaults: 300,000 chars/corpus, 8 pretrain epochs, 8 fine-tune epochs, LoRA rank 8, alpha 16). Measured on this machine, single CPU process, no GPU:

Pretraining (generic corpus, 8 epochs): final train loss 1.7072, final validation loss 1.4246, wall time 208.2s.

Then three conditions were fine-tuned/evaluated against the finance corpus, starting from the identical pretrained checkpoint:

condition	val_loss (generic)	val_loss (finance)	trainable params	total params	train time
no fine-tune	1.4246	3.4090	0	614,784	0.0s
full fine-tune	3.0833	0.7960	614,784 (100%)	614,784	207.3s
LoRA fine-tune (r=8)	2.1970	2.3288	24,576 (3.85%)	639,360	182.0s

What this honestly shows

• No fine-tune is the worst on the target (finance) distribution, as expected — the model has never seen that vocabulary/statistics.
• Full fine-tuning adapts hardest to the finance corpus (0.796 val loss, the best of the three) but pays for it with severe catastrophic forgetting: generic-corpus validation loss more than doubles, from 1.42 to 3.08 — the model is measurably overwriting what it learned during pretraining.
• LoRA fine-tuning trains only 3.85% of the parameters and forgets far less (generic val loss 2.20, vs. full-ft's 3.08) — but at rank 8 it also doesn't recover as much target-task quality as full fine-tuning (2.33 vs. 0.80 on finance). Training time is modestly lower than full fine-tuning (182.0s vs. 207.3s) — a smaller effect than the parameter-count reduction alone would suggest, because the forward pass still runs the full frozen network; only the backward pass and optimizer state shrink.

This is the real, textbook LoRA trade-off at a small rank: parameter efficiency and reduced catastrophic forgetting, in exchange for less aggressive adaptation to a very different target distribution — not a "LoRA wins on every axis" result. Numbers are not tuned to look more favorable; they are reported as measured. A CI smoke-test run (--quick, 20K chars/corpus, tiny model, 2 epochs) reproduces the same qualitative ordering in ~6 seconds and its report is uploaded as a CI artifact on every run.

Why this matters across industries

Parameter-efficient fine-tuning (LoRA and its relatives) is now a standard technique anywhere an organization wants to adapt a foundation model to its own data without the compute, storage, or forgetting cost of a full fine-tune:

• Tech / AI: the default approach for customizing open-weight LLMs per product, customer, or task without maintaining a full model checkpoint per variant.
• Banking / finance: adapting a general model to house terminology, disclosures, and risk/compliance language (the synthetic "finance corpus" above is a toy stand-in for exactly this use case) without retraining from scratch.
• Consulting: client engagements routinely need a model customized to one client's documents/jargon on a tight compute budget — LoRA adapters are cheap to store and swap per client.
• Energy: domain-jargon-heavy technical documentation (well logs, ICS/SCADA manuals, safety procedures) benefits from the same lightweight adaptation pattern used here.

Running it

pip install --index-url https://download.pytorch.org/whl/cpu torch
pip install -e ".[dev]"

pytest tests/ -v                                    # 13 tests, ~2s

python -m lora_ft.cli --quick --output-dir runs/demo # ~6s smoke test
python -m lora_ft.cli --output-dir runs/full         # ~10 min, full comparison above

Or via Docker:

docker build -t lora-ft .
docker run --rm lora-ft

Project layout

src/lora_ft/
  model.py   # TinyTransformerLM: embeddings, positional encoding, causal self-attention,
             # feed-forward, layer norm, residual connections — hand-written
  lora.py    # LoRALinear + apply_lora(): the from-scratch LoRA layer and injection logic
  data.py    # synthetic corpus generators + character tokenizer (honest disclosure above)
  train.py   # block construction, batching, train/eval loops
  cli.py     # pretrain -> {no-ft, full-ft, lora-ft} -> compare, end to end
tests/       # 13 tests: attention masking, LoRA math (hand-computed), frozen/trainable
             # parameter checks, model shape/config checks

License

MIT — see LICENSE.