Chapter 7

Training a Tiny Transformer

Train a small decoder model with inspectable data windows, targets, initialization, optimization, gradients, validation, and checkpoints. Begin by overfitting one tiny batch so correctness is tested before scale is increased.

Training changes a model's parameters so that its next-token probabilities fit the examples in a corpus. A falling loss is useful evidence only after the corpus, target shift, causal mask, parameter set, and update order are known to be correct.

This chapter trains the 368-parameter decoder from Chapter 6. Its architecture does not change:

QuantityValue
vocabulary size8
context length4
model width4
decoder blocks2
attention heads2
MLP width8
tied learned entries368

The fixed forward contract remains

I:(B,T)⟶Z:(B,T,8)⟶L∈R.I:(B,T)\longrightarrow Z:(B,T,8)\longrightarrow \mathcal L\in\mathbb R.

What changes is the experiment around it. Token sequences become aligned input and target batches. PyTorch records the forward computation, calculates a gradient for each registered parameter, and AdamW uses those gradients and its stored state to update the model.

The corpus is deliberately small:

<bos> A B C <eos>
<bos> A B D <eos>
<bos> X Y X <eos>
<bos> Y X Y <eos>

This is not a small sample of natural language. It is a controlled grammar for testing whether the implementation learns intended conditional probabilities. Some contexts have more than one valid next token, so the best possible population loss is not zero.

The executable experiment runs on CPU in a local Python environment. The browser-Python editor used for lighter examples does not include PyTorch. The chapter explains each framework operation and supplies exact checks rather than asking you to trust a framework call.

After this chapter

  • Construct correct context-target batches for causal language modeling.
  • Use tiny-batch overfitting, loss, and gradient checks to validate training.
  • Record and compare reproducible training configurations and samples.

Lessons

  1. 01
    Turn a Token Corpus into Training Examples

    Construct causal language-model examples from token sequences, boundary markers, shifted targets, context windows, and separate validation data.

    1 exercise
  2. 02
    Build Aligned Input and Target Batches

    Trace causal training batch shapes, integer token IDs, target alignment, assertions, and loss reshaping without mixing sequences.

    1 exercise
  3. 03
    Make the Frozen Decoder Trainable in PyTorch

    Implement the frozen tiny decoder with PyTorch modules, tied embeddings, causal attention, pre-norm blocks, and explicit parameter and shape checks.

    1 exercise
  4. 04
    Initialization Sets the First Signal Scale

    Audit Transformer initialization seeds, parameter distributions, activation and logit scale, initial cross-entropy, gradients, and controlled scale ablations.

    1 exercise
  5. 05
    Overfit One Sequence before Training the Corpus

    Overfit one sequence to test Transformer gradients and updates, understand what the test establishes, and diagnose failure before adding data.

    1 exercise
  6. 06
    Run the Forward, Backward, and Update Cycle

    Trace one PyTorch Transformer training step through zero_grad, forward loss, backward, gradient checks, clipping, optimizer state, and logging.

    1 exercise
  7. 07
    Inspect Gradient Norms and Learning-Rate Sensitivity

    Measure Transformer gradient norms, clipping effects, learning-rate sensitivity, loss outcomes, and controlled optimizer ablations.

    1 exercise
  8. 08
    Read Training Loss, Validation Loss, and Diagnostic Predictions

    Interpret Transformer train and validation loss, conditional entropy, irreducible cross-entropy, diagnostic probabilities, and narrow generalization limits.

    1 exercise
  9. 09
    Save Enough State to Repeat or Resume the Run

    Build and test Transformer checkpoints for inference, reproducible experiments, and exact optimizer continuation with model and RNG state.

    1 exercise
  10. 10
    Audit a Training Run Systematically

    Debug tiny Transformer training with a dependency-ordered audit, failure matrix, interventions, metric checks, and controlled ablations.

    1 exercise

Review and practice

  1. •
    Review

    Review tiny Transformer data, shapes, parameters, training cycle, entropy floor, diagnostics, reproducibility, and debugging order.

  2. •
    Exercises

    Solve tiny Transformer training exercises on data, shapes, loss, gradients, optimization, validation, checkpoints, and debugging.

Chapter progress