CapstonePython

Transformer Block from Scratch

One block is the unit of every modern decoder LLM: layer norm, multi head attention, residual, MLP, residual. Pre-LN trains stably without warmup, while the post-LN layout from the original paper does not. This lesson builds both side by side and explains which one survives a 12 layer stack at common learning rates (the demo runs six layers).

TIME: ~90 minutesPREREQUISITES: Phase 19 lessons 30 to 33 (tokenizer, embeddings, attention math, batched data loader)

Learning Objectives

Build a transformer block in PyTorch from the four moving pieces: LayerNorm, multi head causal attention, residual connections, position wise MLP.
Place the LayerNorms in two configurations (pre-LN and post-LN) and explain why one trains stably without warmup.
Implement causal masking inside the multi head attention so token `i` cannot see tokens `j > i`.
Track gradient flow through both variants on a six layer demo stack and read the result, then reason about why it gets worse at twelve layers.
Reuse the block as a drop-in unit when the next lesson assembles a 124 million parameter GPT.

The Problem

A transformer is one block repeated. Get the block wrong once, repeat it twelve times, and you ship a model that diverges in the first epoch or leans on warmup hacks for the rest of training. Both failure modes show up the first time someone stacks blocks naively. In one, the attention layer attends to the future. In the other, the LayerNorm sits where it cannot tame the residual signal at depth.

The fix is mechanical once you see it. The block has exactly two residual paths and exactly two normalization positions. Choose the positions correctly and the rest of the stack is bookkeeping.

Warm-up0/3

1. A decoder only language model is trained to predict each next token. What goes wrong if position i can read tokens after position i?Warm-up

2. A batch holds 2 sequences of 16 token ids, and an embedding layer maps each id to a 384 dimension vector. What shape does the embedding output have?Warm-up

3. Softmax turns a row of attention scores into weights. Which property do the output weights have?Warm-up

The Concept

A decoder only transformer block is a function that takes a tensor of shape (batch, sequence, embedding) and returns a tensor of the same shape. Two sublayers do the work inside it.

Mermaid diagram

This is the pre-LN variant. The LayerNorm sits inside the residual branch, before the sublayer, so the residual connection carries the unnormalized signal forward.

The post-LN variant moves the LayerNorm to after the residual add.

Mermaid diagram

The shapes are identical, but training behavior is not. With post-LN, the gradient flowing back through the residual path has to pass through the LayerNorm. At depth twelve (the demo below runs six layers) and learning rate 3e-4, that gradient shrinks fast enough to need a warmup schedule. Pre-LN leaves the residual path unnormalized, so gradients propagate cleanly to the embedding layer. That is why GPT-2 onward ships pre-LN.

Causal multi head attention

The attention sublayer projects the input three ways into query, key, and value tensors. Each is reshaped from (B, T, D) to (B, H, T, D/H), where H is the head count. Scaled dot product attention computes softmax(Q K^T / sqrt(d_k)) per head. The upper triangle is masked to negative infinity before the softmax, and the result multiplies V. Heads are concatenated back into a single (B, T, D) tensor and projected once more. The mask is what makes the model causal. Forget it and you train a model that cheats.

The MLP

The position wise MLP applies the same two layer network to every token independently. The hidden width is four times the embedding width, the activation is GELU, and a dropout follows the second linear. Tokens do not talk to each other inside the MLP. All token mixing happens in attention.

Residual connections do two things

They make the gradient path additive across depth, which keeps the gradient norm in scale through twelve layers. They also let each block learn an additive update to the running representation instead of a full replacement. Both effects are why the block scales.

Check: The Concept0/17

1. Which two components do the main work inside a decoder transformer block?Define

2. What is pre-LN in a transformer block?Define

3. What is post-LN, the configuration the original 2017 attention paper shipped?Define

4. What does the causal mask in multi head attention do?Define

5. What is the residual stream in a transformer?Define

6. What are the defining properties of the position wise MLP in the block?Define

7. Why does the block need an attention sublayer when the MLP already transforms every token?Explain

8. What is the key difference between pre-LN and post-LN at depth?Explain

9. Why does setting masked logits to negative infinity before the softmax stop token i from using future tokens?Explain

10. What two effects do residual connections have on training?Explain

11. Why would the post-LN gradient problem look worse on a twelve layer stack than on the six layer demo?Explain

12. Why can the block be reused as a drop-in unit when a GPT is assembled?Explain

13. The mask is torch.triu(ones(T, T), diagonal=1) as booleans, where True means masked. For a sequence of 3 tokens, which positions can the token at index 1 attend to?Apply

14. You train a 12 layer post-LN stack at learning rate 3e-4 with no warmup and the loss diverges early. Which change is the most direct fix?Analyze

15. A model reaches near zero training loss within a few steps, yet its generated text is nonsense. Which block bug fits these symptoms best?Analyze

16. A teammate says pre-LN is stable because its LayerNorms normalize the residual stream. What is wrong with that explanation?Analyze

17. You must train a deep decoder at learning rate 3e-4 and do not want to tune a warmup schedule. Which block design fits?Design

Build It

code/main.py implements:

class LayerNorm with learnable scale and shift, biased variance with eps=1e-5, applied per token vector.
class MultiHeadAttention with num_heads, head_dim = d_model // num_heads, a fused QKV projection, a registered causal mask, and attention and residual dropout.
class FeedForward with two linear layers, GELU activation, and dropout.
class TransformerBlock with a pre_ln flag that toggles between the two variants.
A demo that builds a 6 layer pre-LN stack and a 6 layer post-LN stack with identical inputs, then prints (a) the output shape and (b) the gradient norm at the embedding after one backward pass.

