how-to-train-your-gpt

What this skill does

```markdown
---
name: how-to-train-your-gpt
description: A 12-chapter, 3671-line annotated textbook for building a modern LLaMA-style GPT from scratch in PyTorch — tokenizer through inference.
triggers:
  - "build a GPT from scratch"
  - "train a language model"
  - "implement transformer attention"
  - "how does tokenization work"
  - "implement RoPE positional encoding"
  - "write a training loop for LLM"
  - "implement KV cache inference"
  - "understand transformer architecture"
---

# How to Train Your GPT

> Skill by [ara.so](https://ara.so) — Daily 2026 Skills collection.

A 12-chapter interactive textbook that teaches you to build a modern LLaMA 3-style decoder-only Transformer from absolute scratch. Every line of code is annotated with WHAT it does and WHY it's there. Covers BPE tokenization → embeddings → RoPE → multi-head attention → training → inference.

## Installation

```bash
git clone https://github.com/raiyanyahya/how-to-train-your-gpt.git
cd how-to-train-your-gpt

python -m venv gpt_env
source gpt_env/bin/activate        # Mac/Linux
# gpt_env\Scripts\activate         # Windows

pip install torch tiktoken datasets numpy matplotlib
```

Verify GPU availability:
```bash
python -c "import torch; print(f'CUDA: {torch.cuda.is_available()}, Device: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU\"}')"
```

## Project Structure

```
chapters/
├── 00_overview.md          # Big picture: what is a GPT?
├── 01_setup.md             # Environment, PyTorch basics
├── 02_tokenization.md      # BPE walkthrough
├── 03_embeddings.md        # Token embeddings, semantic space
├── 04_positional_encoding.md  # RoPE math + numeric example
├── 05_attention.md         # ⭐ Core: Q,K,V, scaling, causal mask (713 lines)
├── 06_transformer_block.md # RMSNorm, SwiGLU, residuals
├── 07_gpt_model.md         # Full 124M parameter model
├── 08_training.md          # AdamW, cosine warmup, mixed precision
├── 09_inference.md         # KV cache, sampling strategies
├── 10_full_script.md       # Runnable main.py (copy this to run)
└── 11_glossary.md          # Architecture provenance table
```

## Architecture Overview

This project implements a **LLaMA 3-style** decoder-only Transformer with:

| Component | Technique | Why |
|-----------|-----------|-----|
| Positional encoding | RoPE | Relative position, no learned params |
| Normalization | RMSNorm | 15% faster than LayerNorm |
| Activation | SwiGLU | Gated: learns what to pass/block |
| Norm placement | Pre-norm | Stable gradients at 100+ layers |
| Optimizer | AdamW | Better generalization than Adam |
| Tokenizer | BPE | Handles any text including emoji |
| Param sharing | Weight tying | Saves 30% params |
| Training | Mixed precision (fp16) | 2× speed, half memory |

## Key Code Patterns

### 1. BPE Tokenizer

```python
import tiktoken

# Load GPT-4 BPE tokenizer (same family as used by OpenAI)
enc = tiktoken.get_encoding("cl100k_base")

# Encode text to token IDs
tokens = enc.encode("Hello, how are you?")
# → [9906, 11, 1268, 527, 499, 30]

# Decode back to text
text = enc.decode(tokens)
# → "Hello, how are you?"

vocab_size = enc.n_vocab  # 100,277 for cl100k_base
```

### 2. Token Embeddings

```python
import torch
import torch.nn as nn

# Embedding table: maps token ID → dense vector
# vocab_size tokens, each represented as d_model floats
embedding = nn.Embedding(vocab_size, d_model)

# token_ids shape: (batch_size, seq_len)
token_ids = torch.tensor([[9906, 11, 1268]])  # (1, 3)

# x shape: (batch_size, seq_len, d_model)
x = embedding(token_ids)  # (1, 3, 768)
```

### 3. RoPE (Rotary Position Encoding)

```python
def precompute_rope_freqs(d_model: int, max_seq_len: int, theta: float = 10000.0):
    """
    Precompute rotation frequencies for RoPE.
    LLaMA rotates Q and K vectors by position-dependent angles
    so attention scores encode relative distance automatically.
    """
    # Frequency for each pair of dimensions: theta^(-2i/d)
    # Shape: (d_model // 2,)
    freqs = 1.0 / (theta ** (torch.arange(0, d_model, 2).float() / d_model))
    
    # Position indices: 0, 1, 2, ..., max_seq_len-1
    positions = torch.arange(max_seq_len)
    
    # Outer product: each position × each frequency
    # Shape: (max_seq_len, d_model // 2)
    freqs = torch.outer(positions, freqs)
    
    # Convert to complex numbers for efficient rotation
    # Shape: (max_seq_len, d_model // 2)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis


def apply_rope(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """
    Rotate query or key vectors by their position's angle.
    x shape: (batch, seq_len, num_heads, head_dim)
    """
    # Reshape to pairs of floats, treat as complex numbers
    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
    
    # Broadcast freqs_cis: (1, seq_len, 1, head_dim//2)
    freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(2)
    
    # Multiply = rotate in complex plane
    x_rotated = x_complex * freqs_cis
    
    # Back to real numbers
    return torch.view_as_real(x_rotated).flatten(3).type_as(x)
```

### 4. Multi-Head Attention

```python
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads  # e.g. 768 // 12 = 64
        self.scale = self.head_dim ** -0.5    # 1/√d_k: prevents softmax saturation
        
        # Single projection for Q, K, V (more efficient than 3 separate)
        self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor, 
                mask: torch.Tensor = None):
        B, T, C = x.shape  # batch, seq_len, d_model
        
        # Step 1: Project to Q, K, V
        qkv = self.qkv_proj(x)  # (B, T, 3*C)
        q, k, v = qkv.split(C, dim=2)  # each (B, T, C)
        
        # Step 2: Reshape into heads
        # (B, T, C) → (B, T, num_heads, head_dim) → (B, num_heads, T, head_dim)
        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Step 3: Apply RoPE to Q and K (not V)
        q = apply_rope(q.transpose(1, 2), freqs_cis).transpose(1, 2)
        k = apply_rope(k.transpose(1, 2), freqs_cis).transpose(1, 2)
        
        # Step 4: Scaled dot-product attention scores
        # (B, heads, T, head_dim) @ (B, heads, head_dim, T) → (B, heads, T, T)
        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        
        # Step 5: Causal mask — token i cannot attend to token j > i
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # Step 6: Softmax over last dim → attention weights
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # Step 7: Weighted sum of values
        # (B, heads, T, T) @ (B, heads, T, head_dim) → (B, heads, T, head_dim)
        out = torch.matmul(attn_weights, v)
        
        # Step 8: Concatenate heads and project
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(out)
```

### 5. RMSNorm

```python
class RMSNorm(nn.Module):
    """
    Root Mean Square Normalization.
    Faster than LayerNorm: no mean subtraction, just scale by RMS.
    Used by LLaMA, Mistral, Gemma instead of LayerNorm.
    """
    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))  # learned scale
    
    def forward(self, x: torch.Tensor) ->