GPT from Scratch

Build a GPT model from scratch and run it on your MacBook — tokenization to inference, one step at a time.

What This Is

A complete implementation of a decoder-only transformer (the GPT architecture) in ~500 lines of pure PyTorch. No HuggingFace, no tiktoken, no pre-built ML libraries — every component is written from scratch so you can see exactly how it works.

This is the same architecture behind GPT-2, GPT-3, and GPT-4. We just run it smaller: 3.3M parameters instead of 175B, trained on a laptop in minutes instead of a datacenter in months.

Quick Start

# Install the only dependency
pip install torch

# Run the full pipeline: download → tokenize → train → finetune → generate
python run.py

That's it. The pipeline downloads training data, trains a BPE tokenizer, pretrains the model, fine-tunes it on instructions, and generates text — all in one command, ~5 minutes on an M-series Mac.

How to Run Step by Step

# 1. Download training corpus (Project Gutenberg books, ~5.7MB)
python download_data.py

# 2. Train BPE tokenizer on the corpus
python tokenizer.py

# 3. Pretrain the model (next-token prediction)
python train.py

# 4. Fine-tune on instruction-response pairs
python finetune.py

# 5. Interactive text generation
python generate.py

# 6. Run inference test suite (quality, speed, KV cache benchmark)
python inference.py

Project Structure

gpt-from-scratch/
├── model.py                   # GPT architecture: embeddings, attention, transformer blocks
├── tokenizer.py               # BPE tokenizer: train, encode, decode
├── train.py                   # Pretraining: next-token prediction loop
├── finetune.py                # Fine-tuning: instruction-response pairs
├── generate.py                # Inference: generation with KV cache + sampling
├── inference.py               # Test suite: quality, determinism, speed, temperature
├── download_data.py           # Data: downloads public domain books
├── run.py                     # Full pipeline with detailed debug logs
├── train_distributed.py       # DDP training (Post 7)
├── train_fsdp.py              # FSDP training (Post 8)
├── train_tp.py                # Tensor Parallel training (Post 9)
├── inference_disaggregated.py # Disaggregated inference (Post 10)
├── test_distributed.py        # Test suite for all distributed modes
├── data/
│   └── instructions.json      # Sample instruction-response pairs for fine-tuning
├── checkpoints/               # Saved model weights (generated by training)
├── LICENSE                    # MIT
└── .gitignore

How It Works

Architecture

Input text
    ↓
[BPE Tokenizer] → token IDs (integers)
    ↓
[Token Embedding] + [Position Embedding] → vectors
    ↓
[Transformer Block] × N layers
    ├── LayerNorm → Causal Multi-Head Self-Attention → Residual
    └── LayerNorm → Feed-Forward (GELU) → Residual
    ↓
[Final LayerNorm] → [Linear Head] → logits
    ↓
[Sampling (temperature + top-k + top-p)] → next token

Key Components

BPE Tokenizer (tokenizer.py)

• Starts with 256 byte-level tokens
• Iteratively merges the most frequent adjacent pair into a new token
• Produces a fixed vocabulary (512 tokens in our case)
• Handles any UTF-8 text — no unknown tokens

Causal Self-Attention (model.py)

• Each token attends to all previous tokens (never future ones)
• Uses a triangular mask to enforce left-to-right causality
• Multi-head: 4 independent attention patterns computed in parallel
• Scaled dot-product: softmax(QK^T / √d) × V

Pretraining (train.py)

• Objective: predict the next token (cross-entropy loss)
• Optimizer: AdamW with weight decay 0.1
• Schedule: linear warmup → cosine decay
• Gradient clipping at norm 1.0
• Gradient accumulation for effective larger batch sizes

Fine-tuning (finetune.py)

• Same loss function, different data format (instruction → response)
• Padding tokens masked with ignore_index=-100 so loss ignores them
• Lower learning rate (1e-4) and weight decay (0.01) to preserve pretrained knowledge

KV Cache (model.py + generate.py)

• During generation, past tokens' Keys and Values never change
• Cache them to avoid recomputation: O(n²) → O(n) per generation
• Two phases: prefill (process full prompt) → decode (one token at a time)

Sampling (generate.py)

