Training an LLM · Build an LLM From Scratch

The loss made quality measurable; now we minimize it. This unit builds train_model_simple — a compact but complete LLM pretraining loop — runs it for 10 epochs (full passes over the training data) on a tiny corpus, and reads the two curves it produces. The loop is standard deep learning; what’s instructive is watching a GPT go from gibberish to grammar to memorization in 85 steps.

The training loop, step by step

Two nested loops

The outer loop runs num_epochs = 10 full passes over the training data; the inner loop pulls (input, target) batches from ch2's data loader. model.train() enables dropout for learning.

1 / 7

One pass through train_model_simple. The inner cycle (steps 2–5) repeats for every batch; evaluation and sampling punctuate it.

</> train_model_simple and its helpers

def train_model_simple(model, train_loader, val_loader, optimizer, device,
                     num_epochs, eval_freq, eval_iter, start_context, tokenizer):
  train_losses, val_losses, track_tokens_seen = [], [], []
  tokens_seen, global_step = 0, -1

  for epoch in range(num_epochs):
      model.train()  #A
      for input_batch, target_batch in train_loader:
          optimizer.zero_grad()  #B
          loss = calc_loss_batch(input_batch, target_batch, model, device)
          loss.backward()   # Calculate loss gradients
          optimizer.step()  # Update model weights using loss gradients
          tokens_seen += input_batch.numel()
          global_step += 1

          if global_step % eval_freq == 0:  #C
              train_loss, val_loss = evaluate_model(
                  model, train_loader, val_loader, device, eval_iter)
              train_losses.append(train_loss)
              val_losses.append(val_loss)
              track_tokens_seen.append(tokens_seen)
              print(f"Ep {epoch+1} (Step {global_step:06d}): "
                    f"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")

      generate_and_print_sample(  #D
          model, tokenizer, device, start_context
      )
  return train_losses, val_losses, track_tokens_seen

def evaluate_model(model, train_loader, val_loader, device, eval_iter):
  model.eval()  #E
  with torch.no_grad():
      train_loss = calc_loss_loader(train_loader, model, device,
                                    num_batches=eval_iter)
      val_loss = calc_loss_loader(val_loader, model, device,
                                  num_batches=eval_iter)
  model.train()
  return train_loss, val_loss

def generate_and_print_sample(model, tokenizer, device, start_context):
  model.eval()
  context_size = model.pos_emb.weight.shape[0]  #F
  encoded = text_to_token_ids(start_context, tokenizer).to(device)
  with torch.no_grad():
      token_ids = generate_text_simple(
          model=model, idx=encoded,
          max_new_tokens=50, context_size=context_size
      )
  decoded_text = token_ids_to_text(token_ids, tokenizer)
  print(decoded_text.replace("\n", " "))
  model.train()

#A Training mode: dropout active. evaluate_model and the sampler switch to eval() and back — the mode dance matters.
#B The canonical four-beat rhythm: zero_grad → loss → backward → step.
#C Cheap periodic measurement (eval_iter batches only) — the source of the loss-curve checkpoints.
#D One qualitative sample per epoch: 50 greedy tokens from the fixed prompt.
#E Dropout off + no_grad during measurement, so evaluation is deterministic and cheap.
#F A neat trick: read the supported context size straight off the positional-embedding table's row count.

Why AdamW?

Adam optimizers are the popular choice for deep networks, but the loop uses AdamW — a variant with an improved weight decay treatment (decoupled from the adaptive gradient update). Penalizing large weights limits model complexity, and AdamW’s formulation regularizes more effectively — hence its ubiquity in LLM training. Here: lr=0.0004, weight_decay=0.1.

Running it: 10 epochs on a tiny corpus

</> The run — and what it printed

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)

num_epochs = 10
train_losses, val_losses, tokens_seen = train_model_simple(
  model, train_loader, val_loader, optimizer, device,
  num_epochs=num_epochs, eval_freq=5, eval_iter=5,
  start_context="Every effort moves you", tokenizer=tokenizer
)

# Ep 1 (Step 000000): Train loss 9.781, Val loss 9.933
# Ep 1 (Step 000005): Train loss 8.111, Val loss 8.339
# Every effort moves you,,,,,,,,,,,,.  #A
# Ep 2 (Step 000010): Train loss 6.661, Val loss 7.048
# Ep 2 (Step 000015): Train loss 5.961, Val loss 6.616
# Every effort moves you, and, and, and, and, and, and, ...  #B
# ... ...
# Ep 9 (Step 000075): Train loss 0.717, Val loss 6.293
# Ep 9 (Step 000080): Train loss 0.541, Val loss 6.393
# Every effort moves you?" "Yes--quite insensible to the irony.
# She wanted him vindicated--and by me!" ...  #C
# Ep 10 (Step 000085): Train loss 0.391, Val loss 6.452

#A After one epoch: the model has learned that commas are common. Progress, of a sort.
#B Epoch 2: high-frequency words on repeat — the loss is falling but the text is degenerate.
#C Epoch 9: grammatical, coherent English — because it is MEMORIZED verbatim from the training text ("The Verdict"). The val loss tells on it.

The two curves, from the printed checkpoints — hover to compare, and watch them part ways around epoch 2:

training lossvalidation loss

hover the plot to read the values at each checkpoint

Train vs validation loss at the printed checkpoints (the book's output elides epochs 3–8, so the long straight segment between steps 15 and 75 interpolates the gap). Both fall together at first; then training loss keeps sinking toward 0.39 while validation loss flattens near 6.4 — and even creeps back up.

Key idea — the model learned the data, not the language

The diverging curves are the signature of overfitting : with a very, very small training set iterated many times, the model minimizes training loss by memorizing — late-epoch samples are verbatim training-set passages. The validation loss (≈ 6.4 while training loss hits 0.39) is what exposes it. Real pretraining fights this with vastly more data; §5.3’s decoding strategies mitigate the symptom (repetitive, memorized generations) at sampling time.

</> Plotting the curves (the book's matplotlib version)

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator

def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):
  fig, ax1 = plt.subplots(figsize=(5, 3))
  ax1.plot(epochs_seen, train_losses, label="Training loss")
  ax1.plot(epochs_seen, val_losses, linestyle="-.", label="Validation loss")
  ax1.set_xlabel("Epochs")
  ax1.set_ylabel("Loss")
  ax1.legend(loc="upper right")
  ax1.xaxis.set_major_locator(MaxNLocator(integer=True))  #A

  ax2 = ax1.twiny()  #B
  ax2.plot(tokens_seen, train_losses, alpha=0)
  ax2.set_xlabel("Tokens seen")

  fig.tight_layout()
  plt.show()

epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))
plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)

#A Integer-only epoch ticks on the primary x-axis.
#B A twin x-axis maps the same curve to "tokens seen" — the units LLM training progress is usually quoted in.

📝 Check yourself: training the LLM

0 / 5

1.Put one training step in the right order:
2.Why does the training loop call optimizer.zero_grad() every iteration?
3.Why AdamW rather than plain Adam for LLM training?
4.By epoch 10, training loss is 0.391 but validation loss is 6.452. What happened?
5.What do eval_freq and eval_iter control in train_model_simple?