GPT-2 Style Transformer — Built from Scratch

I built this transformer from scratch in PyTorch because I genuinely love NLP and wanted a deep understanding of every component, not just an API wrapper. No HuggingFace, no shortcuts — every attention head, every residual connection, every layer norm written by hand.

Note on model quality: this model is intentionally small (60-dim embeddings, 4 heads, 4 blocks) and trained on a short story. It does not produce fluent text. That is not the point — the goal was to understand why multi-head attention works, what causal masking actually does to the weight matrix, and how residual connections keep gradients healthy across depth. Use a real checkpoint if you want good generation; use this codebase if you want to understand transformers.

Architecture

The model follows the GPT-2 pre-LN design:

token_ids
   │
   ├── Word Embedding  ─┐
   └── Position Embedding ─┘ → sum → Dropout
                                        │
                              ┌─────────▼─────────┐
                              │  Transformer Block │  × N
                              │                   │
                              │  LayerNorm         │
                              │  MultiHead Causal  │
                              │   Attention        │
                              │  Dropout + Residual│
                              │  LayerNorm         │
                              │  FeedForward (GELU)│
                              │  Residual          │
                              └─────────┬─────────┘
                                        │
                              Linear Projection
                                        │
                                     Logits

File Structure

├── config.py         # model hyperparameters
├── model.py          # LayerNorm, MultiHeadAttention, FeedForward, TransformerBlock, GPTModel
├── dataset.py        # TextDataset — sliding window tokenization
├── train.py          # train, evaluate, generate + __main__ entry point
├── api.py            # FastAPI inference server
├── Dockerfile
├── docker-compose.yml
└── requirements.txt

Key Implementation Details

• Causal masking: upper-triangular mask built with torch.full(-inf) + triu to avoid the 0 × -inf = nan trap
• Pre-LN: LayerNorm applied before each sub-layer, which is more stable during training than post-LN
• Scaled dot-product attention: scores divided by sqrt(head_dim), which keeps variance stable regardless of embedding size
• Weight tying: the projection head has no bias, consistent with GPT-2
• No softmax before loss: CrossEntropyLoss handles log-softmax internally — passing probabilities in would compute it twice
• Evaluation: always model.eval() + torch.no_grad() together — one disables dropout, the other skips building the computation graph

What this project teaches you

• How multi-head attention physically works: splitting [b, s, d] into [b, h, s, d/h] and why that lets each head specialise on different relationships
• Why causal masking is an addition (not a multiplication) and what numerical traps look like (0 × -inf = nan)
• How residual connections keep gradients non-zero across many layers
• The difference between nn.Parameter, register_buffer, and a plain Python attribute
• Why model.eval() and torch.no_grad() are two separate and independent concerns
• How CrossEntropyLoss expects raw logits and handles softmax internally

Running locally

pip install -r requirements.txt
uvicorn api:app --reload

Send a generation request:

curl -X POST http://127.0.0.1:8000/generate \
  -H "Content-Type: application/json" \
  -d '{"prompt": "I love", "max_new_tokens": 10}'

Interactive API docs: http://127.0.0.1:8000/docs

Running with Docker

docker compose up --build

The API will be available at http://localhost:8000.

To use a trained checkpoint, place checkpoint.pt in the project root — the compose file mounts it automatically.

Training

Place your training text at the-verdict.txt (or change the path in training.py), then:

python training.py

Future Work

• RoPE (Rotary Position Embeddings): replace learned absolute position embeddings with rotary encodings, which generalise to sequence lengths unseen during training and give the model a relative sense of position for free
• KV Cache: cache past key/value projections during autoregressive generation so each new token only runs one attention pass through the new position — reduces inference cost from O(n²) to O(n)
• Sampling strategies: add top-k and nucleus (top-p) sampling with temperature control instead of greedy argmax, for more natural and varied generation
• Larger training run: scale up embedding_dim, n_blocks, and train on a larger corpus to produce a model that actually generates coherent text