nepoleg/baby-stories

BabyStories GPT Model

Self-made and home-trained GPT model capable of producing small child stories. Inspired by

• "makemore" video curse by A. Karpathy
• "Attention Is All You Need" by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, and Illia Polosukhin
• "How Small Can Language Models Be and Still Speak Coherent English?" by Ronen Eldan and Yuanzhi Li of Microsoft Research, published in 2023 as arXiv:2305.07759

Features

• BPE Tokenizer
• Transformers with Multi-Head Attention
• Masking of training data based on the beginning of the story
• Uses bf16 for training to be able to train on a laptop

Sample output

Once upon a time, there was a little boy named Tim. Tim loved to play outside with his friends. One day, Tim and his friends went to the park to play. At the park, Tim saw a big tree. He wanted to climb it. His friend, Sam, said, "Be careful, Tim!" Tim started to climb the tree. He went up, up, up, until he was very high. When Tim got to the top, he saw a big bird. The bird said, "Hello, Tim!" Tim was surprised. He did not know birds could talk. The bird told Tim that it was a magic bird. The bird said, "I can make one wish come true for you." Tim thought for a moment and said, "I wish for a big ice cream cone!" The magic bird made a big ice cream cone appear. Tim and his friends were so happy. They all shared the ice cream and had a great day at the park.

Last Training summary:

Run 8: story masking, GELU, cosine to 0, 25.3M params

• Config: same as run 7 (vocab 10000, context 512, n_embd 384, n_head 6, n_layer 12, n_ff 1536, ~25.3M params, bf16)
• Story masking: attention stays inside one story. Each token sees only earlier tokens after the last <|endoftext|> in the window (attn_mask instead of is_causal)
• ReLU → GELU in the feed-forward
• Cosine LR to 0 (was 1e-4)
• Val set fixed: the first 640 windows of the valid file (20 batches × 32, 328k tokens), same each eval. Runs 5–7 used 20 random batches, so the comparison is close but not exact
• Paused at step 46k, resumed from best.pt with --proceed (model, optimizer and scheduler state)
• Val loss: 1.59 (10k) → 1.42 (20k) → 1.34 (30k) → 1.30 (40k) → 1.26 (50k) → 1.23 (60k) → 1.21 (70k) → 1.19 (80k) → 1.18 (90k) → 1.1785 (98k, best) → 1.1785 (100k)
• Averages over 10 evals: 71–80k 1.199, 81–90k 1.185, 91–100k 1.179. No plateau: loss kept falling until LR reached ~0
• Train (~1.16) ≈ val, ~3.0 epochs: still no overfitting
• vs run 7: −0.13 per token (−0.03 per char) with the same params and data. The biggest gain since the switch to BPE
• Checkpoint: best.pt (step 98k)

Source code can be found on GitHub