tiny-llm: 24M model trained, and it actually generates coherent stories

Picked up from the last devlog — train.py was working (fp16, grad clipping, val loss, checkpointing), but only tested on 2000 steps. This one covers finishing the training setup, the real 4-hour run, and getting the model talking through a CLI.

Cosine LR with warmup

Added a proper learning rate schedule before committing to a long run:

• Warmup — LR ramps up linearly from 0 to the target (3e-4) over the first 2000 steps, instead of hitting the model with full LR while the weights are still random.
• Cosine decay — after warmup, LR follows a cosine curve down toward 0 over the rest of training, so the model settles into finer adjustments by the end instead of taking large steps the whole way through.

Nothing exotic here — PyTorch doesn’t need a library for this, just a function that computes the right LR for the current step and gets applied to the optimizer’s param groups each iteration.

The real run

474M tokens (the full TinyStories corpus), which comes out to ~231,445 steps at batch_size=8, block_size=256. Ran for about 4 hours on the RTX 2080 Super:
step 0: train loss 11.01 val loss 10.62
step 231000: train loss 1.48 val loss 1.49
step 231400: train loss 1.82 val loss 1.69

Train and val loss tracked closely the entire run — no overfitting. Loss in the 1.4–1.8 range by the end, a solid result for a from-scratch 24M model on this corpus.

sample.py: watching it actually generate

Wrote a script to load a checkpoint and generate text autoregressively — feed in a prompt, predict the next token, append it, repeat. Had to fix two small but important bugs along the way:

• torch.arange and torch.triu inside the model default to creating tensors on CPU even when everything else is on GPU — same device-mismatch class of bug as during training, fixed the same way (device=x.device).
• Loading the checkpoint failed at first because torch.compile prefixes every parameter name with _orig_mod. — stripped it with a quick dict comprehension before calling load_state_dict.

First real generation, prompted with “once upon a time”:

once upon a time, there was a fish named Fin. Fin was very brave and always had a big challenge to do. One day, Fin asked his friend, Sally, “Will you marry me?” Sally said, “Yes, Fin! Let’s go!” Fin and Sally had a big adventure in the ocean…

Not perfect logic (fish swimming through an ocean “full of fish to eat” is a little off), but grammatically solid, consistent character names throughout, and a real narrative arc. Exactly what you’d expect from a well-trained 24M model on TinyStories.

chat.py: a quick interactive loop

Wrapped the generation code in a while True loop so the model loads once and takes prompts repeatedly instead of restarting for every single generation. Worth noting for anyone trying this themselves: this is a base model, not an instruction-tuned chat model — it doesn’t “answer” prompts, it just continues them as text. Say “hello” and it treats that as the opening of a story, not a greeting to respond to. That’s expected behavior at this stage, not a bug — instruction-following is a separate fine-tuning step for later.

What’s next

The full pipeline — prepare.py → train.py → sample.py/chat.py — is confirmed working end-to-end, from raw dataset to a model you can actually talk to (in the “continues your text” sense). Next up: scaling to the 100M-param run on a 2B-token corpus, this time blending real security-focused text (HackTricks, PayloadsAllTheThings, and others) with general web text, plus a retrieval index so the model can pull in relevant passages instead of relying purely on memorized weights.