Skip to content

Fine-tuning (SFT & LoRA)

Stage 2 turns the raw language model into one that follows instructions. Both approaches here use the same data and the same next-token objective — they differ only in which parameters they update.

Watch the SFT walkthrough

Script: sft.py  ·  Loads: ckpt_base.pt  ·  Produces: ckpt_sft.pt  ·  Run: python sft.py [minutes] (default 5)

Updates all ~14.6M weights on (instruction → story) pairs from TinyStories-Instruct.

Script: lora.py  ·  Loads: ckpt_base.pt  ·  Produces: ckpt_lora.pt (adapters only)  ·  Run: python lora.py [minutes] (default 5)

Freezes the base model and trains a tiny low-rank correction (~0.5% of the parameters) on the attention projections.

Supervised fine-tuning

The data: instruction → story

Each TinyStories-Instruct document carries Features, Words, Summary, and a Story: marker. The text up to and including Story: is the prompt; everything after is the completion. An <|endoftext|> is appended to teach the model where to stop.

Features: Dialogue
Words: dog, happy, park
Summary: A dog makes a new friend at the park.
Story:  ← prompt ends here; the story that follows is the completion

Completion masking

The objective is identical to pretraining, but the loss must only count the story tokens. Prompt and padding positions are set to target -1, which F.cross_entropy(..., ignore_index=-1) drops from the mean. The model learns to write a story conditioned on the instruction, not to parrot instructions back.

One-position subtlety

Masking stops one position short of the prompt length: the last prompt token is exactly the position that must predict the first story token, so it has to contribute to the loss.

Hyperparameters

Knob Value Note
Learning rate 2e-4 lower than pretraining — nudge behavior, don't overwrite the base
Batch size 48 right-padded; padding is masked out
Betas / weight decay (0.9, 0.95) / 0.1 same grouping as pretraining

Validation loss is completion-masked too, so it is not comparable to pretraining loss — track its trend: val rising while train falls means the model is memorizing the SFT set.

LoRA: parameter-efficient fine-tuning

Instead of moving all the weights, freeze the base model and learn a low-rank update on each targeted projection:

\[ h = W x + \frac{\alpha}{r}\, B (A x) \qquad A \in \mathbb{R}^{r \times \text{in}},\;\; B \in \mathbb{R}^{\text{out} \times r},\;\; r \ll \text{in} \]
class LoRALinear(nn.Module):
    def forward(self, x):
        return self.base(x) + F.linear(F.linear(x, self.lora_A), self.lora_B) * self.scale
  • B starts at zero, so B·A = 0 and at step 0 the adapted model is exactly the base model — training only moves the small A/B matrices.
  • Adapters go on the q_proj and v_proj of every block — the classic LoRA-paper choice; adding k/o/MLP buys little at this scale.
  • With r = 8, α = 16, only ~0.5% of parameters are trainable, and the saved adapter checkpoint is ~300 KB versus ~60 MB for full weights.
  • The learning rate is a much hotter 1e-3 — the adapters start at zero and are tiny, so they tolerate (and need) a larger step than full fine-tuning.

The training loop is otherwise identical to sft.py; only the set of parameters the optimizer holds changes.

Reload and merge

model, tok = load_lora_model()   # base checkpoint + adapter checkpoint
merge_lora(model)                # fold W += scale · B·A back into each Linear

merge_lora folds every adapter into its base weight, restoring a plain nn.Linear with zero inference overhead and identical outputs — at the cost of no longer being separable from the base.

python sft.py 5     # full fine-tuning  → ckpt_sft.pt
python lora.py 5    # LoRA adapters     → ckpt_lora.pt

📄 Source: sft.py · lora.py