GPT from Scratch

A decoder-only GPT-style transformer implemented from scratch in NumPy, including a custom reverse-mode automatic differentiation engine.

The project implements the complete language-model training pipeline without PyTorch, TensorFlow, JAX, or another automatic-differentiation framework:

The goal is to expose the mechanics behind transformer training rather than hide them behind a deep-learning framework. Core components—including tensor operations, gradient propagation, neural-network layers, causal multi-head self-attention, optimization, and autoregressive generation—are implemented directly on top of NumPy.

Highlights

• Reverse-mode automatic differentiation with a dynamic computation graph and reverse topological traversal.
• Neural-network primitives including linear layers, embeddings, LayerNorm, Dropout, GELU, and cross-entropy loss.
• Decoder-only transformer with causal multi-head self-attention, residual connections, and pre-LayerNorm blocks.
• Adam optimizer with bias-corrected first and second moments.
• End-to-end language modeling from character tokenization and batching through training and autoregressive generation.
• Numerical gradient checks that compare autograd results with finite-difference approximations.

Architecture

Project structure

gpt-from-scratch/
├── tinygpt/
│   ├── tensor.py              # reverse-mode autograd engine
│   ├── nn.py                  # layers, loss, and Adam optimizer
│   ├── transformer.py         # attention, transformer blocks, GPT model
│   └── dataset.py             # character tokenizer and batching
├── examples/shakespeare/      # larger example on real text
│   ├── corpus.txt
│   ├── train_shakespeare.py
│   └── generate_shakespeare.py
├── tests/
│   ├── test_autograd.py       # finite-difference gradient validation
│   └── test_sanity.py         # model and training smoke tests
├── train.py                   # trains the small built-in example
├── generate.py                # samples from saved model weights
└── requirements.txt

Automatic differentiation

tinygpt/tensor.py implements a small Tensor abstraction that records the operation that produced each value and the parent tensors on which it depends. Calling backward() constructs a topological ordering of the computation graph and traverses it in reverse, applying each operation's local derivative and accumulating gradients through the chain rule.

Conceptually, the training path is:

Broadcasting, matrix multiplication, reductions, reshaping, slicing, and nonlinear operations all participate in the graph. The tests include finite-difference checks to validate representative analytical gradients numerically.

Transformer implementation

tinygpt/transformer.py builds a decoder-only language model from the primitives above. Each transformer block uses pre-normalization and two residual paths:

Causal masking prevents each position from attending to future tokens. Multi-head attention projects queries, keys, and values independently, splits the representation across attention heads, applies scaled dot-product attention, and merges the heads before the output projection.

Quickstart

Install the only runtime dependency:

pip install -r requirements.txt

Train the small built-in example and generate text:

python train.py
python generate.py

train.py writes tinygpt_weights.npz, which generate.py loads for sampling.

Example output after training on the small corpus:

--- generated ---

the cat sat on the mat
the dog sat on the rug
the cat chased the mouse

For a larger example using Shakespeare text:

cd examples/shakespeare
python train_shakespeare.py
python generate_shakespeare.py

Validation

Run the included checks with:

python -m tests.test_autograd
python -m tests.test_sanity

The autograd tests compare derivatives produced by the custom engine against central finite differences,

df/dx ≈ [f(x + ε) - f(x - ε)] / (2ε),

for representative scalar, broadcast, and matrix operations. The sanity tests verify model output dimensions and a complete forward/backward/optimizer step.

Design scope

This implementation prioritizes transparency and correctness over training performance. It is intended for studying the mechanics of automatic differentiation and transformer language models rather than large-scale model training.

Current tradeoffs include:

• NumPy/CPU execution only; no GPU kernels or mixed precision.
• The complete computation graph is retained until the backward pass finishes.
• float32 is used throughout the model.
• No gradient checkpointing or gradient clipping.
• No token-embedding / language-model-head weight tying.
• No learning-rate warmup; the example training script uses cosine decay.
• The causal mask is applied directly to attention-score data because the mask itself is non-differentiable and requires no gradient.

Shakespeare-scale text is a practical upper bound for experimentation with this implementation; larger workloads are better suited to optimized frameworks such as PyTorch or JAX.

Possible extensions

• Save optimizer state and resume training from checkpoints.
• Store model hyperparameters alongside serialized weights.
• Add learning-rate warmup and gradient clipping.
• Add embedding / LM-head weight tying.
• Support masked losses for instruction-style prompt/response fine-tuning.
• Benchmark selected operations against equivalent framework implementations.

License

This project is licensed under the MIT License. See LICENSE for details.