Part 2 of the slm-125m Story · A Beginner's Guide

The base model could talk.
It couldn't answer.

If you've read Part 1, you watched slm-125m learn to predict the next token over 2.04 billion words of case law, SEC filings, and web text. That model is fluent. Ask it a question, though, and it just keeps writing in the same register — it has no idea "question" and "answer" are different things. This is the story of teaching it that difference: a 10,000-pair synthetic QA dataset, a full fine-tune on a single small GPU, one real overfitting bug caught mid-run, and the exact math behind all of it.

0
grounded QA pairs, synthesized
0
best validation loss (from 3.13)
~$3
total SFT cost, both phases
Start with why fine-tuning is needed at all
Part 2 · Step 01

From Completion to Conversation

Pretraining teaches a model to continue text. Fine-tuning teaches it a new job: respond to a question with an answer.

Prompt the base model with "What is the capital of Delaware?" and, mechanically, it does exactly what it was trained to do — predict the most statistically likely continuation of that string. In legal and financial text, a sentence that looks like a question is very often followed by more of the same document, not an answer. So the base model might continue with another clause, another citation, another sentence in the same register — because nothing in its training ever taught it that a question deserves a direct, short, answering reply.

This is not a knowledge problem. The information the model needs is very often already somewhere in its 125.8M parameters, absorbed during pretraining. It's a behavior problem — the model has never been shown the shape question → answer as something to reproduce. Fine-tuning fixes that by training on exactly that shape, thousands of times, until it becomes the model's default move whenever it sees a question-shaped prompt.

i

What doesn't change: the vocabulary, the architecture, the knowledge and style already learned from 2.04B pretraining tokens. Fine-tuning here starts from the finished slm-125m-base checkpoint and nudges its existing weights — it does not reset anything or add new capacity.

Part 2 · Step 02

Grounded QA Distillation — Not RAFT

There are two very different ways to build a QA dataset. This project deliberately picked the harder, closed-book one — and it's worth understanding why, before looking at any numbers.

RAFT (retrieval-augmented fine-tuning) trains a model to answer using a passage handed to it at inference time, alongside distractor documents and explicit "I don't know" examples when the answer isn't in the provided context. The model never has to memorize facts — it just has to read.

This project instead does grounded QA distillation: a stronger "teacher" model (Gemini) reads a real passage from the pretraining corpus and writes a question-answer pair whose answer is strictly derivable from that passage — but at training and inference time, slm-125m-qa only ever sees the question, never the source passage. It has to have internalized the answer during pretraining and recall it now, closed-book.

!

Why this matters for what comes later: closed-book recall at 125.8M parameters has a real capacity ceiling — there just isn't room to memorize the long tail of facts in a 2-billion-token legal/financial corpus. Section 8 of this guide shows exactly what that ceiling looks like in real model output, and it's the reason RAFT is flagged as the more promising next experiment at the very end.

Part 2 · Step 03

Loss Masking — Learn to Answer, Not to Ask

Part 1 showed pretraining computing loss at every position, no exceptions. Fine-tuning is the first place in this project's story where that changes.

Each QA example is built as one token sequence with two roles, wrapped in the chat tokens reserved back when the tokenizer was trained: <|user|>question<|assistant|>answer<|eos|>. If loss were computed on every token here the way it was in pretraining, the model would spend part of its gradient budget learning to generate questions — a skill nobody wants from a QA model.

M
Loss masking — question ignored, answer trained
TOKENS
<user>
What is
the
holding?
<asst>
The court
held
for
plaintiff
.
<eos>
LABEL (loss computed here, or ignored)
-100
-100
-100
-100
-100
"held"
"for"
"plaintiff"
"."
"<eos>"
Prompt span — label = -100 (ignored)
Answer span — real label, loss computed

The mechanism, exactly as implemented:

sft_data.py — the actual masking logic
prompt_ids = [user_id] + question_ids + [assistant_id]
response_ids = answer_ids + [eos_id]
full = prompt_ids + response_ids
labels = [-100] * len(prompt_ids) + response_ids
# -100 is PyTorch's cross-entropy "ignore this position" index —
# gradients only flow from tokens the model should learn to generate

