Tiny LLM
GPT-2 Reimplemented From Scratch
A GPT-2-style language model rebuilt line by line to understand how it works
- Loss 6.4 → below 4
- Flash Attention + LoRA
Built with
PyTorchLoRAOpenWebTexttiktoken
Hugging Face Transformers
NumPy
Kaggle
Overview
Large language models are easy to call and hard to understand. During my internship at 3D Smart Factory, I wanted to know exactly what happens inside GPT-2, so I rebuilt a GPT-2-style model from scratch in PyTorch with a small team of fellow interns.
I implemented causal self-attention by hand with a fast path through PyTorch's Flash Attention kernel, pre-LayerNorm decoder blocks, tied input and output embeddings, and my own LoRA layers for parameter-efficient fine-tuning. The model can also load the official GPT-2 weights into the same architecture.
I tokenized about 5 million files from OpenWebText, roughly a quarter of the corpus, with the GPT-2 r50k tokenizer, published that dataset on Kaggle, and trained on Kaggle GPUs. Training and validation loss fell together, from about 6.4 to below 4. It is my most-starred repository, and it was built for understanding, not state-of-the-art perplexity.
The question
Can I reimplement GPT-2 internals from scratch with modern improvements (Flash Attention, pre-LN) and actually train a coherent model?
Key features
- Computes causal self-attention with PyTorch's scaled_dot_product_attention (Flash Attention) and falls back to a manual masked implementation on older PyTorch.
- Uses pre-LayerNorm decoder blocks with a LayerNorm whose bias is optional.
- Adds LoRA adapters to linear and embedding layers, freezing the base weights and merging the low-rank update after training.
- Loads pretrained GPT-2 weights from Hugging Face, resizing the vocabulary for new special tokens and initializing their embeddings from the existing distribution.
- Trains with AdamW (no weight decay on biases and norms), linear warmup then cosine learning-rate decay, gradient accumulation, mixed precision, and gradient clipping.
- Generates text with temperature and top-k sampling.
How it works
Prepare data
About a quarter of OpenWebText is tokenized with the GPT-2 r50k tokenizer and stored as memory-mapped binary token streams for training and validation.
Embed tokens
Token and position embeddings are summed; the token embedding shares its weights with the output layer.
Stack decoder blocks
Each pre-LayerNorm block applies causal self-attention, then an MLP, each with a residual connection.
Train
Next-token cross-entropy with AdamW, warmup and learning-rate decay, gradient accumulation, and mixed precision on Kaggle GPUs.
Fine-tune with LoRA
Base weights are frozen and only small low-rank matrices are trained, then merged back into the weights.
Examples
Tech stack
Model, attention kernels, and training loop
LoRA
Hand-written low-rank adapters for fine-tuning
OpenWebText
Pretraining corpus, about a quarter used
GPT-2 r50k tokenization of the corpus
Loading official pretrained GPT-2 weights
Memory-mapped binary token storage
GPU training and dataset hosting
What I learned
- Flash Attention is a drop-in win on PyTorch ≥ 2.0
- Writing attention by hand first makes the fast kernel a tool, not a black box
Limitations
- Small-scale training; meant for understanding, not SOTA perplexity