File size: 3,683 Bytes
676ffeb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | ---
license: cdla-sharing-1.0
tags:
- text-generation
- gpt
- from-scratch
- pytorch
- mps
- tinystories
datasets:
- roneneldan/TinyStories
pipeline_tag: text-generation
library_name: pytorch
language:
- en
---
# mini-LLM — a 30M-parameter transformer trained overnight on a laptop
A decoder-only GPT written from scratch in plain PyTorch. No `transformers`, no `accelerate`, no
`fastai` — the model, the tokenizer training and the training loop are all in the repository.
**Code: https://github.com/vous99/mini-llm**
Trained on one MacBook Pro (M5 Pro, 24 GB, MPS) in a single overnight run.
## Results
| Metric | Value |
|---|---|
| **Best val loss** | **1.1527** (perplexity 3.17) |
| Loss at initialisation | 9.01 = ln(8192), the uniform-over-vocabulary baseline |
| Steps | 7,561 |
| Tokens seen | 0.74B (1.4 epochs of TinyStories V2) |
| Training time | ~9 hours of compute at 23–25K tokens/s |
| Parameters | 29,893,120 (25.7M non-embedding) |
## Architecture
Modern rather than 2017: what separates this from the original transformer paper is where most of
the interest lies.
| | |
|---|---|
| Layers | 8 |
| Model width | 512 |
| Heads | 8 (head_dim 64) |
| Context | 512 tokens |
| Vocabulary | 8,192, byte-level BPE trained on the corpus itself |
| Positional encoding | RoPE (rotary), not learned embeddings |
| Normalisation | RMSNorm, pre-norm |
| Feed-forward | SwiGLU, hidden 1408 (8/3 × width, rounded to a multiple of 64) |
| Attention | `F.scaled_dot_product_attention`, causal |
| Embeddings | tied — one matrix serves both input and output |
| Dropout | 0.0 |
The 8,192-token vocabulary is a deliberate choice: GPT-2's 50,257 would put a 25M-parameter
embedding table inside a 30M-parameter model.
## Usage
```bash
git clone https://github.com/vous99/mini-llm && cd mini-llm
pip install -r requirements.txt
python -c "
from huggingface_hub import hf_hub_download
import shutil, os
os.makedirs('ckpt', exist_ok=True); os.makedirs('data', exist_ok=True)
shutil.copy(hf_hub_download('vous99/mini-llm', 'best.pt'), 'ckpt/best.pt')
shutil.copy(hf_hub_download('vous99/mini-llm', 'tokenizer.json'), 'data/tokenizer.json')
"
python sample.py
python sample.py --prompt "Once upon a time, a little robot" --n 3
python chat.py # terminal REPL
python serve.py # browser playground on 127.0.0.1:8890
```
**`tokenizer.json` is required.** It is the BPE trained alongside the model; without it the
weights emit token ids, not text.
## Training details
| | |
|---|---|
| Optimizer | AdamW, lr 6e-4, betas (0.9, 0.95), grad clip 1.0 |
| Weight decay | 0.1 on matrices, 0 on 1-D parameters (norms) |
| LR schedule | linear warmup 200 steps → cosine decay to 10% of peak |
| Batch | 24 × 8 gradient accumulation × 512 tokens = 98,304 tokens per step |
| Precision | bfloat16 autocast |
Train and validation loss stayed within 0.01–0.03 of each other for the whole run — no
overfitting, which is why dropout is 0.
## Limitations
- **TinyStories only.** The corpus is synthetic children's stories using a deliberately small
vocabulary. The model writes fluent, coherent short stories in that register and nothing else.
It has no world knowledge, cannot answer questions and cannot follow instructions.
- **512-token context.**
- **No instruction tuning, no RLHF.** This is a base model in the most literal sense.
- **Not safe for production.** It is a study artifact for understanding how a transformer trains.
## Checkpoint contents
`best.pt` is a `torch.save` dict with `model`, `cfg` (the `GPTConfig`), `iter`, `best_val` and
`val_loss`. Load it with `sample.py` from the GitHub repository.
|