Verified, not assumed: real packed examples were decoded back to text after tokenization, confirming the non-masked span matches the answer text exactly, token for token.

Part 2 · Step 04

Full Fine-Tune, Not LoRA

You may have heard that fine-tuning a language model means LoRA adapters — small trainable matrices bolted onto a frozen base. That's a solution to a memory problem this model doesn't have.

LoRA exists because updating every parameter of a 7B–70B model, plus its gradients and optimizer state, needs far more memory than most GPUs have. Freeze the base weights, train only a small low-rank adapter, and the memory problem disappears — at the cost of some capacity to change the model's behavior.

At 125.8M parameters, that trade-off doesn't need to be made. Full fp32 weights, gradients, and AdamW optimizer state together are about 2GB — comfortably inside any modern GPU, let alone the ones this project used. So finetune_run updates every one of the 125.8M parameters directly, exactly like pretraining did, just starting from a trained checkpoint instead of random noise and running for far fewer tokens.

weights (fp32)
~503 MB
503 MB
gradients
~503 MB
503 MB
AdamW state
~1,006 MB
1,006 MB
T4 available
16 GB
16 GB
i

Even the cheapest commonly available GPU (a T4, chosen for this exact job — see Chapter 06) has 16GB of memory, 8× more than this model's full training footprint needs. LoRA would be solving a problem that simply doesn't exist here.

End of Part 2 · Beginning of Part 3

Now the specifics.
How the QA fine-tune was actually built.

Part 2 covered what fine-tuning does, in general. Part 3 is this project's real run — the real dataset, the real teacher model, a real API outage, a real overfitting bug, and the real numbers from the real training curve.

4,761 sampled passages 14,283 raw pairs → 10,000 kept 918 steps, 1× T4, 4.5 minutes → Part 3 starts here
Part 3 · Chapter 01 — Sourcing the Dataset

No new data.
Same corpus, new shape.

The QA pairs aren't sourced from anywhere new — they're synthesized from the same cleaned pretraining corpus, in the same domain proportions, so the model is being asked to demonstrate knowledge it should already have.

SEC
42%
2,000
case law
35%
1,666
fineweb-edu
23%
1,095

Chunk sampling

1,200-character passages (≈300 tokens) are sampled from /data/corpus — the already-cleaned, deduped Phase 2 output, not raw source data — proportional to the realized pretraining mix shown above.

Real logged output: sampled 4761 chunks (target 4761): {'case-law': 1666, 'sec': 2000, 'fineweb-edu': 1095} — the sampler hit its target almost exactly.

How many chunks are "enough"?

Working backwards from a 10,000-pair target, through an assumed 70% survival rate through filtering, at 3 pairs per chunk:

config.py — SFTConfig.num_chunks
num_chunks = target_pairs / filter_survival_rate / qa_per_chunk
          = 10_000 / 0.7 / 34_761 chunks

The 70% survival estimate was a planning number, deliberately conservative. Chapter 04 shows the real survival rate came in at 92.9% — comfortably above plan, which meant the pipeline over-produced and needed a final trim step, not a top-up.

i

The prompt's core instruction, from sft_data.build_prompt(): "Every answer must be derivable strictly from the passage — do not use outside knowledge, and do not invent facts, names, or numbers not present here." This single sentence is what makes the dataset "grounded" — everything downstream (the filter in Chapter 04, the failure mode in Chapter 08) traces back to how well the teacher model actually obeyed it.

Part 3 · Chapter 02 — Choosing (and Fixing) the Teacher

The model in the plan
didn't exist anymore.

The first real obstacle in this phase wasn't a training bug — it was discovering, mid-implementation, that the planned API had quietly changed underneath the project.

Problem 1: the model was deprecated

The original plan called for gemini-2.5-flash. It still appeared in the API's model listing — but every generation call against it returned a 404, "no longer available to new users." The replacement lineup is newer, pricier per token, and — critically — "thinks" by default, billing that reasoning as regular output tokens. A direct test against a one-word reply measured 82 thinking tokens for a single word of actual output.

