Module 2 · Building the Model
Chapter 10 of 17 Run it in your browser

Chapter 10: The Full GPT Architecture¶

Name: Teach Your Computer to Write: Build and Train an LLM from Zero
Author: Truong (Jack) Luu

Where we are in the pipeline Figure 10.1: We connect everything to output a prediction.

We have built all the core components of a Transformer model in previous chapters, so we can now assemble them (Figure 10.1). In this chapter you will:

Build the full GPT model from end to end.
See how data flows through the entire pipeline.
Understand the final raw scores, called logits.

Words to Know

LM Head: The model's last step, which turns each token's internal numbers into one score for every character in the vocabulary.
Logits: The raw scores the model outputs for each possible next token.

Theory¶

The Final Assembly¶

The complete GPT model is surprisingly simple once you have the pieces.

The pipeline flows from embeddings through blocks to the LM head Figure 10.2: The full forward pass of the GPT model.

It flows like a production assembly line (Figure 10.2):

Turn text into a list of numbers (Token IDs).
Look up the meaning and position vectors (Embeddings).
Process them through a stack of Transformer blocks.
Tidy up the numbers one last time (Final LayerNorm).
Map the internal 128 numbers back to the 65 possible characters (LM Head).

Tracing the Flow¶

To trace a concrete example, suppose we feed the model the word "ROMEO" (5 tokens).

Step 1: Input The input is [30, 27, 25, 17, 27]. These are the specific token IDs for the characters "ROMEO", where 'R' is 30, 'O' is 27, 'M' is 25, and 'E' is 17 according to our vocabulary.

Step 2: Embeddings The model converts these 5 IDs into 5 rich vectors (128 numbers each), combining both meaning and position.

Step 3: Transformer Blocks The data passes through 4 blocks. After 4 blocks, the token O (at the end) has looked at all the previous letters. Its 128 numbers now encode something like "I am the last letter of a name from a famous play."

Step 4 & 5: Final Output A final LayerNorm stabilizes the numbers, and then the LM Head (Language Model Head) takes those 128 numbers and turns them into 65 numbers. The number is 65 because our Shakespeare dataset has exactly 65 unique characters.

The internal 128 numbers are projected to 65 vocabulary scores Figure 10.3: The LM Head projects internal data back to our vocabulary.

The model outputs 65 scores for every token in the sequence (Figure 10.3), and each position independently predicts the character that comes next. The highest score at the final position is the model's best guess for what comes after "ROMEO".

What Are Logits?¶

The 65 scores the LM Head outputs are called logits.

Logit is a technical term for "raw score" (Figure 10.4). These numbers can be negative, zero, or very large, and they do not add up to 1, so they are not probabilities yet. In Chapter 11, we will convert these raw scores into proper percentages to generate text. For now, it is enough to know that a higher logit means the model is more confident in that character.

Logit scores for the token O show highest confidence for N and I Figure 10.4: Logits are the raw scores for the next character.

On the big picture map in Figure 10.1, we have now built everything from Text to Next-Token Scores. The model is fully assembled, but right now it only outputs random guesses, so we need to train it.

Code¶

src/ch09_gpt_model.py (excerpt)
class GPT(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.cfg = cfg

        # Look up table for token vectors
        self.token_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        # Look up table for position vectors
        self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)

        # A sequence of transformer blocks
        self.blocks = nn.Sequential(*[
            TransformerBlock() for _ in range(cfg.n_layers)
        ])

        # Final layer normalization and linear layer to output scores
        self.ln_f = nn.LayerNorm(cfg.n_embd)
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)

    def forward(self, token_ids):
        B, T = token_ids.shape
        assert T <= self.cfg.block_size, \
            f"Sequence length {T} exceeds block_size " \
            f"{self.cfg.block_size}"

        tok_emb = self.token_emb(token_ids)
        positions = torch.arange(T, device=token_ids.device)
        pos_emb = self.pos_emb(positions)

        # Combine token and position embeddings
        x = tok_emb + pos_emb

        # Pass through the transformer blocks
        x = self.blocks(x)

        # Final normalization and produce scores
        x = self.ln_f(x)
        logits = self.lm_head(x)
        return logits

Run the script to see the parameter breakdown and a test run.

Terminal

$ python src/ch09_gpt_model.py
Model parameter breakdown:
  Token embedding  :      8,320
  Position embedding:    16,384
  4 Transformer blocks:   791,552
  LM head          :      8,320
  LayerNorm (final):        256
  ---
  TOTAL            :    824,832

Input  shape: torch.Size([2, 10])   (B, T)
Output shape: torch.Size([2, 10, 65])  (B, T, vocab_size)

At each position, model outputs 65 scores.
The highest score = best guess for next character.

First position logits (top 5 scores):
  token 26: 1.861
  token 21: 1.195
  token 44: 1.118
  token 57: 1.071
  token 34: 0.972

Note: these are random (untrained model).

What just happened:

Line 1 assembles the complete GPT class.
Line 12 creates the four blocks, which hold 791,552 of the model's 824,832 parameters (weights). The original GPT-2 Small had 117 million. Our model uses the same architecture, scaled down.
Line 38 finishes the forward pass, where the model outputs 65 logits (scores) for each position.

Shape Check:

Input token IDs: [Batch, Time]
Output logits: [Batch, Time, 65] (Vocab size is 65)

Try It¶

Try It

Open src/ch09_gpt_model.py, find the line config = GPTConfig() and add a new line under it: config.n_layers = 6. Run the script again. Watch how the parameter count increases. The blocks contain the vast majority of the model's "brain".

In Business

Building the final GPT architecture is like assembling a complete production pipeline from modular components. In our house-style writing assistant, we combine a data reader (embeddings), an analysis engine (the blocks), and an output formatter (the LM head). Because it is modular, you can upgrade the engine (add more layers) without rewriting the rest of the pipeline.

Key Takeaways¶

The full model stacks embeddings, Transformer blocks, and a final LM head.
The LM head is the last step, which turns each token's internal numbers into one score for every character in the vocabulary.
The model outputs logits, which are raw scores for the next token and are not probabilities yet.
Our complete model has ~825K parameters, proving that powerful architectures can be built at a small scale.

Check Your Understanding¶

If the input sequence has 10 tokens, how many predictions does the model make?
Why does the LM head output exactly 65 numbers per token?
What is a logit?

Assignment¶

Assignment 10: Shorter text, same model

About 15 minutes. You change one number and write no new code. Work in this chapter's Colab notebook or on your own laptop.

Run it. Run src/ch09_gpt_model.py as it is. Copy three lines: the TOTAL line and the two shape lines, Input shape and Output shape.
Change one number. Near the bottom of the file, find the line B, T = 2, 10. B is how many test sequences the script feeds in, and T is how many tokens each one holds. Change 10 to 5, the length of ROMEO, and run the script again.
Compare. Which numbers in the two shape lines changed, and which did not? Did the TOTAL line change? How many sets of 65 scores does the model now give for each sequence? The five scores at the bottom shift too. They come from an untrained model, so set them aside.
Explain it. In two or three sentences a manager could follow, say what the model hands back when it reads a five-character input such as ROMEO, and why a shorter input does not make the model itself any smaller.

Hand in: both sets of three lines, your answers to step 3, and your sentences.