Coding the GPT Model & Generating Text · Build an LLM From Scratch

All four building blocks exist and the transformer block assembles them. What remains is delightfully little: replace the skeleton’s placeholders, stack the block 12 times, count what we built — and make it talk.

4.6 Coding the GPT model

The full architecture, bottom to top — chapter 2’s embeddings, 12 copies of the block, one final LayerNorm, one output projection:

Recreation of the book’s p. 61 architecture figure. Every arrow preserves (batch, 4, 768) until the output layer projects to the 50,257-token vocabulary.

</> The GPTModel — placeholders replaced, blocks stacked

class GPTModel(nn.Module):
  def __init__(self, cfg):
      super().__init__()
      self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
      self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
      self.drop_emb = nn.Dropout(cfg["drop_rate"])

      self.trf_blocks = nn.Sequential(
          *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])  #A

      self.final_norm = LayerNorm(cfg["emb_dim"])  #B
      self.out_head = nn.Linear(
          cfg["emb_dim"], cfg["vocab_size"], bias=False
      )

  def forward(self, in_idx):
      batch_size, seq_len = in_idx.shape
      tok_embeds = self.tok_emb(in_idx)
      pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
      x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
      x = self.drop_emb(x)
      x = self.trf_blocks(x)
      x = self.final_norm(x)
      logits = self.out_head(x)
      return logits

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)
out = model(batch)  #C
print("Output shape:", out.shape)
# Output shape: torch.Size([2, 4, 50257])

#A The only structural change vs the §4.1 skeleton: REAL TransformerBlocks (and a real LayerNorm). 12 of them, run in sequence.
#B One more LayerNorm after the last block stabilizes the stack's output before the vocabulary projection.
#C Same "Every effort moves you" / "Every day holds a" batch as §4.1 — same output SHAPE, but now the logits come from 12 real blocks of attention + feed-forward processing.

Counting 163 million parameters

</> Parameter count, weight tying, and memory footprint

total_params = sum(p.numel() for p in model.parameters())
print(f"Total number of parameters: {total_params:,}")
# Total number of parameters: 163,009,536  #A

total_params_gpt2 = total_params - sum(p.numel() for p in
model.out_head.parameters())
print(f"Number of trainable parameters considering weight tying: "
    f"{total_params_gpt2:,}")
# Number of trainable parameters considering weight tying: 124,412,160  #B

total_size_bytes = total_params * 4  #C
total_size_mb = total_size_bytes / (1024 * 1024)
print(f"Total size of the model: {total_size_mb:.2f} MB")
# Total size of the model: 621.83 MB

#A More than "124M"! The question the book poses — answered by weight tying below.
#B Subtracting the out_head (50257 × 768 ≈ 38.6M) reproduces the official 124M figure: original GPT-2 TIED the output layer to tok_emb (out_head.weight = tok_emb.weight), reusing one matrix for both jobs. Untied training tends to work better, so the book keeps them separate.
#C float32 = 4 bytes per parameter → 621.83 MB just to STORE the weights. Training state (gradients, optimizer moments) multiplies this further.

Where do 163M parameters actually live? Sort the table — the two vocabulary-sized matrices dominate everything:

Component⇅	Parameters⇅	% of total⇅
Token embedding (tok_emb)	38597376	23.7	+
12 × transformer blocks	85026816	52.2	+
Output head (out_head)	38597376	23.7	+
Positional embedding (pos_emb)	786432	0.5	+
Final LayerNorm	1536	0	+

Parameter budget of the book's GPT-2-small (untied). Click a row for the arithmetic.

The same GPTModel class scales to every GPT-2 size — only three config numbers change:

Model⇅	emb_dim⇅	n_layers⇅	n_heads⇅
GPT2-small (this chapter)	768	12	12	+
GPT2-medium	1024	24	16	+
GPT2-large	1280	36	20	+
GPT2-XL	1600	48	25	+

The GPT-2 family (book exercise data). Depth, width and heads grow together; the class stays identical.

4.7 Generating text

A trained-or-not GPT generates by looping: score the context, pick the next token, append, repeat. First, watch one round of that loop with the book’s real numbers (p. 65 figure):

Encode the context

"Hello, I am" → BPE token IDs [15496, 11, 314, 716], unsqueezed to a (1, 4) batch.

1 / 6

One generation round for "Hello, I am": from token IDs to the next token, "a".

</> generate_text_simple — the greedy generation loop

