Mini GPT Architecture
Build a complete decoder-only Transformer architecture from scratch in PyTorch
What You Will Build
In this milestone project, you will assemble every component of a modern decoder-only Large Language Model (`GPT`) from scratch in PyTorch. Using the causal multi-head self-attention engine you built in Project 9, you will construct a Pre-LayerNorm `TransformerBlock` (`x = x + attn(ln1(x))`, `x = x + mlp(ln2(x))`), stack `N` blocks atop token and positional embeddings, apply final normalization and the language-model output head (`lm_head`), implement scaled residual weight initialization (`0.02 / √(2N)`), and verify the complete architecture by overfitting a tiny batch to near-zero loss and generating text autoregressively.
Why This Project Matters in the AI/LLM Stack
Once you complete Project 10, you have built the exact architectural blueprint of GPT-2/GPT-3/Llama from scratch. In Project 11, you will pair this architecture with your Project 8 BPE tokenizer and a pretraining data pipeline to train `MyLLM v1`.
Learning Objectives
- Define a clean `GPTConfig` dataclass encapsulating `vocab_size`, `block_size`, `n_layer`, `n_head`, `n_embd`, and `dropout`.
- Implement the Pre-Norm `TransformerBlock` combining `CausalMultiHeadAttention`, `LayerNorm` (or `RMSNorm`), and `MLP` (`4 × n_embd` expansion with `GELU`/`SwiGLU`) around residual pathways.
- Assemble the full `MiniGPT` model with embeddings, stacked blocks, `lm_head` weight tying, scaled initialization, and Cross-Entropy loss.
- Verify the architecture with forward-shape checks, initial loss checks (`-ln(1/V)`), single-batch overfit debugging, checkpoint save/load, and autoregressive generation.
Prerequisites
- Project 7 (Character Language Model) and Project 9 (Self-Attention Engine)
- Phase 5 Modules 5.3, 5.4, & 5.5 (Transformer Block, Decoder-Only Architecture, Sampling)
Required Mathematical Concepts
- Pre-Norm Residual Stream: h_l = x_l + MHA(LN(x_l)), x_{l+1} = h_l + FFN(LN(h_l))
- Position-wise Feed-Forward Network: FFN(x) = W₂ · GELU(W₁ x + b₁) + b₂ with d_ff = 4 d_model
- Layer Normalization / RMSNorm across channel dimension d_model
- Residual depth scaling initialization: σ_proj = 0.02 / √(2 · n_layers)
- Weight tying between input token embedding W_E and output projection W_head = W_Eᵀ
Step-by-Step Implementation Tasks (15 Steps)
Check off each engineering task as your code passes its acceptance criteria. Progress is saved automatically in your browser.
Embeddings, Feed-Forward MLP & Pre-LN Transformer Block
Implement token/positional embeddings, LayerNorm, position-wise MLP, and residual Pre-Norm Transformer blocks.
The token embedding table `wte = nn.Embedding(vocab_size, n_embd)` maps integer token IDs `(B, T)` into vectors `(B, T, n_embd)`.
Dot-product attention is permutation-equivariant without positional information; learned positional embeddings `wpe = nn.Embedding(block_size, n_embd)` (or RoPE) inject sequence order `0..T-1`.
Integrate the `CausalMultiHeadAttention` module from Project 9 as the communication sublayer of each Transformer block.
While self-attention routes information across token positions, the position-wise MLP (`n_embd -> 4 * n_embd -> n_embd` with `GELU`) processes each token representation independently.
Applying `LayerNorm` (or `RMSNorm`) before each sublayer (`Pre-LN`) keeps activation variances stable and leaves the residual stream unobstructed for gradient flow.
Residual connections `x = x + sublayer(x)` act as a gradient superhighway: during `.backward()`, `∂(x + f(x))/∂x = I + ∂f/∂x`, allowing gradients to flow directly from the loss to the earliest layers.
Combining `ln_1`, `attn`, `ln_2`, and `mlp` forms a self-contained `TransformerBlock(nn.Module)`.
Full Decoder Stack, Language Model Head & Cross-Entropy Loss
Stack `n_layer` Transformer blocks, add final LayerNorm and weight-tied `lm_head`, and compute sequence Cross-Entropy loss.
Using `nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layer)])` registers all `n_layer` blocks properly with PyTorch's parameter system.
The final linear projection `lm_head = nn.Linear(n_embd, vocab_size, bias=False)` maps normalized hidden states `(B, T, n_embd)` to vocabulary logits `(B, T, vocab_size)`.
When `targets` of shape `(B, T)` are provided to `MiniGPT.forward(idx, targets)`, the model computes `F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))`.
Training Loop, Autoregressive Generation & Checkpointing
Implement the training step, autoregressive token generation, and checkpoint serialization.
Training a Transformer uses `AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)` and gradient norm clipping (`clip_grad_norm_(model.parameters(), 1.0)`).
During `.generate(idx, max_new_tokens, temperature, top_k, top_p)`, the model crops `idx` to `block_size`, computes logits for the last token `logits[:, -1, :]`, applies temperature and Top-k/Top-p filtering, and appends the sampled token.
Saving both `config` (`asdict(model.config)`) and `model.state_dict()` allows reconstructing the exact `MiniGPT` architecture and weights from a single file.
Shape Validation, Mask Checks & Short Training Run Verification
Validate tensor shapes and causal masks and run a short training / single-batch overfit verification test.
Verifying shapes across varied `(B, T)` inputs and confirming prefix causality across the full multi-layer stack ensures no block leaks future information.
Overfitting a single mini-batch of 4–8 sequences to near-zero loss (`< 0.05`) in 50 steps is the gold-standard debugging test before launching a long pretraining run.
Starter & Verification Code Examples
Pre-LN Transformer Block & MiniGPT Core Forward Pass
Shows the clean Pre-LayerNorm residual block structure and weight-tied decoder forward pass.
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
@dataclass
class GPTConfig:
vocab_size: int = 512
block_size: int = 128
n_layer: int = 4
n_head: int = 4
n_embd: int = 128
class TransformerBlock(nn.Module):
def __init__(self, config: GPTConfig) -> None:
super().__init__()
self.ln_1 = nn.LayerNorm(config.n_embd)
self.attn = CausalMultiHeadAttention(config.n_embd, config.n_head, config.block_size)
self.ln_2 = nn.LayerNorm(config.n_embd)
self.mlp = nn.Sequential(
nn.Linear(config.n_embd, 4 * config.n_embd),
nn.GELU(),
nn.Linear(4 * config.n_embd, config.n_embd),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
attn_out, _ = self.attn(self.ln_1(x))
x = x + attn_out
x = x + self.mlp(self.ln_2(x))
return xHow to Test Your Implementation
Run forward shape checks, initial loss sanity check (`≈ ln(V)`), multi-layer causality test, checkpoint reload test, and the 50-step single-batch overfit test.
Initial Untrained Loss Sanity Check (`≈ ln(vocab_size)`)
Confirm weight initialization is properly scaled so initial predictions are near-uniform.
Steps: Instantiate `MiniGPT(GPTConfig(vocab_size=256, block_size=32, n_layer=4, n_head=4, n_embd=64))` and compute loss on random tokens.
Pass Criteria: Initial loss is between `5.40` and `5.70` (`ln(256) ≈ 5.545`).
Single-Batch Overfitting Debugging Exercise
Prove the entire forward, loss, backward, and AdamW pipeline is wired without gradient bugs.
Steps: Train on a fixed `(4, 16)` batch for 50 steps with `lr=1e-3`.
Pass Criteria: Loss decreases from `~5.5` to `< 0.05`.
Overall Project Acceptance Criteria
- Complete decoder-only `MiniGPT` architecture implemented from scratch in PyTorch.
- Pre-Norm residual `TransformerBlock`, `CausalMultiHeadAttention`, `FeedForwardMLP`, weight tying, and scaled initialization verified.
- Forward shapes, multi-layer causal masking, initial loss sanity check, and checkpoint save/load verified by unit tests.
- Single-batch overfitting test drives training loss near zero and generates expected continuations.
Expandable Hints & Debugging Tips
Try implementing each stage on your own first. Expand a hint below if you want to verify tensor shapes, mathematical identities, or common edge cases.
Optional Stretch Challenges
Compare Positional Encodings: Learned `wpe` vs Sinusoidal vs RoPE
Make the positional encoding strategy configurable in `GPTConfig` (`'learned' | 'sinusoidal' | 'rope'`) and compare validation loss and length generalization beyond `block_size`.
What You Learn: Why modern LLMs transitioned from learned absolute positional embeddings to Rotary Positional Embeddings (RoPE).
Modern Llama-Style Upgrades: `RMSNorm` & `SwiGLU` FFN
Replace `LayerNorm` with root-mean-square normalization (`RMSNorm`) and replace the standard 2-layer `GELU` MLP with a 3-matrix gated `SwiGLU` feed-forward network.
What You Learn: How the GPT-2 architecture evolved into the modern Llama/Mistral/Qwen block architecture.
What Your Completed Repository Should Contain
A trainable GPT-style model with randomly initialized parameters.
Deliverable Artifacts
- Complete `MiniGPT` decoder-only Transformer (`minigpt/model.py`)
- AdamW training step, checkpoint manager, and Top-k / Top-p `.generate()` method
- Automated test suite verifying shapes, causality, initial loss, and single-batch overfitting
After Completing This Project, You Should Be Able To...
- Explain causal self-attention and how Pre-Norm residual connections stabilize deep Transformer stacks.
- Implement Transformer blocks combining `LayerNorm`, `CausalMultiHeadAttention`, and position-wise `MLP` sublayers.
- Track token and tensor dimensions (`(B, T)`, `(B, T, n_embd)`, `(B, H, T, T)`, `(B, T, vocab_size)`) across the entire model.
- Train a decoder-only language model from random initialization and verify it by overfitting a single batch.
- Generate text autoregressively using temperature, Top-k, and Top-p nucleus sampling.
- Debug basic training and inference problems including initial loss spikes, mask alignment, and checkpoint reloading.
GitHub Repository Submission
Link your public or private GitHub repository URL to keep a structured portfolio of your from-scratch AI builds. Stored locally in your browser (self-reported; no GitHub tokens or remote code execution).
Suggested Repository: mini-gpt-from-scratch
Recommended Directory Layout:
minigpt/config.py minigpt/model.py minigpt/train.py tests/test_minigpt.py README.md
Recommended README Documentation Checklist
- Architecture diagram & parameter count breakdown by component
- Single-batch overfitting verification curve
- Checkpoint save/load and generation examples