StorySprint 30M

A 30M-parameter decoder-only GPT trained from scratch on TinyStories. It writes short, simple children's stories in English. It comes from a training speedrun: how fast a single NVIDIA T4 can train this model to a fixed validation loss. This upload is the demo model, not a speedrun result: it uses the speedups accepted so far (fp16 autocast, vocab padded to 50,304, torch.compile) but trains far longer than the speedrun does (901M tokens instead of about 45M), on training text whose mis-encoded punctuation was repaired, so the stories are better.

Code, rules and full results: github.com/AbdullahMehmoodAwan/StorySprint
Live demo (runs in your browser): huggingface.co/spaces/AbdullahKhan10/storysprint

Speedrun result

Training this model to val loss 1.95 on one T4 went from 54.75 to 9.70 minutes: 5.64x faster (baseline and final, 3 seeds each, interleaved in one Kaggle session). The gains came from fp16 autocast, a vocab padded to 50,304, torch.compile, SDPA attention, fused AdamW, and a shorter learning-rate schedule with a higher peak. Full results and method: see the GitHub repo.

Model


Architecture	GPT-2 style decoder, 6 layers, 6 heads, 384 dims, 256-token context
Parameters	about 30M (input and output embeddings are tied)
Tokenizer	GPT-2 BPE (`tiktoken` "gpt2"), 50,257 tokens, end-of-text id 50256. The output layer has 50,304 rows; the 47 extra ids are padding for GPU speed and are never real tokens
Training	110,000 steps of 32 x 256 tokens (901M tokens), AdamW (lr 6e-4, 200 warm-up steps, cosine decay to 6e-5 ending on the last step), fp16 autocast, torch.compile, 260.2 minutes on one T4
Training data	the whole TinyStories train split, about two passes. About 5.6% of its stories contain punctuation saved with the wrong text encoding (a quote stored as three garbled characters); those were repaired to plain quotes before training, and the rare stories with other broken characters were dropped
File	`model.web.onnx` (49.9 MB): block weights quantised to int8 (onnxruntime dynamic quantisation), shared embedding / output weight stored in fp16

Evaluation

Mean cross-entropy (nats per token) on fixed batches of 32 x 256 tokens from the validation split.

Model	Batches	Val loss
PyTorch fp32, at the end of training	50	1.6338
PyTorch fp32	10	1.5847
This ONNX file	10 (same)	1.5930

The official validation split still contains the garbled punctuation, which this model never saw, so it scores badly on those 5.6% of stories (3.09 vs 1.82 for the previous upload) and that pulls the number above up. On the 94.4% of validation stories without garbled text (20 fixed batches), it is better than the previous upload: 1.4937 vs 1.5340 (the previous upload was trained on 451M tokens of unrepaired text).

Quantising everything to int8 would make the file 30.6 MB but cost 0.0155 loss; keeping the shared embedding in fp16 cuts that to 0.0083, because that weight produces every output logit.

The held-out test split is not used here. It is evaluated once, for the final speedrun model.

How to use (Python, onnxruntime)

import numpy as np
import onnxruntime as ort
import tiktoken
from huggingface_hub import hf_hub_download

path = hf_hub_download("AbdullahKhan10/storysprint-30m", "model.web.onnx")
session = ort.InferenceSession(path)
enc = tiktoken.get_encoding("gpt2")

ids = enc.encode_ordinary("Once upon a time")
for _ in range(200):
    context = np.array([ids[-256:]], dtype=np.int64)  # the model sees at most 256 tokens
    logits = session.run(["logits"], {"input_ids": context})[0][0, -1]
    logits = logits[:50257] / 0.6  # drop the 47 padding ids; temperature 0.6
    top = np.argsort(logits)[-30:]  # top-k 30
    probs = np.exp(logits[top] - logits[top].max())
    next_id = int(np.random.choice(top, p=probs / probs.sum()))
    if next_id == enc.eot_token:
        break
    ids.append(next_id)
print(enc.decode(ids))

Input input_ids: int64, shape (1, T) with 1 <= T <= 256. Output logits: float32, shape (1, 1, 50304): the next-token logits for the last position only (the export skips the vocab projection for every other position, which makes generation about 2.5x faster). There is no KV cache, so each new token reruns the whole context.

Because its training text was repaired, this model does not produce the garbled punctuation found in TinyStories (0 occurrences in 24 test stories, vs 1 story for the previous upload). The demo page still maps any such sequence back to the intended character, as a safety net.

Intended use

A demo of a small language model trained from scratch, and a reference point for training-speed experiments. It is meant for generating simple children's stories for fun and for learning.

Limitations

It only knows the narrow world of TinyStories: simple vocabulary, short stories, few facts. It is not a general assistant and should not be used for anything factual.
Stories are mostly coherent but still sometimes lose track of characters, repeat themselves, or contradict earlier sentences. It is a 30M-parameter model.
The training data is synthetic (written by GPT-3.5 and GPT-4), so the model inherits that data's style and any of its biases.
It is not safety-tuned. It usually writes harmless stories, but nothing guarantees it.
English only. The context is 256 tokens, so it cannot write long stories coherently.

Downloads last month: -; Downloads are not tracked for this model. How to track

AbdullahKhan10
/

storysprint-30m