sraivante/tinystories-35m

TinyStories 35M (base, chat, smoltalk, chatqa)

A small GPT (decoder-only transformer) trained from scratch on TinyStories on a single Apple M5 laptop, plus a chain of chat fine-tunes.

Variant	Folder	Starts from	What it does
base	base/	scratch	Continues a story from any prompt.
chat	chat/	base	"Tiny", a story bot for kids: small talk, stories on request, QA about a pasted story, polite out-of-scope refusals.
smoltalk	smoltalk/	chat	General assistant-style chat (SFT on smol-smoltalk), keeps storytelling via TinyStories replay.
chatqa	chatqa/	smoltalk	General chat plus answering questions from a passage you give it (RAG-style), multi-turn follow-ups. Best variant for passage QA.

Which one to use: base for raw stories, chat for the kids' story bot, chatqa for chat or answering from a given context.

Model

Architecture	GPT-2 style: 9 layers, 6 heads, 384 dim, pre-LayerNorm, GELU MLP, learned positions, tied input/output embeddings, no biases
Parameters	35.65M (19.3M of them are the tied token embedding)
Context	1024 tokens
Tokenizer	GPT-2 BPE via tiktoken (gpt2), vocab padded to 50,304; `<|endoftext
Weights	safetensors, fp32. lm_head.weight is not stored (tied to wte.weight)

Training

Base (code/train.py, log: logs/train.log)

• Data: 2,119,719 TinyStories training stories (~472M tokens), each followed by <|endoftext|>
• 10,682 steps × 65,536 tokens/step = 700M tokens (~1.5 epochs)
• AdamW (β 0.9/0.95, wd 0.1), lr 1e-3 → 1e-4 cosine, 300 warmup steps, grad clip 1.0, bf16 autocast
• Apple M5 (MPS), ~10.5k tokens/s, ~19 hours
• Final loss: train 1.302, validation 1.340

Chat (code/finetune_chat.py, data builder code/chat_data.py, log: logs/finetune_chat.log)

• Starts from the base model; 3,000 steps, lr 2e-4 → 2e-5
• Synthetic conversations: small talk, story requests (answered with real TinyStories stories), passage QA, simple kid "why" questions, out-of-scope refusals
• Loss only on Assistant turns; every few steps trains on plain TinyStories text to keep storytelling
• Final: chat validation loss 1.048, story loss 1.664 (base: 1.340). The chat model is noticeably weaker as a plain story continuer, so use base for raw story generation.

SmolTalk (code/finetune_sft.py, log: logs/finetune_sft.log)

• Starts from chat; smol-smoltalk (451k conversations, 277M tokens after dropping ones that don't fit 1024 tokens)
• 6,000 steps × 64 sequences: 4,800 chat steps (0.68 epochs) + every 5th step plain TinyStories (replay)
• AdamW, lr 3e-4 → 3e-5 cosine, 200 warmup, fp16 autocast; Kaggle Tesla T4, ~1.1 s/step, ~1h50m
• Final: chat validation loss 2.138 (smol-smoltalk test, 585 convs), story loss 1.251

ChatQA (code/finetune_mix.py, log: logs/finetune_mix.log)

• Starts from smoltalk. Each chat micro-batch is drawn from one source by weight:

  Source	Weight	Train convs	Epochs seen
  UltraChat 200k train_sft	0.65	189,472	1.05
  QuAC (Wikipedia section + follow-ups)	0.15	11,266	4.09
  CoQA (passage + conversational questions)	0.10	7,190	4.27
  smol-smoltalk replay (train[:150000])	0.10	147,198	0.21

• QA conversations put the passage in the first user turn (Context:\n<passage>\n\n<question>); follow-up questions are later turns. Unanswerable QuAC questions (CANNOTANSWER) and CoQA unknown answers become "The context does not say."; QuAC yes/no labels prefix the answer with "Yes." / "No.".

• 6,000 steps (4,800 chat + 1,200 TinyStories replay), lr 2e-4 → 2e-5 cosine, fp16; Kaggle T4, ~1.05 s/step, ~2h

• Best checkpoint = lowest weight-averaged chat validation loss (it was the final step):

  Val loss	start (smoltalk)	chatqa
  UltraChat	3.056	2.407
  QuAC	2.608	0.563
  CoQA	3.580	1.358
  smol-smoltalk	2.074	2.150
  TinyStories	1.251	1.243

  QuAC/CoQA flattened after ~4,500 steps; smol-smoltalk got slightly worse despite the replay.

Results

All outputs are in reports/. Generations are sampled, so a re-run gives different text.

Test	Model	Result
Passage QA, 20 hand-written questions (chat_eval.txt)	chat	15/20 (base model with prompt tricks: 12/20)
Passage QA, 150 held-out generated questions (chat_eval.txt)	chat	135/150 = 90%
HellaSwag val, 10,042 items, acc / acc_norm	random weights	0.255 / 0.250
	base	0.262 / 0.265
	chat	0.263 / 0.259
	smoltalk	0.260 / 0.254
	chatqa	0.263 / 0.260
Grammar correction, 8 sentences (grammar_check_base.txt)	base	7 fixed or unchanged correctly, 1 changed the meaning (She → They)
CoQA held-out, token F1, 3rd question with 2 turns of history (qa_compare.txt)	smoltalk	0.04
	chatqa	0.48
QuAC held-out, token F1 (qa_compare.txt)	smoltalk	0.11
	chatqa	0.26
Says "does not say" on unanswerable questions (qa_compare.txt)	chatqa	20% (and wrongly on 13% of answerable ones)

QA tests use QuAC/CoQA validation[300:] (~80 questions per bucket), never seen in training or checkpoint selection; greedy decoding.

HellaSwag is close to chance (25%). That is expected: it tests general world knowledge, while this model has only read children's stories. It is included as a sanity check, not a strength. The chat fine-tunes (smoltalk, chatqa) do not move it either: SFT teaches answer style, not world knowledge, and all variants are within ~1.5 points of chance.

Example: chatqa (answering from a context)

You : Context:
      The Eiffel Tower is an iron tower in Paris, France. It was built by Gustave Eiffel's company
      and finished in 1889 for the World's Fair. It is about 330 metres tall.

      When was the tower finished?
Bot : 1889.

Held-out CoQA follow-ups: "Who was restraining him?" → Seth. · "Who found the nest?" → Phonny. (smoltalk answered neither).

Example: chat

You : hi
Tiny: Hey! I am Tiny. I can tell stories and answer questions about them.
You : Why do we brush our teeth?
Tiny: We brush our teeth to keep them clean and strong, so they do not hurt.
You : What is the capital of France?
Tiny: I'm sorry, I don't know about that. I know about kids, animals and stories. Shall I tell you a story?
You : I am sad today
Tiny: Oh no, I am sorry you feel sad. Do you want a happy story to cheer you up?

Example: base story (prompt "Once upon a time")

Once upon a time, there was a little girl named Lily. She loved to play in her room with her toys.
One day, she saw her mom cleaning and said, "Mommy, can I help you?" Her mom said, "Sure, Lily.
You can help me wipe the floor."

Example: grammar check (base model scores candidate edits)

IN : the girl want to go hoospital
OUT: the girl wanted to go to the hospital
IN : Tom wanted play with his freind.
OUT: Tom wanted to play with his friends.

All reports

File	Contents
reports/stories_base.txt	5 stories from "Once upon a time"
reports/stories_base_prompts.txt	4 stories from custom opening lines
reports/chat_eval.txt	4 chat conversations, 20 hand-written QA, 150 held-out QA
reports/ask_base.txt	Base model answering questions framed as story dialogue
reports/passage_qa_base.txt	Base model answering questions about a passage
reports/grammar_check_base.txt	8 grammar/spelling corrections
reports/hellaswag_{random,base,chat,smoltalk,chatqa}.txt	HellaSwag benchmark runs
reports/qa_compare.txt	smoltalk vs chatqa: held-out QuAC/CoQA F1, abstention, sample chats

Usage

pip install torch tiktoken safetensors huggingface_hub
# get the code folder, then from inside it:
python from_hub.py --repo sraivante/tinystories-35m           # writes out/final.pt (base), out/chat.pt, out/smoltalk.pt, out/chatqa.pt
python sample.py --ckpt out/final.pt --prompt "Once upon a time"
python chat.py                              # interactive chat with Tiny
python grammar_check.py "He have a big red ball."
./start.sh                                  # web UI at http://127.0.0.1:8765

Or load directly in Python (with code/model.py on the path):

import torch, tiktoken
from from_hub import load
model, cfg = load("sraivante/tinystories-35m", "base")        # or "chat", "smoltalk", "chatqa"
enc = tiktoken.get_encoding("gpt2")
ids = torch.tensor([enc.encode_ordinary("Once upon a time")])
out = model.generate(ids, 200, temperature=0.8, top_k=200, stop_token=enc.eot_token)
print(enc.decode(out[0].tolist()))

Chat prompt format: User: <message>\nAssistant:; stop when the model writes \nUser:. Passage QA with chatqa: User: Context:\n<passage>\n\n<question>\nAssistant:; greedy decoding works best for short answers.

Limitations

• Knows only a little about the world; no reliable facts, maths or general knowledge. General chat answers from smoltalk/chatqa are fluent but often wrong or off-topic.
• Stories can lose track of characters or contradict themselves, especially past ~300 tokens.
• chatqa answers well when the answer is in the context, but usually invents an answer when it is not (says "does not say" only 20% of the time), and sometimes abstains when the answer is there.
• smoltalk/chatqa do not do grammar or spelling correction by generation (chatqa often refuses with "I do not have the ability..."); use grammar_check.py, which scores candidate edits instead.
• Passage QA misses questions that need reasoning across sentences.
• Not safety-tuned beyond the chat model's out-of-scope refusals; not intended for production use.

Files

config.json  configs of all variants (read by from_hub.py)
base/  model.safetensors, config.json
chat/  model.safetensors, config.json
smoltalk/  model.safetensors, config.json
chatqa/  model.safetensors, config.json (includes per-source val losses)
code/  model, training, fine-tuning (finetune_chat.py, finetune_sft.py, finetune_mix.py), evaluation, inference, web UI, from_hub.py
logs/  train.log, finetune_chat.log, finetune_sft.log, finetune_mix.log
reports/  evaluation outputs (see above)
run_reports.sh  regenerates the base/chat reports