Building a 200M parameter LLM from scratch in PyTorch
DEV Community

Building a 200M parameter LLM from scratch in PyTorch

Overview

It's really easy to spin up Unsloth and fine-tune Llama 3 in an afternoon. This author wanted to see what happened when doing things from scratch instead - writing the BPE tokenizer, implementing RoPE, coding the SwiGLU feed-forward blocks, and writing the training loop in pure PyTorch from absolute scratch. The result is EmsyAI (V4), a 196M parameter model trained on a consumer GPU.

Technical Specifications

Active Parameters: 196.7M (roughly 180M in the transformer blocks, the rest in the ~16k vocab embeddings)

Training Metrics:

  • Training Tokens: 1.96 Billion
  • Context Window: 4,096 tokens
  • Hidden Dimension: 1,024
  • Attention: GQA (16 Query / 4 KV)
  • FFN Dimension: 2,816

Pretraining was performed for 15,000 steps using mixed-precision FP16. The training curve remained stable rather than spiking to NaN, ultimately decaying to a validation perplexity of 5.86. The transition from random token garbage at step 100 to recognizable Python syntax at step 10,000 is described as "bizarrely satisfying."

Instruction Tuning & Evaluation

After pretraining, the base model was instruction-tuned with a 2.3M parameter LoRA adapter on the CodeAlpaca dataset. When evaluated against OpenAI's HumanEval benchmark, the model scored exactly 0.0%. While disappointing for a small model trained on 2 billion tokens, this aligns with the expectation that such a compact model cannot solve multi-step algorithmic puzzles.

An audit of the training data revealed that 1,644 sequences in CodeAlpaca leaked HumanEval test logic. Despite seeing these leaked solutions during fine-tuning, EmsyAI still scored 0%, indicating the tiny 2.3M LoRA adapter simply lacked the capacity to rote-memorize those sequences verbatim. The contamination did not enable cheating because the model could not retain the answers anyway.

Architectural Challenges & Solutions

Scaling the architecture to 196M parameters exposed several edge cases missed at the 88M scale. One critical issue was attention logit explosions: at hidden dimension 1024, attention logits would occasionally spike wildly. Standard architectures do not normalize queries and keys before the dot product, creating a stability liability at scale.

Earlier work suggested fixes in Qwen 2.5 and Gemma 2, but those were incorrect - Gemma 2 used attention-logit soft-capping. True stabilization came from QK-Norm, which was widely credited after OLMo 2, Gemma 3 (inspired by Meta's Chameleon), and Qwen 3. Version 5 will implement QK-Norm to prevent these spikes.

Another problem involved tokenizer byte collisions. Although the BPE tokenizer was written from scratch, the fallback byte mappings contained bugs: bytes 0-3 currently overlap with special control tokens. This silent bug does not crash training but degrades performance. Version 5 includes a complete tokenizer rebuild and doubles the vocabulary from 16k to 32k.

Cross-document attention also poses a risk. Currently, the dataloader blindly concatenates text files to reach the 4096 context length, causing the model to attend across unrelated documents that merely share the same training window. The solution requires document-aware packing with proper attention masking.

Deployment & Resources

The final weights have been exported to GGUF format, allowing execution via Ollama. To run the model:

git clone https://github.com/gulding/EmsyAI.git
huggingface-cli download gulding/EmsyAI emsyai-v4-instruct-f32.gguf --local-dir .
ollama create emsyai-v4 -f Modelfile
ollama run emsyai-v4

The full PyTorch training loop, architecture, and tokenizer code are available on GitHub: https://github.com/gulding/EmsyAI.

Read on DEV Community ↗ ← Back to News

Comments

No comments yet. Start the discussion.