LM4-T4 — The Transformer Block and Assembling GPT

Learning objectives

After this lesson, you should be able to:

• Write the equations of a pre-LN Transformer block and apply them in the right order.
• Assemble a full GPT from embeddings, blocks, a final norm and a head.
• Explain where dropout is applied and how to disable it.
• Verify an implementation with shape, causality and parameter-count tests.

The idea: two corrections to the stream

A block takes the residual stream x of shape (B, n, d) and makes two corrections to it. The first correction comes from attention: each position reads from earlier positions and writes what it found into its own vector. The second comes from the feed-forward network: each position transforms its own vector. Both corrections are computed from a normalised copy of the stream and added back to the unnormalised stream.

The order of operations

For each sublayer the recipe is the same: normalise, transform, drop out, add. In code it is one line each: x = x + drop(attn(norm1(x))) and x = x + drop(ffn(norm2(x))). Forgetting the x + turns the block into a plain stack of layers and loses the residual path; putting the norm after the addition gives post-LN, a different model.

Stacking

The model applies the same block structure N times, each with its own weights. Every block returns a tensor of the same shape that it received, so blocks can be placed in an nn.Sequential. The embeddings at the bottom and the output head at the top are the only places where the shape changes meaning (integers to vectors, vectors to vocabulary scores).

Dropout

Dropout randomly sets a fraction of numbers to zero during training and rescales the rest, as regularisation. In GPT-2 it is applied to the embeddings, to the attention weights and to the output of each sublayer, before the addition. In evaluation mode (model.eval) it does nothing. Large pretraining runs on huge datasets often use no dropout; small experiments on tiny text may benefit from it.

Initialisation

Weights start as small random numbers. The default PyTorch initialisation of linear layers is adequate for this track; GPT-2 additionally scales the output projections of deep stacks down by the square root of the number of blocks, so that the residual stream does not grow too fast at the start. Neither affects the structure.

The mathematics: the block and the stack

Let x be the stream entering a block, with shape (B, n, d). With the sublayers Attn and FFN and the layer norms LN₁ and LN₂ (each with its own parameters):

x

the incoming residual stream, shape (B, n, d).

x′

the stream after the attention correction; same shape.

y

the outgoing stream; same shape as x.

Attn

causal multi-head attention (Lesson LM3-T5) with input and output width d.

FFN

the position-wise feed-forward network (Lesson LM4-T3).

The model composes N such blocks. With embeddings x₀ (Lesson LM4-T1) and blocks y = B_l(x):

Parameter count of a block. Adding attention (4d² + d with no bias on Q, K and V), the feed-forward network (8d² + 5d) and two layer norms (4d):

For d = 768 that is 7,077,888 + 7,680 = 7,085,568. With biases on Q, K and V (as in the released GPT-2 weights), add 3d = 2,304.

Causality is preserved through the stack. Only attention mixes positions, and its mask lets position i see positions ≤ i. Layer norm, the FFN, dropout and the residual additions act on each position separately. So the output at position i depends only on the inputs at positions ≤ i in every block and therefore in the whole stack.

A worked example: a compact GPT in PyTorch

The following model uses PyTorch’s own LayerNorm, GELU and scaled_dot_product_attention to stay short; the notebook writes each part from scratch. Configuration: vocabulary 100, context 16, width 32, 4 heads, 2 blocks.

import torch
import torch.nn as nn
import torch.nn.functional as F

class Block(nn.Module):
    def __init__(self, d, heads, p=0.0):
        super().__init__()
        self.h = heads
        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.proj = nn.Linear(d, d)
        self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(approximate="tanh"), nn.Linear(4 * d, d))
        self.drop = nn.Dropout(p)

    def attn(self, x):
        B, n, d = x.shape
        q, k, v = self.qkv(x).view(B, n, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)   # each (B, H, n, dh)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.proj(out.transpose(1, 2).reshape(B, n, d))

    def forward(self, x):
        x = x + self.drop(self.attn(self.n1(x)))     # first correction: attention
        x = x + self.drop(self.ff(self.n2(x)))       # second correction: feed-forward
        return x