Run it:

bash

python3 code/main.py

The output is a shape check on both stacks and the gradient norms side by side. The pre-LN stack's embedding gradient is about 200 times larger than the post-LN stack's at the same learning rate (210x in the demo run). That gap is the empirical signal that pre-LN trains without warmup.

Check: Build It0/9

1. In the Build It demo, what does comparing the six layer pre-LN and post-LN stacks show?Apply

2. The demo config has d_model=192 and num_heads=6. For a batch of 2 with 10 tokens, what are head_dim and the shape of q after the view and transpose in MultiHeadAttention?Apply

3. With d_model=192, what is the output width of self.qkv, and what does qkv.split(self.d_model, dim=-1) return?Apply

4. TransformerBlock.forward runs with pre_ln=True, but attn and mlp both return all-zero tensors for input x. What does the block return, and how does pre_ln=False compare?Apply

5. A test replaces the input vectors at positions 5 to 7 of an 8 token sequence with random values and reruns the causal attention layer. Which output positions change?Apply

6. MultiHeadAttention is built with context_length=8 and then called on a (1, 16, 32) tensor. What happens?Apply

7. A FeedForward is built with d_model=32 and the default mlp_expansion. What are the layer sizes?Apply

8. What happens when MultiHeadAttention is built with d_model=64 and num_heads=5?Apply

9. A LayerNorm(64) with the default scale and shift gets x = randn(2, 16, 64) * 5 + 3. What do the statistics over the last dimension look like in the output?Apply

Stack

torch for the tensor math, autograd, and nn.Module plumbing.
No transformers, no pretrained weights. The block is built from primitives.

Production patterns in the wild

Three patterns turn the textbook block into something you can ship.

Fused QKV projection. Three separate linear layers cost three kernel launches and three matmuls. One linear layer of width 3 * d_model does the same work in a single launch, then splits the output along the last axis. The fused path is faster on every accelerator and matches what the reference implementations of GPT-2, LLaMA, and Mistral ship.

Registered causal mask buffer. The mask depends only on the maximum context length. Allocate it once at construction with register_buffer, slice the active window on each forward pass, and skip the per-call allocation. Forget this and the mask becomes an allocator hot spot at long context.

Dropout in three places, none on the residual stream. Dropout belongs after the attention softmax (attention dropout), after the attention output projection, and after the second linear of the MLP (the last two are residual dropout, applied to each sublayer's output before it is added back). A dropout on the residual stream itself breaks the additive identity that lets the gradient flow at depth. Some early implementations got this wrong and paid for it with brittle training.

Check: Production patterns in the wild0/5

1. What is a fused QKV projection?Define

2. Why use a fused QKV projection instead of three separate linears?Explain

3. Why register the causal mask as a buffer rather than allocating it per forward call?Explain

4. Which dropout placement would break the additive identity the block relies on at depth?Analyze

5. Profiling a long context run shows allocator time growing with every forward call. Which mistake in the attention module matches?Analyze

Use It

The block plugs straight into the GPT assembly in lesson 35 without modification.
Pre-LN is what every modern open weights LLM uses. Post-LN is what the original 2017 attention paper used. Knowing both is enough to read any decoder architecture you will meet.
Swap the GELU for SiLU and you have the LLaMA family activation. Swap the LayerNorm for RMSNorm and you have the LLaMA family normalization. The skeleton stays the same.

Check: Use It0/2

1. What changes turn this block into the LLaMA family skeleton?Define

2. You replace nn.LayerNorm with a hand written RMSNorm and still want the block to drop into the GPT assembly. Which constraint must the replacement meet?Design

Exercises

Set the existing use_bias flag in BlockConfig to False. Modern open weights LLMs ship without biases on the linear layers. Measure the parameter savings in a 12 layer 768 dim model.
Replace nn.LayerNorm with a hand rolled RMSNorm and verify the output shape is unchanged.
Add a flag that returns the attention weights for the first head as a (B, T, T) tensor. Plot the upper triangle to confirm it is zero after softmax.
Build a sanity check that feeds a (2, 16, 384) tensor with H=6 through both variants and asserts the forward outputs are different (for example, not torch.allclose) when weights are initialized identically and dropout is set to zero.

Key Terms

Term	What people say	What it actually means
Pre-LN	"Pre norm"	LayerNorm inside the residual branch, before each sublayer; the residual carries the unnormalized signal
Post-LN	"Post norm"	LayerNorm after the residual add; what the 2017 paper shipped and what needs warmup
Causal mask	"Triangle mask"	The upper triangle of the attention logits set to negative infinity so token i cannot read token j when j is greater than i
Fused QKV	"Combined projection"	One linear of width 3D instead of three linears of width D; one kernel, one matmul
Residual stream	"Skip connection"	The unnormalized tensor that flows top to bottom through every block; what each block adds to

Author's Note

I'd keep pre_ln=True as the default and leave the flag in place, so the post-LN variant stays one toggle away. Running both stacks on identical inputs is the fastest way to see the gradient gap for yourself, and the demo does exactly that. I'd also keep dropout off the residual stream itself: it breaks the additive identity, and everything the block does at depth leans on that identity.