def generate_text_simple(model, idx, max_new_tokens, context_size):
  # idx is (batch, n_tokens) array of indices in the current context
  for _ in range(max_new_tokens):

      # Crop current context if it exceeds the supported context size
      idx_cond = idx[:, -context_size:]  #A

      # Get the predictions
      with torch.no_grad():  #B
          logits = model(idx_cond)

      # Focus only on the last time step
      # (batch, n_tokens, vocab_size) becomes (batch, vocab_size)
      logits = logits[:, -1, :]

      # Apply softmax to get probabilities
      probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)

      # Get the idx of the vocab entry with the highest probability value
      idx_next = torch.argmax(probas, dim=-1, keepdim=True)  #C

      # Append sampled index to the running sequence
      idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)

  return idx

#A The model supports at most context_size (1024) tokens — keep only the most recent ones. E.g. if the LLM supported only 5 tokens and the context is 10, just the last 5 are used.
#B Inference only: no_grad() skips gradient tracking, saving memory and time.
#C Greedy decoding: always the argmax. (Since softmax preserves ranking, argmax on the raw logits would give the identical token — the softmax is pedagogical here.)

Greedy decoding — and what ch5 changes

Always choosing the most likely token is greedy decoding : deterministic, simple, and prone to repetitive text. Chapter 5 introduces sampling techniques (temperature scaling, top-k) that modify the softmax outputs so the model doesn’t always select the most likely token — introducing variability and creativity into the generated text. That’s when the “optional” softmax stops being optional.

Running it — an untrained GPT speaks

</> Generating 6 tokens from the untrained model

start_context = "Hello, I am"

encoded = tokenizer.encode(start_context)
print("encoded:", encoded)
# encoded: [15496, 11, 314, 716]

encoded_tensor = torch.tensor(encoded).unsqueeze(0)  #A

model.eval()  #B
out = generate_text_simple(
  model=model,
  idx=encoded_tensor,
  max_new_tokens=6,
  context_size=GPT_CONFIG_124M["context_length"]
)
print("Output:", out)
# Output: tensor([[15496,   11,  314,  716, 27018, 24086, 47843, 30961, 42348,  7267]])
print("Output length:", len(out[0]))
# Output length: 10

decoded_text = tokenizer.decode(out.squeeze(0).tolist())
print(decoded_text)
# Hello, I am Featureiman Byeswickattribute argue

#A unsqueeze(0) adds the batch dimension: (4,) → (1, 4).
#B model.eval() switches to inference mode, disabling random components like dropout that are only used during training. Essential before generating — forgetting it makes outputs noisy AND non-deterministic.

Four prompt tokens + six generated = ten tokens, decoding to “Hello, I am Featureiman Byeswickattribute argue”. Perfect gibberish — and exactly the right result: with random weights, every next-token distribution is meaningless. The plumbing works end to end; coherence is what training buys, and that is chapter 5’s entire job.

4.8 Summary

Layer normalization stabilizes training by giving each layer’s outputs a consistent mean and variance.
Shortcut connections skip one or more layers by adding a layer’s input to its output, mitigating vanishing gradients in deep networks.
Transformer blocks combine masked multi-head attention with a GELU feed-forward network; GPT models are LLMs with many repeated transformer blocks — 124M to 1,542M parameters, all expressible by one GPTModel class.
Text generation decodes the model’s output tensors one token at a time, each prediction conditioned on the growing context — and with model.eval() set, deterministically so.
Untrained, a GPT generates incoherent text — underscoring why the next stage, pretraining, matters.

Stage 1 complete

Data pipeline (ch2) ✓ · attention (ch3) ✓ · architecture (ch4) ✓. The GPTModel built here — 163M parameters, weight-tying -aware, generation-ready — is the exact object chapter 5 will pretrain, evaluate, and finally load with OpenAI’s published GPT-2 weights.

📝 Check yourself: the full GPT & generation

0 / 5

1.The book's GPTModel has 163,009,536 parameters, yet GPT-2 is called a "124M" model. Where does the difference come from?
2.Which component owns the largest share of a transformer block's ~7.1M parameters?
3.In generate_text_simple, why does the code look only at logits[:, -1, :]?
4.The book notes torch.softmax before argmax is technically unnecessary. Why does the result not change without it?
5.Untrained, the model completes "Hello, I am" with "Featureiman Byeswickattribute argue". What does this demonstrate?