Train a Small Language Model from Scratch

Lab 16 built the transformer architecture. This lab builds the entire training pipeline around it — real tokenizer, real corpus, LR schedule, validation loop, checkpointing, perplexity tracking, sampling. Everything you'd find in a production training run at 100–1000× this scale.

We're training on TinyStories (Eldan & Li, 2023 — arXiv:2305.07759). The dataset is a breakthrough contribution: ~2M short English stories at a 3–4-year-old's vocabulary level, synthetically generated by GPT-3.5/4. Eldan and Li showed that with this data you can train a 10–30M param model that produces coherent, grammatical, internally-consistent stories — debunking the idea that you need a frontier-scale model for coherent text. The paper kicked off an entire wave of research into "SLM" (Small Language Model) training.

The key insight

At very small scales, model quality is almost entirely bottlenecked by the difficulty of the training data. On web-crawled text, a 10M-param model produces gibberish. On TinyStories (curated simple data), the same 10M-param model produces correct children's stories. Data > parameters at this scale.

This maps to Microsoft's Phi-3 strategy at larger scale: filter/generate high-quality data → train a much smaller model to frontier quality.

References to read while training runs

• TinyStories (Eldan & Li, 2023) — the paper behind our corpus.
• Chinchilla (Hoffmann et al., 2022) — the scaling law paper. Models should be trained with ~20× more tokens than parameters. Our 2M-param model ↔ ~40M tokens of data (we have ~3M for speed, so we're under-trained).
• Karpathy's nanoGPT — the canonical readable reference training loop. Our loop tracks its structure.
• Andrej Karpathy — Let's reproduce GPT-2 (124M) — 4-hour video reproducing GPT-2 from scratch including all the optimizer tricks.

What you'll build

1. Data pipeline — stream TinyStories, tokenize with GPT-2 BPE, pack into efficient training tensors, held-out validation split
2. Model + optimizer — compact GPT (~1.5M params), AdamW, linear warmup → cosine decay LR schedule
3. Training loop — loss tracking, periodic validation, perplexity reporting
4. Generation + checkpointing — before/after samples, save + reload the model

Runtime target on RTX 3060 Ti: ~2 min.

---

Production training does three things in the data layer and we'll mirror all three:

1. Real tokenizer. Char-level toys (Lab 16) train fast but produce terrible text. Real LMs use BPE — we'll reuse GPT-2's pretrained tokenizer (50,257 vocab). Same tokenizer used by GPT-3, Codex, many Llama variants.
2. Packed concatenation. Rather than per-story padding, real training concatenates all documents with an EOS token between them and chops into fixed-length sequences. Much less padding waste, ~2× throughput.
3. Train/val split. Held-out tokens for tracking generalization. ~1% of the corpus is plenty for perplexity tracking.

import torch
import torch.nn as nn
import torch.nn.functional as F
from datasets import load_dataset
from transformers import AutoTokenizer
import time
import math

device = torch.device('cuda')
torch.manual_seed(1337)

# 1. Load a slice of TinyStories. Full set is ~2M stories (~4GB tokenized);
#    we take the first 5000 for fast iteration. Stream to avoid the full download.
print('Loading TinyStories (streaming first 5000 entries)...')
ds = load_dataset('roneneldan/TinyStories', split='train')
stories = []
for i, ex in enumerate(ds):
    if i >= 5000:
        break
    stories.append(ex['text'])

total_chars = sum(len(s) for s in stories)
print(f'Loaded {len(stories):,} stories, {total_chars:,} chars total')

# 2. Tokenize + pack into contiguous token streams
tokenizer = AutoTokenizer.from_pretrained('gpt2')
EOS = tokenizer.eos_token_id
print(f'Vocab: {tokenizer.vocab_size:,}, EOS token id: {EOS}')

# Tokenize each story, concatenate with EOS between them
print('Tokenizing...')
all_ids = []
for s in stories:
    all_ids.extend(tokenizer.encode(s) + [EOS])

all_tokens = torch.tensor(all_ids, dtype=torch.long)

# Held-out val split — last ~1% of the corpus
split = int(0.99 * len(all_tokens))
train_tokens = all_tokens[:split].to(device)
val_tokens = all_tokens[split:].to(device)

print(f'Train: {len(train_tokens):,} tokens')
print(f'Val:   {len(val_tokens):,} tokens')
print(f'Total tokens: {len(all_tokens):,} (vs 20 * params rule-of-thumb for a compute-optimal model)')