Basic training loop with cross-entropy loss · CS336 Language Modeling From Scratch

Deliverableartifact

A train.py script that trains the GPT model on tokenized text for several steps, reports train/val loss, and includes optional gradient clipping.

Goals

Train your model for a few iterations on a small corpus and show loss decreasing over steps.
Demonstrate gradient clipping by printing gradient norms before and after clipping.
Implement a validation loop that computes average loss over a held-out split.
Identify and handle common training issues like NaNs by adding checks and printing diagnostics.

Resources

Overview: what you’re building

You’re going to build a small train.py script that:

Takes tokenized text and chops it into batches of sequences
Feeds them to a GPT-style model
Uses cross-entropy loss between predicted next-token logits and shifted targets
Runs a standard PyTorch training loop
Optionally applies gradient clipping
Logs train/val loss every few steps

Think of this as the “spine” of training: once you understand this, you can plug in bigger models, datasets, schedulers, etc.

We’ll assume:

You already have:
- model = GPTModel(...) (any nn.Module that takes (batch, seq) token IDs and returns (batch, seq, vocab_size) logits)
- A tokenized 1D torch.LongTensor array of token IDs: data
You’re comfortable running a Python script from the command line.

1 / 6