Vaibhaav Tiwari · 2026-11-13 · notable
deepseek-v3-jax — DeepSeek V3 in 190 lines of pure JAX
In short
deepseek-v3-jax is an MIT-licensed, 190-line JAX version of the DeepSeek V3 architecture. The 36.2M-parameter model trained in about 2.6 hours on a MacBook Air M5. The author calls it an architecture demo, not a strong language model.
Vaibhaav Tiwari · 2026-11-13 · 3 sources checked
deepseek-v3-jax rebuilds DeepSeek V3's key ideas in 190 lines of plain JAX that you can read in one sitting.
Key specs
| Parameters | 36.2M |
|---|---|
| Lines of code | 190 |
What is it?
deepseek-v3-jax is a small, readable copy of the DeepSeek V3 design written in pure JAX, with no Flax, Equinox or PyTorch. It has Multi-head Latent Attention, a DeepSeekMoE layer with 1 shared and 16 routed experts (top-4), bias-based load balancing without an auxiliary loss, and multi-token prediction. The model has 36.2M parameters, 13.6M of them active per token.
How does it work?
Parameters are plain dictionaries of arrays, and compute runs in bf16 on f32 master weights. The author trained for 12,000 steps (about 98M tokens from 207 books) in roughly 2.6 hours on a MacBook Air M5 with 24 GB, using the jax-metal plugin. A per-expert bias changes which experts the router picks, not their weights. At step 2000 each expert got 5.7–6.9% of tokens with the bias, against 0.09–18.6% without it.
Why does it matter?
The write-up shows that balancing experts with a bias cost almost no accuracy: validation loss was 3.921 with it and 3.923 without. It also shows Multi-head Latent Attention caching 144 numbers per token per layer, against 480 for standard attention in this model. The author says the output reads like detective fiction but the stories do not hold together.
Who is it for?
ML engineers learning how DeepSeek V3 works
Try it
git clone https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax.git && cd deepseek-v3-jax && uv venv -p 3.12 .venv