rishabhguptajs/tinystories-10m-mlx

slm-lab TinyStories 10M

A 10M-parameter GPT trained from scratch in MLX on a fanless MacBook Air M4, in 2h40m. Own byte-level BPE tokenizer, own architecture code, own training loop. No pretrained weights were used.

By Rishabh Gupta.

Results

[image]

Bits per byte (lower is better) on the first 1,000 stories of the TinyStoriesV2-GPT4 validation split. Each model is conditioned on its own BOS token and scored on the story text only. Bits per byte is used because the models have different tokenizers.

Caveats: the TinyStories-* models were trained on TinyStories v1 (GPT-3.5 + GPT-4 stories). This model was trained on the v2 GPT-4 split that the evaluation uses, which is a home advantage. stories15M was trained for much longer on datacenter GPUs.

Model

• —6 layers, width 320, 5 heads of 64, context 256, vocab 8192 (ids 0-8190 are bytes + BPE merges, id 8191 is <|endoftext|>)
• —RoPE, RMSNorm (pre-norm), QK-norm, GELU MLP (4x), tied input/output embeddings, zero-init residual projections

Training

• —Data: TinyStoriesV2-GPT4 train (539.6M tokens with this tokenizer, 4.06 bytes/token). 164M tokens seen (10,000 steps x 64 x 256).
• —Optimizer: Muon (lr 0.02, Nesterov, Newton-Schulz in bf16) for hidden matrices; AdamW (lr 3e-3, betas 0.9/0.95, wd 0.1 on the embedding) for the embedding and norm gains. Grad clip 1.0.
• —Schedule: warmup 100 steps, flat, then linear decay to 0 over the last 20% (WSD).
• —bf16 compute with fp32 master weights. ~17-20k tokens/s sustained on the M4 Air.

Short equal-token ablations (8.2M tokens each) that led to this recipe, in bits/byte: AdamW + ctx 512 baseline 0.974; ctx 256: 0.955; + AdamW lr 3e-3: 0.952; Muon: 0.689;

• —QK-norm and zero-init: 0.658; + WSD: 0.655. These were single-seed short runs, so treat them as early signal.

Usage (Apple Silicon)

pip install mlx numpy
python generate.py --prompt "Once upon a time" --n 2

Sample:

Once upon a time, there was a little dog named Spot. Spot loved to sit and watch the birds outside. One day, Spot's owner, Tom, took him to the park. [...] Spot was curious and wanted to see what was inside. He climbed up to the box and tried to open it. But the box was shut very tight.

Files

Reproduce

cd training
pip install mlx numpy
python prepare.py          # downloads TinyStoriesV2-GPT4, trains the tokenizer, writes data/*.bin (~4 min)
python train.py --out run1 # 10k steps, ~2h40m on an M4 MacBook Air
python compare.py --runs run1

Limitations

• —It writes simple children's stories in English, and nothing else. Stories drift logically and have occasional grammar slips.
• —It was trained on TinyStories, which was generated by GPT-3.5/GPT-4. This model is a benchmark and learning artifact, not a general-purpose model.
• —The tokenizer is English-centric (it encodes Devanagari at 1 token per byte).

License

Weights are released under CDLA-Sharing-1.0, matching the TinyStories dataset license.