189 lines, 39 tests pass: I ran DeepSeek-V3 in pure JAX
A Show HN post this week claims the whole DeepSeek-V3 architecture — MLA, DeepSeekMoE, aux-loss-free balancing, multi-token prediction — in 190 lines of pure JAX, pretrained from scratch on a MacBook Air. I counted 189 lines of actual code, ran all 39 of the repo's tests on a CPU, and they passed in 91 seconds. The released checkpoint's SHA-256 matches, its final validation loss is 3.507, and the median training speed in its own log is 10,677 tokens/s. Every headline number checks out — except one, which needs a footnote.
What it is
The repo is Vaibhaav-Tiwari/deepseek-v3-jax, posted November 13. It is a 36.2M-parameter toy, not a real V3: 6 layers, d_model 384, vocab 8,192, trained on 207 public-domain detective novels. No Flax, no Equinox, no PyTorch — params are plain nested dicts of arrays, and the whole model is one file, deepseek_v3.py.
The ideas it squeezes in, per the README:
| Idea from the paper | Implementation |
|---|---|
| Multi-head Latent Attention (MLA) | One compressed latent per token, cached instead of full K/V |
| Decoupled RoPE | One rotary key shared by all heads, kept outside the latent |
| DeepSeekMoE | 1 shared expert + 16 routed, top-4 per token |
| Aux-loss-free balancing | A per-expert bias nudges routing toward idle experts |
| Sequence-wise balance loss | The paper's small complementary loss, α = 1e-4 |
| Multi-token prediction (MTP) | An extra module predicts the token after next, λ = 0.3 |
The line count: 189, not 190
wc -l says 265 lines. Counting the way a skeptic would — dropping blank lines, comment-only lines, and docstrings:
total=265 blank=56 comment-only=16 docstring=4 code=189 189 lines of code. The README says 190. That is a rounding error, not a lie — one line either way depending on whether you count the module docstring. The claim survives contact with wc.
The tests: 39 passed on a CPU in 91 seconds
The repo ships tests.py with checks for the parts that would actually break: MLA is causal, the MLA cache matches a full forward pass, routing picks top-k and normalizes, the bias changes which experts are picked but not the gate values, MTP gets gradient into the shared head, and the model can overfit one batch. I installed CPU JAX 0.4.34 on this AMD EPYC box (the jax-metal requirement is Mac-only, so I skipped it) and ran the suite:
JAX_PLATFORMS=cpu .venv/bin/pytest tests.py -q ....................................... [100%] 39 passed in 90.95s 39 passed, 0 failed. The most useful test is test_mla_cache_matches_full_forward — it is the one that would catch a wrong cache, and the cache is the whole point of MLA.
The training log: every number reproduces
The v1.0 release ships the checkpoint, the tokenized data, and metrics-main.jsonl — 480 logged rows. I read the log instead of trusting the README:
step loss val tok/s 25 10.77 — 10,473 250 — 5.018 — 2000 — 3.964 — 12000 3.76 3.507 8,890 - Validation loss 5.02 → 3.51. The README says 5.02 at step 250 and 3.51 at the end. The log says 5.018 and 3.507. Exact.
- Speed ~10.7k tokens/s (median). Median of all 480 logged
tok_per_svalues: 10,676.5. Exact. - 12,000 steps, ~98M tokens, ~2.6 h. Config is batch 32 × seq 256 = 8,192 tokens/step, so 12,000 steps is 98.3M tokens. At the logged median speed that is 2.55 hours. The README's 2.6 hours is the same number with rounding.
- 21.7M training tokens.
train.binis 43,379,858 bytes; tokens are stored as uint16, so 21,689,929 tokens. Exact. - Checkpoint integrity.
ckpt.pklis 144,818,333 bytes and its SHA-256 matches the release'sSHA256SUMSexactly. The tokenizer matches too.
The one footnote: 13.6M active per token
Total parameters: I counted them from a fresh init_params — 36,203,072, exactly the README's 36.2M. The active-per-token number is where it gets interesting:
total: 36,203,072 (36.2M) MTP module: 1,783,616 (train-only, unused at inference) routed experts: 23,592,960 (16 per MoE layer, top-4 used) active / token: 16,724,736 (16.7M, counting the output head) active / token: 13,579,008 (13.6M, excluding the 3.1M output head) The README's 13.6M holds only if you don't count the 3,145,728-parameter output head. Count it — it runs on every token at inference — and the number is 16.7M. This is a convention choice, not a fabrication, but the honest statement is "13.6M of transformer body, 16.7M including the head." The untied embedding and head each cost 3.1M parameters, which is 17% of the model on a vocab of 8,192.
The MLA cache claim checks out cleanly: 144 numbers cached per token per layer (kv_rank 128 + d_rope 16) versus 480 for full K/V (6 heads × (32 + 16 + 32)). A 3.3x smaller cache, exactly as advertised.
Expert balance: the bias does what the paper says
The README claims expert load of 5.7–6.9% per expert at step 2000 with the bias on (even would be 6.25%), versus 0.09–18.6% without it. Recomputed from the published log across all five MoE layers, I get 5.38–7.16% at step 2000 — slightly wider than the README's range, same picture: no expert starved, none hogging. At step 25, before the bias had done its work, loads ranged from 0.006% to 24%. The collapse-then-balance arc is in the data, not just in the prose.
Generation: it writes like a detective novel, mostly
I ran the released checkpoint with the repo's own sample.py, prompt "The fog came in off the", seed 0:
--- sample 1 --- The fog came in off the rain from the great, cloudy sun. The darkness was bright and dark. Holmes led the way into the garden and up to the window. ... "I hope that you will pardon my troubling you, Holmes," said he. "By the way, it is that your friend Watson is an invalid..." Coherent genre pastiche, correct character names, occasional nonsense ("the darkness was bright and dark"). The README's quoted line — "The fog came in off the bay, and the great trees were spreading and the stars were shining" — is a cherry-picked sample, not the deterministic output: sampling at temperature 0.8 with top-k 50 means every run differs. That is fine for a 36M model trained on Hammett; the point is the architecture runs, not that it writes well.
What I could not verify
The "trained on a MacBook Air M5" claim is unverifiable from here — I have no MacBook, and the log only proves the speed (10,677 tok/s median) is consistent with the 2.6-hour claim. The jax-metal requirement means the README's install path is Mac-first; on Linux you need CPU JAX and patience. And this is a 36M-parameter architecture demo: it proves the ideas fit in 190 lines, not that a 190-line model is any good. The validation loss of 3.51 on held-out crime fiction is honest about that.
The takeaway
The most surprising thing about this repo is not that DeepSeek-V3 fits in 190 lines. It is that the author shipped the receipts: a test suite that actually tests the tricky parts, a training log with 480 rows, and a checkpoint whose hash matches. I re-ran all of it on a CPU I own, and the only correction I found is a 3.1M-parameter footnote on one number. More repos should be auditable like this.
Sources: Vaibhaav-Tiwari/deepseek-v3-jax (repo, created November 13, 2026); v1.0 release (checkpoint, data, metrics); write-up. All test runs, line counts, parameter counts, and log arithmetic are my own, run November 15, 2026 on an AMD EPYC CPU box with JAX 0.4.34 (CPU).
