--- title: "Transformer from scratch in PyTorch, block by block" description: "A working decoder-only transformer in 120 lines. Pre-norm, residual streams, why the MLP is 4x wide, weight tying, stable initialisation." url: https://articles.sythra.ai/articles/transformer-from-scratch-pytorch slug: transformer-from-scratch-pytorch author: "vaibhavkothari" author_url: https://articles.sythra.ai/writers/vaibhavkothari date_published: 2026-08-29T09:26:53.619Z date_modified: 2026-08-29T10:22:15.966Z topics: ["Machine Learning", "Ai", "Python", "Deep Learning"] reading_time_minutes: 10 publisher: "Sythra" publisher_url: https://sythra.ai access: free language: en --- # Transformer from scratch in PyTorch, block by block > A working decoder-only transformer in 120 lines. Pre-norm, residual streams, why the MLP is 4x wide, weight tying, stable initialisation. Source: https://articles.sythra.ai/articles/transformer-from-scratch-pytorch · Author: vaibhavkothari · Published: 2026-08-29 · Reading time: 10 min · Topics: Machine Learning, Ai, Python, Deep Learning Attention is the famous part. It is also less than half of a transformer, and most of the design decisions that make a deep stack trainable live in the boring parts around it. This is a complete decoder-only transformer — the GPT architecture — in about 120 lines, with the reasoning for each block. It trains. ## The shape of the thing ```text tokens ──► embedding + positions │ ├─► [ Block 1 ] ─┐ ├─► [ Block 2 ] │ N identical blocks ├─► ... │ └─► [ Block N ] ─┘ │ final LayerNorm │ linear ──► logits over vocabulary each Block: x = x + attention(norm(x)) <- residual x = x + mlp(norm(x)) <- residual ``` Two sub-layers, each wrapped in a residual connection with normalisation applied *before* the sub-layer. That pattern is the whole architecture, repeated. ## Multi-head causal self-attention ```python import math import torch import torch.nn as nn import torch.nn.functional as F class CausalSelfAttention(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() assert d_model % n_heads == 0 self.n_heads = n_heads self.d_head = d_model // n_heads # One projection producing Q, K and V together — a single fused matmul # is meaningfully faster than three separate ones. self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) self.proj = nn.Linear(d_model, d_model, bias=False) self.drop = nn.Dropout(dropout) self.dropout_p = dropout def forward(self, x): B, T, C = x.shape # batch, time, channels q, k, v = self.qkv(x).split(C, dim=2) # (B, T, C) -> (B, n_heads, T, d_head) q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2) k = k.view(B, T, self.n_heads, self.d_head).transpose(1, 2) v = v.view(B, T, self.n_heads, self.d_head).transpose(1, 2) # is_causal builds the mask internally and uses a fused kernel # (FlashAttention where available) — same maths, far less memory. out = F.scaled_dot_product_attention( q, k, v, dropout_p=self.dropout_p if self.training else 0.0, is_causal=True, ) out = out.transpose(1, 2).contiguous().view(B, T, C) return self.drop(self.proj(out)) ``` Three decisions worth naming. **Fused QKV.** Three separate `nn.Linear` calls produce the same numbers, but one `(d_model, 3*d_model)` matmul is a single kernel launch on a bigger matrix — measurably faster. **`scaled_dot_product_attention` instead of hand-rolled.** It dispatches to FlashAttention when the shapes allow. Identical output, dramatically lower memory, because the `(T, T)` matrix is never materialised. **`is_causal=True` rather than a manual mask.** Same result, no mask tensor to allocate or move to the right device. ## The MLP, and why it is 4× wide ```python class MLP(nn.Module): def __init__(self, d_model, expansion=4, dropout=0.1): super().__init__() hidden = expansion * d_model self.fc1 = nn.Linear(d_model, hidden) self.fc2 = nn.Linear(hidden, d_model) self.drop = nn.Dropout(dropout) def forward(self, x): return self.drop(self.fc2(F.gelu(self.fc1(x)))) ``` Attention *moves* information between positions. It applies no per-position nonlinearity worth speaking of — the output is a weighted average of values. The MLP is where each position does actual computation on what it gathered. Expanding to `4 × d_model` and back gives the nonlinearity room to work in. The expansion factor of 4 is empirical convention rather than theory, and this block is where **roughly two thirds of the parameters live** — at `d_model = 768`, the MLP holds 4.7M parameters against attention's 2.4M. GELU rather than ReLU: smooth, non-zero gradient for small negative inputs, and consistently slightly better in transformers. ## Pre-norm vs post-norm: the change that made deep stacks trainable ```python class Block(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.attn = CausalSelfAttention(d_model, n_heads, dropout) self.ln2 = nn.LayerNorm(d_model) self.mlp = MLP(d_model, dropout=dropout) def forward(self, x): x = x + self.attn(self.ln1(x)) # PRE-norm: normalise, then sub-layer x = x + self.mlp(self.ln2(x)) return x ``` The original 2017 paper put the norm *after* the residual add: ```text post-norm (original): x = LayerNorm(x + attn(x)) pre-norm (modern): x = x + attn(LayerNorm(x)) ``` The difference looks cosmetic and is not. In pre-norm, the residual path from input to output is **completely unobstructed** — no normalisation sits on it. Gradients flow from the loss to layer 1 without passing through 48 LayerNorms. ```text post-norm: x --> [+] --> LN --> [+] --> LN --> ... gradient crosses every LN ^ ^ attn mlp pre-norm: x --> [+] --------> [+] --------> ... clean residual highway ^ ^ LN,attn LN,mlp ``` Post-norm transformers deeper than about 12 layers need careful learning-rate warmup or they diverge. Pre-norm trains stably at 100+ layers. Essentially every model since GPT-2 uses pre-norm, which is why the final `LayerNorm` before the output head is necessary — without it the residual stream leaves the last block unnormalised. ## The residual stream Worth stating explicitly, because it reframes the architecture usefully. ```text x0 = embeddings x1 = x0 + attn_1(...) + mlp_1(...) x2 = x1 + attn_2(...) + mlp_2(...) ... xN = x0 + sum of every sub-layer's contribution ``` Every block **adds** to a shared stream rather than transforming it. The residual stream is a communication channel that all layers read from and write to. Early layers write positional and lexical information; later layers read it and write semantic information. Two practical consequences follow. Gradients reach early layers directly through the addition chain, which is why depth works at all. And layers can be **partially independent** — this is why techniques like layer dropout and early exit work without the model collapsing. ## Positional information Attention is permutation-invariant: shuffle the tokens and the score matrix permutes identically. Order must be injected. ```python class Transformer(nn.Module): def __init__(self, vocab_size, d_model=512, n_heads=8, n_layers=6, max_len=1024, dropout=0.1): super().__init__() self.tok_emb = nn.Embedding(vocab_size, d_model) self.pos_emb = nn.Embedding(max_len, d_model) # learned positions self.drop = nn.Dropout(dropout) self.blocks = nn.ModuleList( Block(d_model, n_heads, dropout) for _ in range(n_layers) ) self.ln_f = nn.LayerNorm(d_model) self.head = nn.Linear(d_model, vocab_size, bias=False) # Weight tying: input embedding and output projection share weights. # Saves vocab_size * d_model parameters and usually improves quality. self.head.weight = self.tok_emb.weight self.apply(self._init) # Scale residual-path output projections by 1/sqrt(2*n_layers) so the # residual stream's variance does not grow with depth. for name, p in self.named_parameters(): if name.endswith("proj.weight") or name.endswith("fc2.weight"): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * n_layers)) def _init(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, idx, targets=None): B, T = idx.shape pos = torch.arange(T, device=idx.device) x = self.drop(self.tok_emb(idx) + self.pos_emb(pos)) for block in self.blocks: x = block(x) x = self.ln_f(x) logits = self.head(x) loss = None if targets is not None: loss = F.cross_entropy( logits.view(-1, logits.size(-1)), targets.view(-1) ) return logits, loss ``` Three details in that constructor do real work. **Weight tying** (`self.head.weight = self.tok_emb.weight`). Input and output embeddings represent the same relationship in opposite directions. Sharing them saves `vocab_size × d_model` parameters — 38M for a 50k vocabulary at 768 dims — and typically improves quality slightly. **Scaled residual initialisation.** Each block adds to the residual stream, so variance accumulates with depth. Scaling the output projections by `1/sqrt(2·n_layers)` — two residual writes per block — keeps the stream's scale stable from layer 1 to layer 48. Without this, deep models start with an exploding activation scale and need a long warmup to recover. **Learned position embeddings** are the simplest option. They do not extrapolate beyond `max_len` at all. Modern models use RoPE instead, which encodes *relative* position directly into the attention computation and extrapolates far better. ## Generation ```python @torch.no_grad() def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None): model.eval() for _ in range(max_new_tokens): idx_cond = idx[:, -model.pos_emb.num_embeddings:] # crop to context logits, _ = model(idx_cond) logits = logits[:, -1, :] / temperature # last position only if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float("inf") probs = F.softmax(logits, dim=-1) idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1) return idx ``` Note the inefficiency: every step re-runs the whole sequence. Real inference uses a **KV cache** — keys and values for previous positions do not change, so you compute them once and append. That turns generation from `O(T^2)` per token into `O(T)`, and it is the single most important inference optimisation. ## Training loop ```python model = Transformer(vocab_size=50257, d_model=512, n_heads=8, n_layers=6) opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1, betas=(0.9, 0.95)) # Warmup then cosine decay — transformers are unstable at full LR from step 0 warmup, total = 2000, 100_000 sched = torch.optim.lr_scheduler.LambdaLR( opt, lambda s: s / warmup if s < warmup else 0.5 * (1 + math.cos(math.pi * (s - warmup) / (total - warmup))), ) for step, (x, y) in enumerate(loader): _, loss = model(x, y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # essential opt.step() sched.step() opt.zero_grad(set_to_none=True) ``` Two non-negotiables in there. **Gradient clipping at 1.0** — transformers produce occasional huge gradients, and one unclipped spike destroys a run that took days. **Learning-rate warmup** — Adam's second-moment estimate is unreliable in the first few hundred steps, and a full learning rate on top of that unreliability diverges. `betas=(0.9, 0.95)` rather than the PyTorch default `0.999` is the standard transformer setting: a shorter second-moment window adapts faster to the changing loss landscape. ## Parameter budget For `d_model = 768`, `n_layers = 12`, `vocab = 50257`: | Component | Formula | Parameters | |---|---|---| | Token embedding (tied with head) | `V × d` | 38.6M | | Position embedding | `max_len × d` | 0.8M | | Attention, per layer | `4d²` | 2.4M | | MLP, per layer | `8d²` | 4.7M | | **Total** | | **~124M** | Two things stand out. The MLP is **twice** the attention parameters per layer, which surprises people who think of transformers as "attention models". And at this scale the embedding table alone is nearly a third of the model — which is why weight tying matters and why vocabulary size is an architectural decision, not a preprocessing detail. ```remember # Remember this Pre-norm -> `x = x + attn(norm(x))`, keeps the residual path clean The MLP -> 4x wide, and about two thirds of the parameters Weight tying -> input embedding shared with the output projection --- Warmup and gradient clipping at 1.0 are not optional. One unclipped spike destroys a run that took days. ``` ## Frequently asked questions ### Why pre-norm instead of post-norm? Pre-norm leaves the residual path completely clean, so gradients reach early layers without crossing every normalisation layer. Post-norm models deeper than about 12 layers need careful warmup to train at all; pre-norm is stable at 100+ layers. Every modern model uses pre-norm. ### Why is the MLP four times wider than the model dimension? Attention moves information between positions but applies almost no per-position nonlinearity. The MLP is where each position computes on what it gathered, and the 4× expansion gives that computation room. The factor is empirical convention, and this block holds about two thirds of the parameters. ### What is weight tying and should I use it? Sharing the input embedding matrix with the output projection. It saves `vocab_size × d_model` parameters and usually improves quality slightly, since both matrices encode the same token-to-vector relationship. Standard practice for language models. ### Why do transformers need learning-rate warmup? Adam's second-moment estimate is unreliable in the first hundreds of steps, and applying a full learning rate on top of unreliable statistics diverges. Warm up linearly over 1–5% of training, then decay. Pair it with gradient clipping at 1.0 — one unclipped gradient spike can destroy a multi-day run. ### Should I use learned positional embeddings? They are the simplest thing that works, and they do not extrapolate past the trained maximum length at all. If long-context behaviour matters, use RoPE, which encodes relative position inside the attention computation and generalises much better beyond the training length. ## Next reading on Sythra Articles - [The attention mechanism, worked out by hand](https://articles.sythra.ai/articles/attention-mechanism-visualized) - [CNNs from scratch in NumPy](https://articles.sythra.ai/articles/cnn-from-scratch-numpy) - [Embeddings explained mathematically](https://articles.sythra.ai/articles/embeddings-explained-mathematically) ## Glossary (terms defined in this article) - **RoPE** (advanced) — Rotary Position Embedding: rotates query and key vectors by an angle proportional to position, so attention scores depend on relative distance