A DeepSeek-Style LLM from Scratch
DeepSeek's architecture ideas, rebuilt at 33M parameters
- Year
- 2025–26
- Kind
- Models
- Status
- In progress
Trains at 74% of the GPU's fp32 peak
A decoder-only language model written from scratch in PyTorch, with DeepSeek's V2 to V4.1 components as config switches, each measured against the same baseline.
The base is a modern dense transformer: grouped-query attention with QK-norm and a sigmoid output gate, SwiGLU, tied embeddings and logit soft-capping, trained with Muon. On top of it sit DeepSeek's components, each one a config switch: multi-head latent attention, DeepSeekMoE with auxiliary-loss-free load balancing, manifold-constrained hyper-connections (mHC), Engram n-gram memory and multi-token prediction. An ablation sweep trains each one against the same baseline on TinyStories to measure what actually helps at this size.
The problem
DeepSeek's papers from V2 to V4.1 each add architecture ideas: latent attention, expert balancing without an auxiliary loss, hyper-connections, hashed n-gram memory, a Muon variant. I wanted to build each one from its paper and measure which still help at 33M parameters, with a 4 GB laptop GPU and a few hours of a rented one.
How it works
- 01
Tokenize once, then read fixed windows
prepare_data.py streams a dataset through GPT-2's BPE tokenizer into flat uint16 files. Training cuts them into fixed windows and shuffles them with a permutation seeded by the epoch, so batch i is identical on every run. That is what makes exact resume after a crash possible.
- 02
Start from a current dense block
The baseline is a decoder-only transformer built the way recent open models are: RMSNorm, rotary positions, SwiGLU and grouped-query attention, where several query heads share one key/value head. It adds QK-norm, a gated attention output and logit soft-capping. The TinyStories config is 8 layers at width 384.
- 03
Switch in DeepSeek's attention and memory
Multi-head latent attention (MLA) caches one small latent and a rotary key per token, and rebuilds every head's keys and values from them. Engram hashes each token's last two and three tokens into large tables and gates what it finds into the residual stream, at constant cost per token.
- 04
Route each token to two of eight experts
DeepSeekMoE replaces the feed-forward layer in every block after the first with a shared expert plus routed experts, 2 of 8 per token in the ablation. Instead of an auxiliary loss, a per-expert bias nudged toward equal load after each step changes which experts are picked but not their weights, so balancing never fights the language-modeling loss.
- 05
Widen the residual stream to four
mHC (manifold-constrained hyper-connections) keeps four copies of the residual stream. Around every sublayer, small per-token maps choose how much of each copy to read, where to write the output and how to mix the copies. Sinkhorn-Knopp makes the mixing matrix's rows and columns each sum to 1, which keeps the gain through depth bounded.
- 06
Train with Muon on a token budget
Muon replaces each weight matrix's momentum update with the nearest orthogonal matrix, found by Newton-Schulz iterations: a few matrix multiplies that push every singular value toward 1. Embeddings, the output head and norms stay on AdamW. The learning rate warms up, holds, then decays over the last 20 to 25% of the token budget.
Results
- ~6,000
- tokens per second training the 33M model in fp32 on a 4 GB GTX 1650 Ti
- 74%
- of the card's fp32 peak at that speed
- < 2 GB
- of GPU memory for that run
Decisions
One switch per component
An ablation means something only if one thing changes at a time. ablate.py trains the base, the base plus each component, and all of them, on the same data and token budget, and skips finished runs so a crashed sweep resumes.
fp32 on older GPUs, after measuring
On the GTX 1650 Ti, fp16 matmuls measured 0.41 TFLOPS against 1.78 for fp32, and bfloat16 is only emulated, though torch.cuda.is_bf16_supported() returns True. So training uses bfloat16 on Ampere and newer, and plain fp32 before that.
No sparse attention
DeepSeek's sparse and compressed attention exist to make 64K to 1M-token contexts affordable. At 512 to 1,024 tokens their top-k selection would keep every token anyway.
What I learned
- Full fp32 logits over a 50,304-token vocabulary don't fit in 4 GB, so the loss is computed in chunks with recomputation.
- The mHC paper gives no bias initialization. I inferred one from DeepSeek-V4's released weights, and a test checks that a fresh mHC layer acts like a plain residual.