class GPT(nn.Module):
    def __init__(self, V=100, L=16, d=32, heads=4, layers=2):
        super().__init__()
        self.tok, self.pos = nn.Embedding(V, d), nn.Embedding(L, d)
        self.blocks = nn.Sequential(*[Block(d, heads) for _ in range(layers)])
        self.norm, self.head = nn.LayerNorm(d), nn.Linear(d, V, bias=False)

    def forward(self, idx):
        x = self.tok(idx) + self.pos(torch.arange(idx.shape[1]))
        return self.head(self.norm(self.blocks(x)))

torch.manual_seed(0)
m = GPT().eval()
idx = torch.randint(0, 100, (2, 10))
print(tuple(m(idx).shape))                                   # (2, 10, 100)

idx2 = idx.clone(); idx2[:, -1] = (idx2[:, -1] + 1) % 100    # change only the last token
print(torch.allclose(m(idx)[:, :-1], m(idx2)[:, :-1], atol=1e-6))   # True: earlier outputs are unchanged

d, V, L, N = 32, 100, 16, 2
expected = V * d + L * d + N * (12 * d * d + 10 * d) + 2 * d + V * d
print(sum(p.numel() for p in m.parameters()), expected)      # 32192 32192

The two printed counts, 32,192 each, agree because this implementation has no bias on the Q, K, V projection and one bias on the output projection, as in the formula. The first print confirms the output shape (B, n, V), and the second is the leakage test: changing the last token does not change any earlier output.

Engineering implications and limitations

Write the shape and leakage tests first. The two checks above catch most mistakes: a wrong residual, a wrong reshape, a missing mask. Keep them in the repository as unit tests.

Match a reference. To be sure your reimplementation is exactly GPT-2, load the published weights into it and compare logits with a reference implementation on the same input (Module 5 does this). Parameter counts alone do not prove the wiring is correct.

Eval mode matters. Forgetting model.eval leaves dropout on, making outputs random and tests flaky. Forgetting to switch back with model.train() disables dropout in training.

Share code between blocks, not weights. Each block has its own parameters. A list comprehension builds N distinct blocks; multiplying a list ([Block()] * N) would reuse the same block N times, a classic bug that still runs.

Gradient checkpointing. Activations from every block are stored for the backward pass. To fit larger models or batches, frameworks can recompute a block’s activations instead (checkpointing), trading about a third more compute for much less memory.

Compile and fuse. The block is a good target for compiler optimisations (torch.compile, fused attention kernels). Optimise only after the unfused version passes its tests.

For robots. Embedded and on-robot models are usually small: a few blocks with narrow width. The same block code applies; what changes is the budget for parameters, memory and latency (Module 10).

Common misconceptions

Engineering takeaways

1. A pre-LN block is x′ = x + Dropout(Attn(LN₁(x))) followed by y = x′ + Dropout(FFN(LN₂(x′))); the shape (B, n, d) never changes.
2. A GPT is embeddings, N such blocks, a final norm and a linear head; only attention mixes positions, so the whole stack stays causal.
3. One block has 12d² + 10d parameters (7,085,568 for d = 768, without Q, K, V biases).
4. Dropout is applied to embeddings, attention weights and sublayer outputs, and is off in eval mode.
5. Test the output shape, causality (change a later token, earlier outputs must not move) and the exact parameter count.

References and further reading

See Chapter 4 of the book (connecting attention and linear layers in a Transformer block, and coding the GPT model) for the full implementation this lesson summarises.

• Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.
• Build a Large Language Model (From Scratch), Chapter 4: Implementing a GPT Model from Scratch (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: the output logits, where the parameters are, and the tensor shapes of the complete model.

Related content

Nothing below is required to finish this lesson.

• Before
  LM4-T3: Feed-Forward Layers and GELU

  Explain the position-wise feed-forward network, its expansion factor and its nonlinearity.

• Lab
  LM4-L3: Assemble a GPT Model in PyTorch

  Implement the GPT components, run a forward pass, count parameters and generate text from random weights.