Unit 3.2 · about 3 h
305 / 324Feeding the model: batches, initialization and first steps
Silent data bugs, from ordered batches to bad initialization, waste every GPU hour after them. Two tiny controls on CPU catch them in minutes before any real run.
- implement
- measure
You can now
- Build a batch sampler that draws uniform random windows with correct input-target pairing.
- Apply the assignment initialization recipe per module type and explain each choice.
- Run overfit-one-batch and shuffled-label controls and read the verdict from each.
Covers
- Sequence batching, shuffling and data-loader state
- Initialization and stable first-step controls
- Tiny-set overfit and shuffled-label controls
Study, in this order
3 steps
- CS336 Assignment 1 handout: data loadingStanford CS336 Assignment 1 - training and decoding contract
Section 5.1, Data Loader (pp. 34 to 35):
(batch_size, context_length)input and target pairs, targets offset by one, random start indices and memory-mapped loading withnp.memmap - CS336 Assignment 1 handout: initialization recipeStanford CS336 Assignment 1 - training and decoding contract
Section 3.3.1, Parameter Initialization (p. 17): truncated normal for linear weights and embeddings with
torch.nn.init.trunc_normal_, ones forRMSNorm - CS336 Assignment 1 handout: debugging controlsStanford CS336 Assignment 1 - training and decoding contract
Section 7.2.3, Tips and tricks for debugging model architectures (p. 40): overfit a single minibatch to near-zero loss, inspect tensor shapes in a debugger, monitor activation, weight and gradient norms
Focus on
- Each row draws its own random start in \([0, n-m)\), so a start never runs past the end of the token array.
- Inputs and targets share the same device as the model. The handout CPU tip names
cpuandmpsexplicitly. - Linear, embedding and norm layers get different initial scales because they play different roles in the first forward pass.
- The two controls answer different questions. Driving one tiny batch to near-zero loss proves the model and optimizer can fit at all. Training on shuffled targets should leave validation loss near the unigram baseline; doing much better means the answer leaks in, for example through a broken causal mask or a missing shift.
Reproduce closed-book
get_batch from scratch: memory-mapped load, uniform random starts, and \((B, m)\) input and target tensors with targets offset by one.Traps
- Sampling a start index up to \(n\) instead of \(n-m\). The window reads past the array end or wraps into unrelated tokens.
- Initializing every weight to zero or one constant. Symmetric weights get identical gradients and never break symmetry.
- Judging initialization from the first-step loss alone. A healthy init shows in stable norms and in the overfit control, not in one number.
Practice
CS336 Assignment 1: `data_loading`
Problem (data_loading) in Section 5.1; connect your sampler to run_get_batch in tests/adapters.py and pass test_get_batch in tests/test_data.py
The test draws 1000 batches from arange(100) and checks shapes, the offset-by-one relation and uniform start coverage. Implement it in your own repository.
Self-test
-
The token array has \(n=100\) entries and context length \(m=7\). Which start indices are valid?
0 through 92, which is 93 values. The window needs \(m\) tokens plus one target, so the last start is \(n-m-1=92\). -
You train on shuffled targets and validation loss falls almost as fast as with the real targets. What does that tell you?
The model is reading the answer through a leak, such as a causal mask that lets position \(i\) see position \(i+1\), or inputs and targets taken from the same unshifted window. Shuffled targets carry no signal beyond token frequencies, so validation loss should stall near the unigram baseline. -
Overfit-one-batch stalls at high loss. What does that verdict mean?
The pipeline or model is broken, not the data. Check tensor shapes, the label shift, the learning rate and whether any parameter is frozen or detached. -
Why does the recipe give
RMSNormones but small truncated weights to linear and embedding layers?Norms start as identity so they pass activations through, while small weights keep first-step activations and gradients finite. Large or constant init risks exploding or vanishing gradients before step one.
Completion criterion
Stop when
Your sampler passes test_get_batch, your loop drives one tiny batch to near-zero loss, and a run on shuffled targets leaves validation loss near the unigram baseline.
Leads to
- 3.3 AdamW, schedules and the global token batch Hand the batches to
AdamWand count the global token batch. - 3.4 Keeping training numerically healthy Watch the first steps for exploding norms and overflow.
Comes after
- 1.7 Filtering, mixtures and deterministic shardsFeed batches from a versioned token corpus.
- 3.1 The next-token objective: shifts, masks and perplexityFeed correctly shifted and masked batches from a real data loader.
Resources for this unit
- Stanford CS336 Assignment 1 - training and decoding contract 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.
Also in this section
- LLMs from Scratch companion - training and generation 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.
- Modern LLM Notebook - configuration, loss and generation 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 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 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 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.