gemma-2-2b-legal-merged
google/gemma-2-2b-it fine-tuned with 4-bit QLoRA for grounded question answering
over US case law and SEC filings. Merged FP16 weights, usable directly with transformers.
Read this first
This model answers only from a passage you supply. It is not a closed-book legal model — ask it a question without source text and it will invent an answer.
Results
Held-out validation, 1,003 examples, greedy decoding. Every row uses the same metrics, so they are directly comparable:
| Model | Token F1 | Exact | F1 ≥ 0.5 | F1 < 0.2 | Numeric fidelity |
|---|---|---|---|---|---|
| 125M trained from scratch, full FT | 0.590 | 11.4% | 60.1% | 10.5% | 98.1% |
| gemma-2-2b-it zero-shot | 0.720 | 12.4% | 82.2% | 2.9% | 97.3% |
| This model (QLoRA 4-bit) | 0.856 | 36.0% | 94.9% | 0.4% | 99.8% |
On what the fine-tune actually bought. Token F1 is scored against reference answers written by the teacher model, so some of that gain is the model learning the teacher's phrasing rather than getting more answers right — the exact-match jump in particular. The honest signal is numeric fidelity, which checks every number the model emits against the passage rather than the reference and so cannot be gamed by style: it rose 97.3% → 99.8%, and clearly-wrong answers fell 2.9% → 0.4%. That part is real.
Training
| Method | QLoRA, 4-bit NF4 + double quantization, bf16 compute |
| LoRA | r=16, alpha=32, all 7 linear projections |
| Trainable | 20,766,720 / 2,614,341,888 = 0.79% |
| Data | 8,000 teacher-distilled grounded QA pairs |
| Schedule | 1 epoch, 249 steps, effective batch 32, lr 2e-04 |
| Hardware | 1× H100, 39.9 min, $2.62 |
| Peak GPU | 46.9 GB |
What 4-bit actually saved
Measured at identical batch size, 4-bit peaked at 25.3 GB against bf16's
31.0 GB — an 18% saving, not the 8× the "4-bit" label suggests. The reason
is that bitsandbytes quantizes nn.Linear but not nn.Embedding, and Gemma 2's
256,128-token vocabulary makes the embedding table 22.6% of all
parameters, which stays in bf16:
quantized (NF4) 1.01 GB
NOT quantized 2.36 GB <- embeddings, norms
Two different 4-bit formats
- Training used bitsandbytes NF4. It is CUDA-only and cannot run on CPU.
- Serving uses GGUF Q4_K_M via llama.cpp, a separate quantization of the merged weights. That is what the GGUF repo contains.
An adapter trained against a 4-bit base does not merge losslessly into FP16 — it partly learned to correct quantization error that is absent after merging.
Gotchas if you fine-tune Gemma 2 yourself
- Use
attn_implementation="eager". FlashAttention-2 does not implement Gemma 2's logit soft-capping and fails silently — you get a plausible loss curve on a model that is not Gemma 2. - Watch the logits tensor, not the weights. With a 256,128 vocab, cross-entropy at batch 32 × seq 1024 is 33.6 GB of logits alone and OOMs an 80 GB card before any weight loads. Quantization does not help; a smaller micro-batch does.
- Gemma 2's chat template has no
systemrole and raises if given one. Fold the system prompt into the first user turn — in training and inference alike.
Limitations
- Requires the source passage; no useful closed-book recall.
- ~0% of answers are still clearly wrong.
- Training data is synthetic, from a single teacher, and inherits its biases.
- English only. Not legal or financial advice.