mini-10m

Simple decoder-only Transformer (~10.5M params) with a 16k max context window, trained on the first 10,000 rows of:

karpathy/tinystories-gpt4-clean

Architecture

model.py uses:

Byte-level vocabulary (vocab_size=256)
7 Transformer blocks
d_model=352, n_heads=8
RoPE positional encoding (max_seq_len=16384)
RMSNorm + GELU MLP
Tied input/output embeddings

Parameter count is approximately 10.5M.

Setup

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Quick dry run

python train.py --dry-run

Train on 10k TinyStories rows

python train.py \
  --dataset karpathy/tinystories-gpt4-clean \
  --num-rows 10000 \
  --cache-dir .cache/huggingface \
  --max-seq-len 16384 \
  --train-seq-len 2048 \
  --batch-size 2 \
  --grad-accum 8 \
  --max-steps 500

Notes:

--max-seq-len 16384 sets model context capability.
--train-seq-len can be smaller (for memory); this is common for small runs.
Data loading uses Hugging Face streaming, so only the requested rows are consumed.
Checkpoints and config are written to runs/tiny10m/.

Inference

Use the latest checkpoint in a run directory:

python infer.py \
  --run-dir runs/tiny10m \
  --prompt "Once upon a time" \
  --max-new-tokens 200 \
  --temperature 0.8 \
  --top-k 40 \
  --top-p 0.95

Or target an exact checkpoint:

python infer.py --checkpoint runs/tiny10m/step_000500.pt --prompt "The little robot"