tiny-llm: train.py works, and the model is actually learning

model.py and prepare.py were done, but nothing had actually been trained yet. This devlog covers train.py and the first real smoke test on the 24M-parameters config.

Writing train.py

Wrote the training loop the same way as the model piece by piece, understanding each part before moving on:

• get_batch — pulls random windows from the memmapped .bin files. x is a block of tokens, y is the same block shifted by one position (next-token prediction). Understanding why y is just x shifted by 1 was one of those small moments where next-token prediction actually clicked.
• Mixed precision (fp16 + GradScaler) — my GPU (RTX 2080 Super, Turing) doesn’t have real bf16 tensor cores, so fp16 with a GradScaler is the right call. The scaler exists because fp16 has a much narrower range than fp32 — small gradients can underflow to zero and silently vanish. The scaler multiplies the loss up before backward() and unscales before step(), so nothing gets lost in the process.
• Gradient clipping — caps the gradient norm at 1.0 before the optimizer step, to stop occasional exploding gradients from wrecking a good run.
• Validation loss — evaluates on a held-out batch from val.bin every 100 steps, wrapped in model.eval() + @torch.no_grad(), so I can actually tell if the model is generalizing instead of just memorizing.
• Checkpointing — saves model + optimizer state every 1000 steps. Training a real run for tens of hours without this would be a great way to lose everything to one crash.

Debugging session

Hit two device-mismatch errors early on — torch.arange and torch.triu both default to creating tensors on CPU, even when the rest of the model is on the GPU. Fixed by explicitly passing device=x.device wherever a new tensor gets created inside the model. Small bug, but a good reminder that PyTorch doesn’t infer device placement for you.

First real run

Ran the 24M-param smoke test (n_layer=6, n_head=6, n_embd=384, block_size=256) on TinyStories, ~470M tokens, for 2000 steps on my RTX 2080 Super:
step 0: train loss 11.01 val loss 10.62
step 1000: train loss 3.34 val loss 3.47
step 1900: train loss 2.89 val loss 3.10

Loss starts right around ln(50257) ≈ 10.8, which is exactly what you’d expect from a model guessing uniformly over the vocabulary before it’s learned anything. Train and val loss track each other closely the whole way down — no sign of overfitting, which is what I wanted to confirm before committing to a much longer run.

What’s next

The full pipeline is confirmed working end-to-end: prepare.py → get_batch → model.py → train.py, all on my own hardware. Next step is adding a cosine learning rate schedule with warmup, then moving on to the actual 100M-param run.

PS: train.py had some AI usage