Skip to content

Repository files navigation

ZeroToGPT

We built a GPT from scratch and pretrained it on a 7-billion-token corpus, as undergraduates, without a research lab or massive compute.

Open In Colab

134M parameters, every component written out in PyTorch, pretrained on a 7B-token FineWeb-Edu corpus using a single free-tier GPU. No nn.Transformer, no HuggingFace model class. Open the notebook and continue the run yourself.

Overview

Parameters 134,077,440 (95,283,456 non-embedding)
Model 8-block decoder-only transformer, 768 wide, 12 heads
Dataset 7B tokens, FineWeb-Edu CC-MAIN-2024-10
Hardware one Tesla P100 16 GB, free Kaggle session
Held-out perplexity 36.12 on FineWeb-Edu, context 128
Throughput 10,200 tok/s, 31.7% MFU, 12.24 GB peak of 16
Weights umerateeq/zerotogpt-134m, 538 MB
Write-up ZeroToGPT on Level Up Coding
Post LinkedIn

Architecture

Type decoder-only transformer Blocks 8
Residual width 768 Attention heads 12, head dim 64
Attention multi-head causal self-attention Masking causal, upper triangular
Feed-forward 768 to 3072 to 768, ReLU Normalization pre-norm LayerNorm, eps = 1e-5
Token embedding nn.Embedding(50257, 768) Positional encoding learned absolute
Tokenizer GPT-2 BPE (tiktoken) Vocabulary 50,257
Context 128 Dropout 0.1
Output head untied, nn.Linear(768, 50257, bias=False) Sampling greedy, temperature with top-k

Every module is implemented directly: multi-head causal self-attention, the causal mask, LayerNorm, the feed-forward block, the residual wiring, the sampler and the training loop.


Training

Optimizer AdamW Learning rate 4e-4
Weight decay 0.1 Grad clip global norm 1.0
Precision fp16 + GradScaler Sequence length 128
Batch 32 x 128 = 4,096 tok/step Tokens seen 1.23B of a 7B corpus
Throughput 10,200 tok/s Peak GPU memory 12.24 GB of 16
Achieved 5.93 TFLOP/s MFU 31.7% of the P100's 18.7 TFLOP/s fp16 peak

Data pipeline

Stage Implementation
Source FineWeb-Edu CC-MAIN-2024-10, streamed from Hugging Face
Tokenizer GPT-2 BPE via tiktoken, encode_ordinary
Storage pre-tokenized into a flat .bin file, written through np.memmap
Dtype uint16, since GPT-2's highest token id is 50256 and fits in two bytes
Batching memmap reopened per batch, pinned, copied to GPU asynchronously

Results

Dataset Perplexity Context
Held-out FineWeb-Edu 36.12 128
TinyStories 35.41 128
WikiText-2 184.96 128

Scored on full raw test sets with non-overlapping windows. Every number is reproducible in one click, see Evaluate the pretrained model below.


Continue the pretraining run

Open In Colab

Open the notebook in Colab and hit Run all. Nothing to install. It pulls the pre-tokenized train.bin and validation.bin and the released weights.pth, so the run picks up pretraining from the published checkpoint and plots the loss curve as it goes. Skip the "Load my weights" cell to train from scratch instead.

This is a base model: it completes text.


Evaluate the pretrained model

Open In Colab

Hit Run all. It fetches the released checkpoint and the held-out split, scores every dataset in the results table above, and prints each number next to the command that produced it. Nothing to install, and it verifies the numbers in this README rather than asking you to take them on trust.

To run the scorer directly instead:

python evaluate.py --ckpt weights8b_300epoch.pth --all --data-bin validation.bin

evaluate.py also scores any single dataset on its own, and runs GPT-2-small through the identical scoring function as a calibration baseline. Recorded output: RESULTS.md.


Credits

With thanks to:

Data: FineWeb-Edu.

About

A 134M-parameter GPT written from scratch in PyTorch, no nn.Transformer and no HuggingFace, pretrained on a 7B tokens corpus of FineWeb-Edu, where it consumed 1.2B tokens on one free Kaggle P100. Held-out perplexity 38.89 against a Chinchilla prediction of 38.70.

Topics

Resources

Stars

6 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages