Instruction Finetuning Loop Execution, Convergence, and Capstone Stress Testing
How is the complete instruction finetuning training loop structured in PyTorch with AdamW, learning rate schedules, and masked loss?
How this lesson unfolds.
This is the capstone chapter of Build LLM From Scratch. Here, we bring together every single component built throughout the curriculum into an end-to-end instruction training loop.
How is the complete instruction finetuning training loop structured in PyTorch with AdamW, learning rate schedules, and masked loss?
Imagine practicing responses for a quiz bowl. Your coach reads the question, but you are only scored on the answers that come out of your mouth. In the training loop, the prompt is visible for context, but the error calculation is masked out so you are only graded on the answers you generate.
- $$\mathcal{L} = -\frac{1}{\sum_{i, t} \mathbb{I}[y_{i, t} \neq -100]} \sum_{i, t \mid y_{i, t} \neq -100} \ln P(y_{i, t} \mid x_{i, \le t})$$
- $$The instruction training loop feeds dynamically collated batches through the forward model, calculates cross-entropy loss ignoring prompt tokens via ignore_index=-100, clips gradients to prevent instability, and updates parameters with weight-decayed AdamW.$$
- Learning goalAssemble the complete instruction training loop in PyTorch with masked cross-entropy, gradient clipping, and AdamW.
- Helpful foundationsMasked target collation · Cross-entropy loss calculation
- Common trapApplying weight decay to LayerNorm weights and bias vectors; doing so degrades normalization stability and harms convergence.
- Next connectionVerify initial baseline loss before triggering updates in Screen 2.
More context & mechanics
The end-to-end instruction training loop executes four tightly choreographed phases per batch:
- Dynamic Forward Pass: Input token IDs X ∈ R^(B × T) pass through the causal transformer to produce unnormalized logit tensor Z ∈ R^(B × T × V).
- Masked Target Alignment: Targets Y ∈ R^(B × T) have prompt tokens and padding set to -100:
L = CrossEntropyLoss(Z[:, :-1, :], Y[:, 1:], ignore\_index=-100)
- Backward Gradient Computation: ∇_θ L is computed via PyTorch autograd.
- Gradient Clipping & Optimizer Step:
g_(clipped) = clip\_grad\_norm\_(∇_θ L, max\_norm=1.0)
θ_(t+1) = θ_t - η_t ( (m̂_t) / (√(v̂_t) + ε) + λ θ_t )
Where λ = 0.01 weight decay is applied only to weight matrices (excluding 1D LayerNorm scales and biases).
A batch of B=2 sequences with sequence length T=4:
- Masked target tensor Y = [[-100, -100, 45, 88], [-100, 12, 19, -100]].
- Unmasked tokens: Token at index 2 & 3 in seq 1; Token at index 1 & 2 in seq 2 arrow 4 active loss tokens out of 8 total tokens.
- Cross-entropy losses for the 4 active tokens: [1.20, 0.80, 1.50, 0.90].
- Batch Loss:
L = (1.20 + 0.80 + 1.50 + 0.90) / (4) = (4.40) / (4) = 1.10
The -100 tokens contribute exactly zero gradient, shielding the pre-trained prompt representations from destructive updates.
Standard core training loop pattern powering HuggingFace TRL Trainer, PyTorch Lightning, and raw PyTorch research scripts.
Loshchilov & Hutter introduced decoupled weight decay (AdamW) in 2017, solving standard Adam's weight decay decay decay decay decay flaws.
Mask the prompts with minus hundred; clip the gradients at one: AdamW guides the step, clean answers are won.
What is the expected initial cross-entropy loss before finetuning begins, and how do we verify loss masking integrity?
Before launching a rocket, engineers perform a static countdown test to ensure all sensors report nominal readings. Checking the initial loss before step 1 tells you instantly if your tokenizer, target masks, or data loader have fatal bugs before you burn GPU hours.
- $$\mathcal{L}_{\text{random}} = -\ln\left(\frac{1}{V}\right) = \ln(V) \approx \ln(50257) \approx 10.825$$
- Learning goalPerform pre-flight loss verification to detect uninitialized weights, label leakage, and masking failures.
- Helpful foundationsInstruction training loop · Natural logarithm of vocabulary size
- Common trapSkipping initial loss verification; catching a corrupted target mask at step 0 saves hours of wasted training runs.
- Next connectionObserve loss convergence curves across training epochs in Screen 3.
More context & mechanics
Prior to launching gradient updates, two quantitative sanity checks must be passed:
- Pretrained Baseline Loss Check:
For a pre-trained model like GPT-2 (124M) evaluated on instruction responses, initial cross-entropy loss is typically in the range:
L_(init) ∈ [2.5, 3.8]
If L ≈ -ln(1/V) = ln(50257) ≈ 10.82, the pre-trained weights were not loaded properly (model is acting like random initialization).
If L < 1.0, data leakage has occurred (targets leaked into inputs).
- Masking Integrity Audit:
Check that the count of -100 tokens in target tensor Y exactly equals the prompt token count plus padding token count:
N_(masked) = N_(prompt) + N_(pad)
If N_(masked) = 0, the loss function is mistakenly training on the prompts.
We load GPT-2 on a sample Alpaca prompt: 'Instruction: Write a haiku about rain.'
- Vocabulary size V = 50,257.
- Expected random loss: ln(50257) ≈ 10.825.
- We compute forward pass loss on the prompt with targets properly masked:
L_(observed) = 3.142
- Validation Audit Verdict:
- 3.142 ≪ 10.825 arrow Pretrained weights successfully active.
- 3.142 > 1.0 arrow No trivial label leakage.
- Number of -100 tokens in target = 18, matching prompt length (18 tokens).
Sanity check passed with 100\% confidence; proceed to training loop.
Standard automated pre-flight assertion in production deep learning pipelines (PyTorch Lightning `on_train_start`).
Sanity checking initial loss against ln(V) was popularized by Andrej Karpathy's 'A Recipe for Training Neural Networks' (2019).
Ten is random, three is trained, one is leaking: check before you train, save GPU pain.
How do training and validation loss evolve during instruction finetuning, and how do we detect catastrophic forgetting?
Learning a new language accent is like instruction tuning. At first, you speak awkwardly. Then your accent improves rapidly. But if you practice exclusively with one small script for too long, you memorize those exact lines and forget how to speak freely! Watching validation loss tells you when to stop before you become a parrot.
- $$\text{Epoch}^* = \arg\min_{e} \mathcal{L}_{\text{val}}(e)$$
- Learning goalInterpret instruction training loss trajectories, identify the early stopping inflection point, and mitigate catastrophic forgetting.
- Helpful foundationsPre-finetuning loss audit · Validation dataset splitting
- Common trapTraining until training loss reaches zero; zero loss in generative LLMs indicates severe rote memorization and degraded generalization.
- Next connectionPersist optimal model states with atomic checkpointing in Screen 4.
More context & mechanics
Instruction training dynamics differ markedly from pretraining:
- Phase 1: Format Adaptation (Steps 1–200): Loss drops precipitously from ≈ 3.2 to ≈ 1.5 as the model learns delimiter tokens (`### Response:`) and stylistic cadence.
- Phase 2: Task Refinement (Steps 200–1000): Gradual descent from ≈ 1.5 to ≈ 0.8 as reasoning syntax is refined.
- Phase 3: Overfitting Danger Zone (Steps 1000+): Training loss continues downward to < 0.4, but validation loss curves upward.
Catastrophic Forgetting:
Excessive finetuning destroys pre-trained linguistic associations. Regularly evaluate on general language benchmarks (e.g., Wikitext perplexity) alongside instruction loss to verify base capabilities remain intact.
A 5-epoch run on 1,000 instruction pairs with evaluation every epoch:
- Epoch 0 (Baseline): Train Loss: 3.25 | Val Loss: 3.28
- Epoch 1: Train Loss: 1.82 | Val Loss: 1.86 (Massive format gain)
- Epoch 2: Train Loss: 1.15 | Val Loss: 1.21 (Optimal alignment point)
- Epoch 3: Train Loss: 0.78 | Val Loss: 1.19 (Global Val Minimum)
- Epoch 4: Train Loss: 0.48 | Val Loss: 1.34 (Overfitting begins)
- Epoch 5: Train Loss: 0.29 | Val Loss: 1.58 (Severe memorization)
Action: The early stopping protocol halts training at Epoch 3, saving checkpoint `checkpoint_epoch₃.pt`.
Standard early stopping and checkpoint selection protocol in all deep learning production platforms.
Catastrophic forgetting in neural networks was first formalized by McCloskey & Cohen (1989) and remains a central challenge in modern alignment.
Format learns fast, memory rots slow: stop at the validation dip, never chase zero loss.
How do model generation samples qualitatively transform across training epochs, and how do we persist complete model checkpoints atomically?
Watching a model learn across epochs is like watching a child learn to write. At Epoch 0, it rambles and repeats random words. By Epoch 1, it writes complete sentences. By Epoch 2, it follows the exact format requested! Saving full checkpoints preserves this progress so you can reload any stage.
- $$\text{Checkpoint} = \{ \Theta_M, \mathbf{m}_{t}, \mathbf{v}_{t}, \eta_t, \text{step} \}$$
- Learning goalImplement complete atomic checkpoint saving and recovery, and track qualitative sample generation evolution.
- Helpful foundationsInstruction training dynamics · PyTorch state dictionary serialization
- Common trapSaving only model weights during training; without optimizer and scheduler states, resuming a run causes severe learning rate and momentum discontinuities.
- Next connectionStress-test the finalized instruction model against out-of-distribution prompts in Screen 5.
More context & mechanics
Atomic checkpoint persistence ensures reproducibility and crash recovery:
```python
checkpoint =
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'val_loss': val_loss
torch.save(checkpoint, f"checkpoint_epoch_(epoch).pt")
```
Qualitative Evolution of Sample Output Across Epochs:
- Prompt: `### Instruction: Give three tips for healthy sleep.`
- Epoch 0 (Base Model): 'Give three tips for healthy sleep and wake up early and sleep is good for health and healthy sleep is...' (Rambling, unhelpful continuation).
- Epoch 1: 'Here are tips for sleep:
- Go to bed early.
- Do not drink coffee.
- Sleep in dark room.' (Format acquired, blunt syntax).
- Epoch 2 (Optimal): 'Here are three actionable tips for healthy sleep:
- Maintain a consistent sleep schedule even on weekends.
- Avoid blue light from screens 1 hour before bed.
- Keep your bedroom cool (around 65°F/18°C) and dark.' (Polished, structured, highly helpful).
A complete checkpoint file for GPT-2 (124M):
- Model weights: 124.4 × 10⁶ × 4 bytes ≈ 497.6 MB
- Optimizer states (AdamW momentum + variance): 124.4 × 10⁶ × 8 bytes ≈ 995.2 MB
- Metadata (epoch, step, scheduler state): ≈ 4 KB
- Total Checkpoint Size: ≈ 1.49 GB
To save storage, production pipelines store:
- Full checkpoint (1.49 GB) only for the latest step (crash recovery).
- Weights-only export (`model.state_dict()`, 498 MB) for evaluated checkpoints.
Standard checkpointing practice across PyTorch Lightning, Deepspeed, and Hugging Face Accelerate.
Standardized through PyTorch's serialization specification (`torch.save` via pickle/zipfile format).
Save weights for inference, save optimizer for resumption: inspect samples with human eyes to see true learning.
How do we stress-test a finished instruction-finetuned LLM across diverse domains to verify robust generalization?
Graduating from driving school doesn't just mean driving on the test track. You have to drive in rain, on highways, at night, and in heavy traffic. Multi-domain stress testing takes your finetuned model out of its cozy training sandbox onto the open road.
- $$\text{Generalization} = \frac{1}{|\mathcal{D}_{\text{OOD}}|} \sum_{i=1}^{|\mathcal{D}_{\text{OOD}}|} \mathbb{I}[\text{Pass}(\mathcal{M}(x_i), y_i)]$$
- Learning goalDesign and execute a multi-domain stress test suite evaluating creative, coding, logic, and constraint compliance capabilities.
- Helpful foundationsInstruction training loop · Sample evolution and checkpoints
- Common trapAssuming in-domain validation performance guarantees real-world generalization; small models degrade sharply on negative constraints and multi-step logic.
- Next connectionGrand Finale of the Build LLM From Scratch Curriculum: from raw characters to autonomous instruction-following intelligence!
More context & mechanics
Stress testing evaluates zero-shot transfer across five unseen domains:
- Domain A — Creative Synthesis: Metaphorical explanations and creative storytelling.
- Domain B — Code & Structured Syntax: Python function generation with docstrings and return types.
- Domain C — Multi-Step Logic: Simple arithmetic and deduction prompts.
- Domain D — Negative Constraints: Instructions with explicit prohibitions ('Explain photosysthesis without using the letter e').
- Domain E — Adversarial Prompt Injection: Attempts to bypass instruction delimiters using nested pseudo-tags.
Evaluation Metrics:
- Constraint Compliance Rate: Percentage of hard negative constraints successfully honored.
- Instruction Completion Rate: Percentage of prompts receiving a fully formed, non-truncated answer.
- Hallucination Frequency: Manual and automated factual veracity audit.
We stress-test our finalized GPT-2 Instruction model on 50 out-of-distribution test prompts (10 per domain):
- Domain A (Creative): 9/10 pass (90\%)
- Domain B (Code): 7/10 pass (70\%) — small 124M model makes minor syntax slips in complex algorithms.
- Domain C (Logic): 6/10 pass (60\%) — limited parameter capacity constrains multi-step arithmetic.
- Domain D (Negative Constraints): 5/10 pass (50\%) — common limitation in small models without DPO/RLHF.
- Domain E (Adversarial Robustness): 8/10 pass (80\%) — prompt templates resist shallow delimiter hijacking.
- Overall Multi-Domain Generalization:
Generalization Score = (9 + 7 + 6 + 5 + 8) / (50) = 35/50 = 70.0\%
The model demonstrates genuine general-purpose instruction following, completing the entire journey from raw bytes to a functional instruction-tuned language model!
Red-teaming and safety acceptance gates prior to publishing open-source model weights or deploying customer-facing AI features.
Multi-domain stress testing was standardized by the BIG-bench collaboration (Srivastava et al., 2022) with 200+ consortium tasks.
Test in the rain, test in the snow: a model that only passes training knows nothing at all.