LLM-Lab
A from-scratch, config-driven language-model lab — tokenizer → decoder-only transformer → scaling experiments — on a mission to train a genuinely useful medium language model on real web data.
Overview
LLM-Lab builds a modern decoder-only language model from random initialization — no pretrained weights — and treats every stage as a reproducible experiment.
The whole pipeline is implemented by hand: byte-level BPE tokenization, a stateless LSTM baseline and a GPT-style transformer, a shared training harness (AdamW, warmup + cosine LR, checkpoint/resume, mixed precision), TensorBoard logging, and autoregressive generation.
The point isn't one model — it's a clean, reproducible experimental system for understanding how data, architecture, optimization and compute turn random parameters into text. The ambition is a real Medium Language Model: a ~350M base, instruction-tuned and aligned.
At a glance
Motivation
How far can I go?
The real question wasn't whether I could train a language model — it was how far I can go. I wanted to find the constraints: with the resources I have, how far does the loss keep falling, which axis — model size, data, or compute — buys the most, and where do the returns plateau?
So rather than chase a single result, the pipeline is rebuilt from scratch and run as a series of small, deliberate experiments — treating training runs as science and documenting every curve and constraint along the way.
Architecture
A config-driven pipeline, end to end
Data. Corpus loaders (single-document Shakespeare; multi-document TinyStories and FineWeb-Edu) with <|endoftext|> document boundaries, byte-level BPE, next-token windows, and a train/val split. Token ids are stored in the smallest lossless dtype (int16), keeping large corpora compact.
Models. A stateless LSTM baseline and a GPT (decoder-only transformer: pre-LN blocks, learned positions, GELU MLP, tied embeddings, causal attention) behind one forward(input_ids) → logits contract, so the same harness trains both.
Training harness (architecture-agnostic): AdamW, linear warmup + cosine decay, gradient clipping, periodic eval, checkpoint/resume (scheduler state included), early stopping, optional AMP, and throughput/compute accounting (tokens, FLOPs, wall-time).
Reproducibility. Everything is config-driven — thin scripts/ over a tested src/llm_lab library — and every artifact (checkpoint, curve, log) is traceable to the config that produced it.
Results
Two architectures, two corpora
Both architectures were validated end-to-end across two corpora before any scaling work began.
| Model | Corpus | Params | Val perplexity |
|---|---|---|---|
| LSTM | Tiny Shakespeare | 1.58M | 25.4 |
| GPT | Tiny Shakespeare | 0.93M | 24.9 |
| LSTM | TinyStories | 2.17M | 10.2 |
| GPT | TinyStories | 2.15M | 5.44 |
- →Scaling the corpus roughly halved perplexity; param-matched, the transformer reached ~half the LSTM's perplexity
- →Small corpora memorized (train ↓, val ↑); more data generalized better
- →Mixed precision was numerically identical (loss matched to 4 decimals) and ~1.7× faster on a T4
Scaling
One axis at a time
With the pipeline validated, the real question was how loss responds to scale. Three sweeps, each changing a single axis — model size, data budget, compute budget. The solid curves are validation loss (left axis); the dashed lines are each run's cosine learning-rate schedule (right axis).
Model size
51.9M → 88.3M params · fixed 164M tokens · 20k steps · TinyStories
- →Bigger models get harder to keep stable — xl3 (88M) dips to ~1.93, spikes +0.10 at step 15,500, then recovers; xl2 descends smoothly
- →It is an instability, not overfitting — validation spikes transiently and recovers rather than rising monotonically
- →Not simply “LR too high”: xl3's LR (6e-4) is already lower than xl2's (8e-4) — the larger model is just closer to the edge of stable training
Key finding
Data budget
22M → 164M unique tokens · GPT XL 27.4M · 20k steps
- →More data → lower validation loss (generalization): best val loss improves 2.104 → 1.969 → 1.925 as D grows 22M → 110M, then flattens (d164m: 1.94)
- →Less data → lower training loss (memorization): d22m re-reads the same tokens and overfits — its val loss bottoms at 2.104 then rises to 2.178 while train loss keeps falling
- →The benefit saturates around ~1.5 epochs (110M tokens); beyond that the run is compute-limited, not data-limited
Compute budget
82M → 655M tokens · GPT XL 27.4M · 164M unique tokens
- →More compute → better, monotonically: best val loss 2.038 → 1.936 → 1.865 → 1.811 (perplexity 7.68 → 6.11) as the budget grows 10k → 80k steps
- →Strongly diminishing returns — the gain per doubling shrinks: −0.102 → −0.071 → −0.055 nats
- →A power law fits — 10× the compute buys only ~12% lower loss
- →No overfitting at D = 164M even at 4 epochs; more data pushes the overfitting point further out
The three axes, side by side
| Axis | Range | Result |
|---|---|---|
| Model size | 51.9M → 88.3M params | monotone, flattening; larger models need smaller learning rates |
| Data | 22M → 164M tokens | monotone, saturates ≈ 110M for 27M model |
| Compute | 82M → 655M tokens | monotone, power law α ≈ 0.06 |
Source: TensorBoard logs. Single seed per run; no error bars.
Chinchilla
The scaled-up run — xl4
Applying the lesson — scale data and compute with the model — the next investigation is xl4: a GPT-2-style 355M-parameter model (21 layers, 18 heads, d=1152, 16k vocab, 1024 context) trained compute-optimally at ~20 tokens/parameter — a target of ≈7.1B tokens over ~433k steps.
It is still training — the curves cover ~47% of the budget — but the loss is still falling with no plateau and no instability.
- →Loss is still falling — no plateau yet, and no instability
- →The 1024-context setting carried over cleanly from context/xl3_ctx1024
- →At ~47% it already beats every earlier run's absolute loss, as expected from the much larger parameter and compute budget
Challenges
What building it taught me
Learning rate scales with model size
Larger models need smaller learning rates — the 88M model used a lower LR than the 52M model and sat closer to the edge of stable training. Stability is a first-class concern as models grow.
One axis at a time flattens out
Moving only model size, data, or compute each showed diminishing returns — the Chinchilla lesson: scale data and compute with the model.
You can't compare tokenizers with per-token loss
Bigger tokens carry more information, so loss must be normalized to bits-per-byte — which is what revealed the sweet-spot vocabulary.
Hardware gotchas are real
torch.cuda.is_bf16_supported() returns True on a T4 via software emulation; you must exclude emulation to detect real tensor cores. AMP then gave the expected speedup.
Data engineering matters at scale
Multi-billion-token corpora need int16 storage (and later sharding/mmap); the small-scale pipeline stays deliberately simple until it needs more.