GPT-2 small, pretrained from scratch on FineWeb-Edu

GPT-2 small (124M parameters) trained from random initialisation on the FineWeb-Edu sample-10BT sample, with the model and training loop written directly in PyTorch. Training details and plots: pretraining write-up.

Results

Files

• gpt2_step_018890.safetensors: the final weights.
• config.json: the training config.
• checkpoints/step_*.pt: the last three resume checkpoints, with optimizer and RNG state, for continuing training with pretraining/train.py --resume.

Loading

These weights use the GPT2 class from gpt2-repro, not Hugging Face Transformers. Clone gpt2-repro and install its dependencies with uv, then run the code below from the repository root:

git clone https://github.com/rahuliyer/gpt2-repro
cd gpt2-repro
uv sync

import tiktoken
from huggingface_hub import hf_hub_download

from inference import generate_text
from model import GPT2, GPT2SmallConfig

path = hf_hub_download("rahuliyer/gpt2-small-fineweb-edu", "gpt2_step_018890.safetensors")
model = GPT2.load_from_checkpoint(GPT2SmallConfig(), path)

tokenizer = tiktoken.get_encoding("gpt2")
print(generate_text(model, tokenizer, "The history of the printing press", max_length=100))

Limitations

A base language model: it continues text and does not follow instructions. It was trained on educational web text in English only and can produce factually wrong or biased text.