Two fixes, both load-bearing

  1. Use gemini-3.1-flash-lite, not the full gemini-3.6-flash — extracting a grounded Q&A pair from a short passage is simple instruction-following, not a task that benefits from a larger reasoning model. Lite is ~5× cheaper for no measurable quality loss here.
  2. Explicitly disable thinking: thinking_config.thinking_budget = 0. Confirmed via a direct API test — with this set, thoughtsTokenCount doesn't appear in the response at all.

The cost math

Batch API tier
$0.125/1M in
$0.75/1M output
Estimated, 10K pairs
~$2.81
at Batch API pricing

This estimate assumed the 50%-discount Batch API would actually be used for the full run. Chapter 03 explains why that assumption didn't survive contact with reality — and why the real spend came in somewhat higher.

Part 3 · Chapter 03 — Generation at Scale

The API stalled.
For over an hour.

The cheaper path was tried first, in good faith, and abandoned only after it demonstrably stopped being trustworthy — a real engineering judgment call, not a shortcut.

What happened: Gemini's async Batch API — 50% cheaper, designed for exactly this kind of bulk job — worked perfectly on a 10-request smoke test. On the real 7,143-request batch, pendingRequestCount stuck at 100% and updateTime never advanced past createTime, for over an hour, despite real token usage appearing on the account. The API's own SLA is "up to 24 hours" with no incremental progress signal — there was no way to tell "working slowly" from "silently stuck" from the outside.

Decision: cancel it, and switch to real-time parallel calls — roughly 2× the token cost with no batch discount, but visible, verifiable, cancellable progress. This mirrors the exact "one worker per shard" philosophy Phase 1 of pretraining used for cleaning millions of documents.

1

sample_chunks()

Draws 4,761 passages from the pretraining corpus, proportional to the realized data mix.

2

generate_qa_shard() × 48

One Modal worker container per ~100-chunk shard, each calling Gemini sequentially with retries — 48 containers running in parallel.

3

merge_qa_shards()

Combines all 48 shard outputs into one raw QA file.

process_qa()

Validation, dedup, and chat-format conversion — Chapter 04.

Real logged result: generation complete: 4761 ok, 0 failed — a 100% success rate across 4,761 requests, and every single chunk yielded exactly 3 parsed pairs (14,283 = 4,761 × 3 exactly). The JSON-array prompt format was reliable enough in practice that the fallback markdown-fence parser in sft_data.parse_qa_response() never actually had to trigger.

Part 3 · Chapter 04 — Filtering, With Math

14,283 raw pairs.
10,000 survived.

Three checks, all free — no more API calls — turn the teacher's raw output into a clean, deduplicated training set.

The grounding check

This is the one that actually enforces the prompt's "don't invent facts" instruction. For every answer, count what fraction of its content words also appear in the source passage:

GROUNDING OVERLAP — sft_data.grounding_overlap()
passage: "...the Company's SEC filings report a 12% increase in net revenue..." · answer: "Net revenue increased 12% year over year."
"net"
✓ in passage
"revenue"
✓ in passage
"increased"
✓ in passage
"12%"
✓ in passage
overlap = 4/4 = 1.00 ≥ 0.3 threshold →KEPT

Only 18 of 14,283 pairs (0.13%) failed this check — the grounding instruction in the prompt worked almost perfectly on its own; the filter mainly exists to catch the rare exception, not to do the bulk of the work.

Exact and near-duplicate removal

Exact dedup normalizes question text (lowercase, punctuation stripped) and drops repeats. Near-dup removal embeds every surviving question with sentence-transformers/all-MiniLM-L6-v2 and greedily drops any question whose cosine similarity to an already-kept question is ≥ 0.92.

StageCountDropped
Parsed from teacher output14,283
Grounding check (≥30% content-word overlap)14,26518 (0.13%)
Exact-duplicate questions13,811454
Near-duplicate (embedding cosine ≥ 0.92)13,261550
Randomly trimmed to target10,0003,261
i

92.9% overall survival — well above the 70% planning estimate from Chapter 01, which is exactly why a final random-trim step was needed at all: generation produced more usable pairs than the target required. Final source mix after trimming: SEC 4,184 (41.8%) · case law 3,379 (33.8%) · fineweb-edu 2,437 (24.4%) — close to the intended 42/35/23.

Part 3 · Chapter 05–06 — Tokenizing and Training Setup

9,800 examples.
Zero truncated.

Tokenization here uses the loss-masking scheme from Part 2 for real, on real data — and the training config makes one more "smaller model, smaller GPU" call, following the same reasoning pretraining used.

Train / val split
9,800/200
98/2
Avg. real tokens
40.6
out of 256 budget
Max real tokens
106
longest example
Truncated
0
the "concise answer" prompt worked

Efficient batching — cropping to the batch, not the budget

Stored examples are padded to a fixed 256 tokens for uniform storage, but the real average is 40.6 — training against the full 256-width every batch would waste roughly 84% of every forward pass on pure <|pad|> tokens. sft_train.sft_batch_iterator crops each batch down to its own longest real example instead:

sft_train.py — dynamic batch cropping
real_len = max(1, int((ids != pad_id).sum(axis=1).max()))
ids = ids[:, :real_len]  # e.g. 256 → 64, a real ~4× compute reduction

Which GPU, and why (again)

The same reasoning Part 2 applied to pretraining (1× H100, not 8×) gets applied here in the opposite direction: this job is small enough that a faster GPU wouldn't even help.

Total compute~1.2M tokens over the whole run — a rounding error next to pretraining's 2.04B
Bottleneckfixed per-step overhead (Python loop, kernel launch), not raw FLOPs, at this batch/model size
GPU chosenT4 — confirmed empirically at ~0.2s/step; an A100 or H100 would not run faster here, only cost more per hour
Precisionfp16 autocast + torch.amp.GradScaler — T4 is Turing-generation hardware with no native bf16 tensor cores, unlike pretraining's H100; fp16 needs gradient scaling to avoid underflow, which is why this loop has a scaler and the pretraining loop didn't
S
SFTTrainConfig
every fine-tuning hyperparameter
epochs3configured — turned out to be more than needed, see Chapter 07
batch_size32examples per optimizer step
lr / min_lr3e-5 / 3e-6cosine schedule, 20× lower peak than pretraining's 6e-4 — fine-tuning nudges an already-good model, it doesn't relearn from scratch
warmup_fraction0.03= 3% of total steps spent ramping up to peak LR
weight_decay0.01lighter than pretraining's 0.1 — less regularization pressure needed for a short run
beta1 / beta20.9 / 0.95same AdamW moment decay as pretraining
eval_every_steps100validation loss checkpoint cadence — the single most important setting in this run, see Chapter 07
gpuT4Turing — fp16 + GradScaler, not bf16
Part 3 · Chapter 07 — The Overfitting Bug

Train loss said success.
Val loss said otherwise.

This is the most important lesson in this entire fine-tuning phase — and it very nearly shipped as a silent mistake.

The first full run (918 steps, 3 configured epochs) looked, by training loss alone, like a clean success: loss fell smoothly and dramatically, from 3.01 down toward ~1.0. If that were the only number being watched, the run would have been called done and the final-step checkpoint pushed straight to Hugging Face.

Validation loss told a different story. On the first attempt, it bottomed out around step 400–500 — under two of the three configured epochs — and then began rising for the rest of training, even as train loss kept falling. This is the textbook signature of overfitting: the model stopped learning generalizable question-answering behavior and started memorizing the specific training examples. The qualitative generations confirmed it — samples from step 900 read noticeably more hallucinatory than samples from earlier in training.

The part that made this expensive, not just informative: the training loop's checkpoint-saving logic overwrote the same path on every save. Once training moved past the true optimum, there was no checkpoint left on disk that captured it — the actual best model was gone, and the run had to be repeated from scratch to recover it.

the fix — track the best state_dict in memory, write once
best_val_loss = float("inf")
best_state = None

for step in range(total_steps):
    # ...forward, backward, optimizer step...
    if step % eval_every_steps == 0:
        vloss = evaluate(model)
        if vloss < best_val_loss:
            best_val_loss = vloss
            # deep copy to CPU — doesn't touch disk, doesn't slow training
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}

# only write to disk ONCE, after training finishes entirely
model.load_state_dict(best_state)
model.save_pretrained(config.SFT_BEST_CKPT_DIR)

