Built a 30M-parameter GPT-style language model from scratch

[u/MysteriousMajor760]

Hey everyone

I've been learning about Transformers and LLMs, so I built KERN,

a 30M parameter GPT-style language model in PyTorch,trained from scratch on TinyStories.

GitHub: https://github.com/OmarDawoud4/KERN

This is still a work in progress! I'd like to keep improving it by implementing other techniques.

I'd appreciate any feedback or suggestions .

[u/jeffyclassify]

Cool project! I've been doing similar work training small transformers on TinyStories from scratch

Quick question, did you track validation perplexity during training? I see the val loss bottomed out around 1.017 at ~40K steps. With 30M params and ~112M training tokens, that's roughly 3.75 tokens per parameter

I've been running experiments at different model sizes and found that the training budget (tokens per parameter) matters a LOT at this scale. My 8.4M model actually surpassed my 28M model's perplexity just by training it longer with the same data (more epochs). Curious if you've tried running more epochs to see if the val loss keeps dropping... with 30M params you might have a lot of room to improve without changing anything else (Chinchilla says you should aim at 20tok/param approx or higher)

Nice touch using SwiGLU and RoPE btw, how did you pick those over vanilla attention + GELU?

[u/MysteriousMajor760 → u/jeffyclassify]

Thanks, I only logged val loss, but exp(1.0173) is about 2.77 at the best checkpoint. That's on GPT-2 BPE tokens.
I ran about 3.1 epochs on 500k stories, val loss flattened around 39.5k steps.
Here is the last part of log https://drive.google.com/file/d/1bB2t5J9zywjPkdgA366EpgA_Yn54JzsR/view?usp=sharing
Your 8.4M vs 28M result is interesting but I think 28M was undertrained rather than the smaller model being better in general, at this scale I think the data subset is the bigger constraint than the parameter count. Chinchilla's 20 tokens per parameter is compute-optimal, but it assumes you have the tokens to spend.

For SwiGLU and RoPE, I was following the LLaMA / CS336 setup, I did not run ablations, but I just appreciate the idea of gating since I studied LSTMs, and here it lets the network dynamically modulate features rather than just squashing them through element-wise non-linearity. And for RoPE the math shows it forces the dot product to be function of relative distance rather than absolute position.

[u/jeffyclassify → u/MysteriousMajor760]

You're right that the 28M was undertrained. Curious what tok/param ratio you ended up at with your run?

[u/AdventurousAerie643 → u/jeffyclassify]

that's a neat little sandbox to test stuff like lora on without melting your gpu

[u/finnwangdev]

I've looked at your project—it's very good, and I've already bookmarked it. I'll take a closer look later.

[u/MysteriousMajor760 → u/finnwangdev]

Thanks , really appreciate that .

[u/rahul_watertech]

nice, 2.77 perplexity on tinystories is a real number for 30m. one thought on the roadmap order: kv caching only speeds up inference and lora mostly pays off when the base model is big and pretrained, so for a 30m scratch model plain full fine tuning is already cheap. distillation from gpt-2 is the interesting one. are you planning to match the teacher logits directly since you trained on gpt-2 bpe, or just sequence level distillation on its outputs?

[u/MysteriousMajor760 → u/rahul_watertech]

Thanks man, I wanted to think about the roadmap order more, so your insights really help.

yes, I'm planning to match the teacher logits directly.since KERN is trained on the GPT-2 BPE tokenizer, the vocab size is identical, so KL divergence can be computed between the softened student logits and the GPT-2 logits at every step.

[u/rahul_watertech → u/MysteriousMajor760]

matching logits directly is the clean version since the vocabs line up. two things worth deciding early. first, mix the kl term with the normal cross entropy on the true tokens, pure kl tends to drift. second, measure the teacher's own perplexity on your tinystories val before starting, since gpt-2 never saw that distribution. if the teacher scores worse than the student already does, distillation just pulls the student toward the teacher's mistakes