StorySprint
A 30M-parameter GPT trained from scratch to write short children's stories, and a training speedrun: reaching the same validation loss on a single NVIDIA T4 in as little time as possible. The model below runs entirely in your browser.
- Parameters
- 30M
- Architecture
- 6 layers, 256-token context
- Demo model training
- 901M tokens
- Inference
- In-browser, ONNX
Try the model
Examples
Sampling settings
Loading model...
The model could not be loaded.
Speedrun results
5.64x faster: from 54.75 to 9.70 minutes of training on one Kaggle T4 to reach validation loss 1.95 (baseline and final, 3 seeds each, interleaved in the same session). Test loss of the final model, evaluated once: 1.9225.
Each step below was measured against a control run of the previous configuration in the same session, because GPU speed varies by up to about 12% between sessions.
| Configuration | Time to target | Step speedup | Total |
|---|---|---|---|
| Baseline: fp32, hand-written attention, AdamW | 51.68 min | - | 1.00x |
| + fp16 autocast and GradScaler | 23.34 min | 2.37x | 2.37x |
| + vocabulary padded to 50,304 | 16.32 min | 1.29x | 3.05x |
| + torch.compile | 13.54 min | 1.31x | 3.99x |
| + SDPA attention, fused AdamW | 13.17 min | 1.06x | 4.23x |
| + LR decay ending at the target | 11.65 min | 1.13x | 4.78x |
| + peak learning rate 1e-3 | 10.47 min | 1.11x | 5.33x |
| + peak learning rate 1.5e-3 (final) | 9.87 min | 1.06x | 5.65x |
Each time is from the session that tested that step, so rows come from different machines; step speedups compare against the same-session control and the total is their product. The first four steps make the GPU faster; the learning-rate steps make the model need fewer tokens (45M down to 35M). The demo model is separate: the same speedups, trained for 901M tokens on text with repaired punctuation, instead of stopping at the target.
Method
Model and data
A GPT-2 style decoder with 6 layers, 6 heads and 384 dimensions, trained from random weights on TinyStories with the GPT-2 tokenizer.
Measurement
Every step is timed between GPU synchronisations; validation runs off the clock on 50 fixed batches. Every run is logged with its git commit, config, seed and GPU.
In the browser
The model is exported to ONNX, compressed to 50 MB (int8 weights, fp16 embedding) and run with onnxruntime-web. Prompts are processed on your device and never sent to a server.