Re-run, with best-checkpoint tracking — real logged data

best: step 300 · val 2.1224

Determinism confirmed: the re-run's training loss trajectory was bit-for-bit identical to the first attempt's — same seed, same data order, same everything except the checkpointing logic. That match is what makes it possible to say with confidence that the fix changed nothing except which checkpoint got saved, not how training itself behaved.

Two smaller bugs, caught the same week

ModuleNotFoundError: No module named 'pydantic' — crashed the SFT tokenization step, which shares a file with the inference endpoint code but doesn't share its Modal image
guarded try/except import fallback
"pad token is same as eos token" warning from generate() — looked like a tokenizer bug, wasn't; the two IDs were genuinely distinct
pass an explicit attention_mask
Part 3 · Chapter 08 — Results, and What the Model Actually Learned

2.12 loss.
Fluent — and confidently wrong.

The number is a clear win. What the number hides is more interesting, and it's the honest headline of this whole phase.

Best checkpoint
step 300
of 918 — under 1 of 3 epochs
Best val loss
2.1224
perplexity ≈ 8.3
Pre-finetune baseline
3.1299
perplexity ≈ 22.9
Training time
4.5 min
1× T4, full run

Real qualitative samples

Three real generations from the deployed slm-125m-qa model, taken from the live demo:

QuestionExpected (grounded)Model generated
What are the primary factors that cause Polycystic ovary syndrome?A combination of genetic and environmental factors."...a lack of adequate nutrition, poor diet and weight gain, and an increased risk for developing breast cancer" — fluent register, wrong facts.
What two requirements are generally needed for the formation of a contract?A bargain involving mutual assent to the exchange, and consideration."They must be signed by both parties and must contain provisions that require each party to sign" — plausible legal boilerplate, not the actual doctrine.
Why was a rule to show cause issued for the appeal in No. 562, 2013?To direct the appellee to show cause why the appeal should not be dismissed as interlocutory."The court granted the motion because it determined that the appellant failed to provide any evidence of actual prejudice or misconduct" — coherent legal reasoning shape, wrong specific facts.
!

The pattern, every time: the model reliably reproduces the register and structure of a correct answer in each domain — formal legal phrasing, SEC-style business description, encyclopedic health-topic tone — while frequently inventing the specific facts inside that structure. This is exactly the closed-book capacity ceiling flagged back in Chapter "Grounded QA Distillation": 125.8M parameters is not enough to memorize the long tail of specific facts across a 2-billion-token corpus, so the model falls back on plausible pattern completion whenever the true answer isn't reliably stored in its weights.

This reframes what "2.1224 validation loss" actually measures: the model got very good at the shape of an answer — short, on-topic, grammatically clean, correctly terminated — which is genuinely what cross-entropy loss on held-out data rewards. It is not, on its own, a claim about factual accuracy. Loss went down because predicting plausible-sounding tokens in the right register got easier; it did not go down because the model became a more reliable source of truth.

Part 3 · Chapter 09 — Shipping It

Two models,
one demo.

The QA fine-tune deploys through the exact same three-service pattern Part 2 used for the base model — nothing new had to be invented to ship it.

Hugging Face
slm-125m-qa, public

The step-300 best checkpoint, pushed with an auto-generated model card documenting the exact training config and metrics above.

Modal
/generate_qa endpoint

A second route on the same ASGI app that already served the base model, with its own in-memory model cache alongside the base model's.

Vercel
Live QA testing box

The demo site's browser calls /generate_qa directly, right next to the original base-model completion demo, so both models are testable side by side.

Summary

Three ideas.
One model that answers.

Loss masking
Learn to answer, not to ask

Unlike pretraining's every-token loss, SFT sets -100 over the question span — gradients only ever flow from the answer the model should learn to generate.

Watch validation, not train
Best checkpoint: step 300 of 918

Train loss fell the whole way to ~1.0. Validation loss bottomed under one epoch — the gap between those two curves is the entire overfitting story.

Register, not facts
The honest headline result

Fluent, correctly structured, frequently wrong on specifics — the expected shape of closed-book distillation at 125M parameters, not a bug to hide.

Now go test it yourself. Both checkpoints and the live demo are one click away.