tiny-llm: the 100M model is trained, and RAG is wired in

The 100M parameter run is done.
976,562 steps on a 2B-token corpus, and chat.py now pulls retrieved passages before generating instead of relying on memorized weights.

The problem, and the fix

Started out wanting this model trained purely on security content (HackTricks, PayloadsAllTheThings, SecLists, GTFOBins, CheatSheetSeries, wstg).
All of that combined comes to about 6.73M tokens. Compute-optimal for a 100M model is around 2B tokens, real security wikis just don’t reach that scale, and no amount of repo-hunting was going to fix that 300x gap.

The solution i found was blend the 6.73M security tokens into a 2B-token corpus, topped up with FineWeb-Edu for general English coverage. The model gets real security-domain exposure plus enough general text to actually hold together grammatically, not by inventing security content that doesn’t exist, but by being upfront about what’s actually available.

The training run

d_model=768, n_heads=12, n_layer=12 4x the previous 24M config. Config changes were straightforward, getting a resumable 50-hour run right was the actual work:

• Checkpoint pruning — early on, checkpoints/ silently grew to 561GB before a torch.save call failed with a disk-full error mid-write. Fixed by keeping only the 5 most recent checkpoints on every save, removing older ones automatically.
• torch.compile ordering — resuming from a checkpoint failed with a state_dict key mismatch the first time, because model = torch.compile(model) was running before the resume logic loaded the weights. Compiled models get every parameter name prefixed with _orig_mod., so loading a plain checkpoint into a compiled model (or vice versa) breaks. Fixed by moving the resume block to before the torch.compile call — load into the “bare” model first, compile after.
• Verified resume actually works — killed the process mid-run, restarted, confirmed it picked up exactly where it left off (same step, same loss) before trusting it for the full 50 hours unattended.

Final loss landed around 3.5–3.8 (train and val tracking closely) — expected for a corpus this heterogeneous (security docs + general web text) compared to the tighter 1.4–1.8 range on TinyStories.

Wiring in retrieval

A 100M model isn’t going to reliably “know” security facts from 6.73M tokens seen once inside a 2B-token corpus — so instead of hoping the weights memorized it, chat.py now retrieves first:

1. User’s prompt goes to a local SQLite FTS5 search index built over the same repos/ content used in training
2. Top-k matching passages come back with source, heading, and BM25 relevance score
3. Those passages get prepended to the prompt as context before the model generates

Straightforward in concept, but ran into block_size=256 limits fast, retrieved passages plus the question can easily blow past what the model was trained to handle positionally. Fixed by truncating the encoded prompt to the last 256 tokens before generation, keeping the question and “Answer:” cue intact even if some of the retrieved context gets cut from the front.

What actually happens now

Asked “What is SQL injection?” — with retrieval wired in, the model opens with something closer to on-topic (“one of the most powerful tools available to attackers”) instead of pure noise, but drifts into plausible-sounding but fabricated details a few sentences in. That’s the honest result at this scale: retrieval gives it real information to start from, but staying faithful to that context for 100 generated tokens is past what 100M parameters can reliably do. Better than without retrieval, not a fix for the underlying scale limit.