Section 3 of 5 · LLM Foundations

Causal Training, Checkpointing & Generation

A correct forward pass is only the beginning. This section fixes target alignment, valid-token weighting, full training state and sampling semantics so that a tiny model can train, resume and generate reproducibly.

7 units 25 topics 0/7

Units

7 in this section

3.1 The next-token objective: shifts, masks and perplexity
derive · implement
4 h
3.2 Feeding the model: batches, initialization and first steps
implement · measure
3 h
3.3 AdamW, schedules and the global token batch
derive · implement
4 h
3.4 Keeping training numerically healthy
explain · measure
3 h
3.5 Checkpoints that really resume
implement · measure
4 h
3.6 Decoding policies: from greedy to nucleus sampling
derive · implement
4 h
3.7 Stopping rules and reproducible generation
implement
3 h

Before moving on

Checkpoint

Overfit one tiny batch, prove target and mask alignment, train and resume a compact decoder with a stated equivalence level, then reproduce fixed-seed greedy, temperature, top-k and top-p outputs while explaining each distributional change.

Deliberately out of scope

Leave for later

Large-scale pretraining and distributed checkpointing
SFT, LoRA, DPO, RLHF and RLVR
Production decoding servers and fleet scheduling

Resources for this section

Shared by every unit above

LLMs from Scratch companion - training and generation
Sebastian Raschkaimplementation

primary

Use the ch05 training and generation code for the small pretraining loop. Add full checkpoint state, explicit loss masks and resume-equivalence tests that the base example does not supply.
Stanford CS336 Assignment 1 - training and decoding contract
Stanford Universityassessment

required

Use cross_entropy, adamw, adamw_accounting, learning_rate_schedule, gradient_clipping, data_loading, checkpointing, training_together and decoding specifications plus public tests. Debug on CPU and implement independently in your own repository.
Modern LLM Notebook - configuration, loss and generation
Walking Labsimplementation

support

Use English notebooks 09-model-config, 10-training-loss and 20-generation to visualize configuration ownership, shifted loss and decoding policies. Change edge cases and cross-check probability calculations against your own implementation.
PyTorch saving, loading and reproducibility references
PyTorchreference

support

Use the state-dict and reproducibility notes to enumerate all state needed for the supported resume guarantee. Test model, optimizer, schedule, scaler, data order and RNG state across an interruption.
Hugging Face LLM Course - library translation
Hugging Faceapplication

optional

Only after the blank-file loop passes, map tokenization, model configuration, data collation, generation configuration and stopping behavior into Transformers APIs. Keep the from-scratch tests as the authority.
PyTorch Automatic Mixed Precision recipe
PyTorchreference

support

Use the autocast, GradScaler, unscale-before-clip, scaler checkpointing and inf-NaN troubleshooting sections to add mixed precision to the blank-file loop. Keep the CS336 clipping contract as the authority for the clipping rule.

Apply the section

Mini-project · Train, resume and sample the tiny decoder (10 to 14 h)

Turn the decoder into a trained model you can trust: a loss you can derive by hand, a loop that resumes exactly, and decoding policies whose distributions you can check. Train on the packed split from section 1 on a CPU or a free GPU notebook.

Steps

Implement cross-entropy, gradient clipping, AdamW, the warmup-cosine schedule, get_batch and checkpointing, and pass the CS336 tests in test_nn_utils.py, test_optimizer.py, test_data.py and test_serialization.py.
Run two controls before any real run: drive one tiny batch to near-zero loss, and confirm that a run on shuffled targets leaves validation loss near the unigram baseline.
Train on your section 1 split with a fixed seed, logging loss, gradient norm, learning rate and tokens seen, with a validation curve from the held-out split.
Save complete state, including model, optimizer, step, RNG and data cursor, and show that an interrupted run matches an uninterrupted one to a stated level.
Implement greedy, temperature, top-\(k\) and top-\(p\) sampling with EOS and length stops, and reproduce fixed-seed outputs exactly.
Train on a synthetic Markov chain with a known entropy rate: validation loss should approach that rate, and a loss clearly below it means the labels leak. On CPU, a resumed run should match an uninterrupted one exactly.

Deliverables

The training script and config, with the checkpoint format documented.
The training log and validation curve for your run, plus the two control runs.
A resume-equivalence result: the steps compared, the tolerance and why it is the right one.
Sample outputs for each decoding policy from one prompt and seed, with the config that reproduces them.

Checks

external-tests · CS336 tests for cross-entropy, clipping, `AdamW`, the schedule, data loading and checkpointing
rubric · Self-grade with the rubric below

Rubric

Criterion	Pass	Strong
Objective	Your cross-entropy passes the public test, including the overflow case.	You derive one masked loss by hand and match your code, and the shuffled-target control behaves as you predicted.
Training	Validation loss falls on your split, and on the Markov chain it approaches the entropy rate without going clearly below it.	The log shows learning rate, gradient norm and tokens seen, and you explain one anomaly from it.
Resume	An interrupted run continues without a loss jump.	Losses after resume match the uninterrupted run exactly on CPU, and your stated GPU tolerance is justified.
Decoding	Each policy runs and stops on `EOS` or the length cap.	Fixed-seed outputs reproduce exactly, and you show by hand how temperature followed by top-\(p\) changes the kept set at one step.

Feeds into: Capstone artifact 2: the trained checkpoint and its decoding policies. Section 4 evaluates this checkpoint, and section 5 must reproduce its greedy outputs with a cache.