• Temperature: scales logits before softmax (low=focused, high=creative)
• Top-k: only consider the k most probable tokens
• Top-p (nucleus): keep tokens until cumulative probability exceeds p
• Defaults: temperature=0.8, top_k=40, top_p=0.9

Model Configuration

vocab_size:      512        tokens in the BPE vocabulary
context_length:  128        maximum sequence length
n_layers:        4          transformer blocks
n_heads:         4          parallel attention heads per layer
embed_dim:       256        dimension of all representations
head_dim:        64         embed_dim / n_heads
parameters:      3.3M       total trainable parameters

Training Details

• Corpus: ~5.7MB of public domain English text (Shakespeare + Gutenberg novels)
• Tokens: ~2.7 million after BPE encoding
• Device: Apple Silicon MPS (falls back to CPU if unavailable)
• Training time: ~1 minute for 2000 steps
• Final loss: ~3.4 (from 6.3 random baseline)

Verbose Mode

run.py includes detailed debug logging that shows exactly what happens at each stage:

• Tokenization: compression ratio, example encodings
• Architecture: layer breakdown, parameter counts, forward pass trace
• Training: loss curves, learning rate schedule, gradient norms, tokens/sec
• Generation: token-by-token sampling with top-3 candidates and probabilities

This makes it useful as a learning tool — you can see the internals of every step.

Distributed Training & Inference

Same model, same math — scaled to multiple GPUs. All scripts work in single-device mode (MacBook) and distributed mode (multi-GPU/multi-node).

DDP (Distributed Data Parallel) — Post 7

# Single device (same as train.py)
python train_distributed.py

# Multi-GPU (4 GPUs)
torchrun --nproc_per_node=4 train_distributed.py

# Multi-node (2 nodes × 4 GPUs)
torchrun --nnodes=2 --nproc_per_node=4 --master_addr=node0 train_distributed.py

# Test locally without GPUs
torchrun --nproc_per_node=2 train_distributed.py --backend gloo

FSDP (Fully Sharded Data Parallel) — Post 8

# Multi-GPU with full sharding
torchrun --nproc_per_node=4 train_fsdp.py --sharding_strategy full

# HYBRID_SHARD: FSDP within node, DDP across nodes
torchrun --nnodes=4 --nproc_per_node=8 train_fsdp.py --sharding_strategy hybrid

# Test locally (falls back to DDP since FSDP requires CUDA)
torchrun --nproc_per_node=2 train_fsdp.py --backend gloo

Tensor Parallelism — Post 9

# TP=4 within one node
torchrun --nproc_per_node=4 train_tp.py --tp_size=4

# 3D parallelism: TP=8 within node, DP=4 across nodes
torchrun --nnodes=4 --nproc_per_node=8 train_tp.py --tp_size=8 --dp_size=4

# Test locally with TP=2
torchrun --nproc_per_node=2 train_tp.py --tp_size=2 --backend gloo

Disaggregated Inference — Post 10

# Simulate prefill/decode separation (single device)
python inference_disaggregated.py

# Benchmark colocated vs disaggregated
python inference_disaggregated.py --mode benchmark

# With KV cache quantization (int8, 2× transfer reduction)
python inference_disaggregated.py --quantize_kv

Testing Distributed Code Locally

All distributed training can be tested on a MacBook without GPUs:

# Run the full test suite
python test_distributed.py

# Test individual modes
python test_distributed.py --test ddp
python test_distributed.py --test fsdp
python test_distributed.py --test tp
python test_distributed.py --test disaggregated

torchrun with --backend gloo creates multiple CPU processes that communicate via shared memory — identical communication patterns (AllReduce, process groups) to a real multi-GPU cluster.

What This Does NOT Include

Things omitted for clarity (they matter at scale, not for understanding):

• Flash Attention (optimization, same math)
• Rotary Position Embeddings (RoPE — we use learned positional)
• Mixture of Experts (MoE)
• RLHF (we do supervised fine-tuning only)
• Proper end-of-sequence handling

Requirements

• Python 3.10+
• PyTorch 2.0+ (the only dependency)
• Apple Silicon Mac recommended (uses MPS backend), works on any CPU

Follow Along

This repo is part of a post series on thelastprogrammers.com — each post walks through one step with full explanations of the architecture, critical code, and design decisions.

License

MIT