The architecture from the DeepSeek-V3 technical report in 190 lines of pure JAX, pretrained from scratch on a MacBook Air on public-domain detective fiction.
No Flax, no Equinox, no PyTorch. Params are plain dicts of arrays and every part of the model is a function you can read top to bottom in deepseek_v3.py.
Read the write-up: DeepSeek V3 in 190 lines of pure JAX, with diagrams, the code explained piece by piece, and the results.
| Idea from the paper | What it does | Where |
|---|---|---|
| Multi-head Latent Attention (MLA) | Caches one small latent per token instead of full keys and values: 480 → 144 numbers per token per layer here | mla, init_cache |
| Decoupled RoPE | One rotary key shared by all heads, kept outside the compressed latent | rope, apply_rope, mla |
| DeepSeekMoE | 1 shared expert that always runs + 16 routed experts, top-4 per token | route, moe |
| Auxiliary-loss-free balancing | A per-expert bias nudges routing toward idle experts; it changes which experts are picked, not their weights | update_bias, route |
| Sequence-wise balance loss | The small complementary loss from the paper (α = 1e-4) | seq_balance_loss |
| Multi-token prediction (MTP) | An extra module also predicts the token after next, weighted λ = 0.3 in the loss | forward, loss_fn |
| Training and generation | bf16 compute on f32 master weights; prefill + decode with the MLA cache | make_train_step, generate |
One run of the default config on a MacBook Air M5 (24 GB):
| Model | 36.2M parameters, 13.6M active per token |
| Data | 207 books, 21.7M training tokens (Hammett repeated 3×); validation is held-out crime fiction |
| Training | 12,000 steps, ~98M tokens (~4.5 epochs), ~2.6 h |
| Speed | ~10.7k tokens/s (median) |
| Validation loss | 5.02 (step 250) → 3.51 |
| Expert load, with bias | 5.7–6.9% per expert (even is 6.25%), at step 2000 |
| Expert load, without bias | 0.09–18.6% per expert, at step 2000 |
With balancing on and off, validation loss at step 2000 was the same (3.921 vs 3.923). The bias evens out the experts without costing anything, which is the paper's claim.
Its own words, from the prompt "The fog came in off the":
The fog came in off the bay, and the great trees were spreading and the stars were shining.
This is an architecture demo, not a strong language model. See the write-up for what worked, what didn't, and what's missing.
The v1.0 release has the exact checkpoint and data from the published run. After installing (step 1):
mkdir -p data runs/main
curl -L -o runs/main/ckpt.pkl https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/releases/download/v1.0/ckpt.pkl # 138 MB
curl -L -o data/tokenizer.json https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/releases/download/v1.0/tokenizer.json
.venv/bin/python sample.py --prompt "The fog came in off the" --num 3The checkpoint is a Python pickle, which can run code when loaded. Only load checkpoints from sources you trust.
You need an Apple Silicon Mac, Python 3.12 and uv. The published run used 24 GB of memory. No Mac? See Running on CPU.
1. Install
git clone https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax.git && cd deepseek-v3-jax
uv venv -p 3.12 .venv
uv pip install --python .venv/bin/python -r requirements.txtEvery version in requirements.txt is pinned. jax-metal 0.1.1 only works with jax==0.4.34, so don't upgrade JAX.
2. Check your setup
JAX_PLATFORMS=cpu .venv/bin/pytest tests.py -q # 39 passed
.venv/bin/python smoke_metal.py # optional: what jax-metal supports on your MacIn smoke_metal.py, "grad through 2-axis contraction" fails on jax-metal 0.1.1. That's expected, and the model is written to avoid it.
3. Build the dataset
.venv/bin/python prepare_data.pyThis downloads the 207 books in sources.txt (~120 MB, about 7 minutes), cleans them, trains an 8,192-token BPE tokenizer and writes data/train.bin, data/val.bin and data/tokenizer.json. It should print:
train: 211 docs, 83.4M chars, 21.69M tokens
val: 4 docs, 1.1M chars, 0.30M tokens
The published run's data had these hashes:
shasum -a 256 data/train.bin data/val.bin data/tokenizer.jsond37aca97e707bb81b6e4d266cc1adfd1f3e3556457c6eb724e495118a734e0bb data/train.bin
b7ac285b115e6793b5d36aa5b4bc6e81a5fccfc7af0908a306d3828ebca9c31c data/val.bin
c5a142b4006a967c880dda44ad792719d53c16d947da3b6c2332d6132d9967cc data/tokenizer.json
You will most likely get different hashes. The Hammett books are pinned to exact Standard Ebooks commits and come out identical, but Project Gutenberg files can't be pinned, and Gutenberg has since re-issued 28 of the 203 books with small edits (mostly accents stripped, like "José" → "Jose"). The token counts come out the same, so training should behave the same.
To train on the exact same bytes, skip prepare_data.py and download the published data instead:
mkdir -p data
for f in train.bin val.bin tokenizer.json; do curl -L -o data/$f https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/releases/download/v1.0/$f; done
shasum -a 256 data/*.bin data/tokenizer.json # should match the hashes above4. Train
.venv/bin/python train.py # the main run: 12k steps, ~2.6 h
.venv/bin/python train.py --steps 2000 --out runs/bias_on # balancing ablation, with the bias
.venv/bin/python train.py --steps 2000 --out runs/bias_off --no-bias # ...and without itEach run writes metrics.jsonl (loss, validation loss, per-expert load, biases, tokens/s) and ckpt.pkl (saved at every evaluation) into its --out folder, which defaults to runs/main. Seeds are fixed, but GPU math on Metal isn't bit-for-bit deterministic, so expect your numbers to land close to the ones above rather than match them exactly.
5. Generate
.venv/bin/python sample.py --prompt "The fog came in off the" --num 3Defaults: temperature 0.8, top-k 50, 300 new tokens per sample (--temp, --top-k, --n).
6. Make the figures
.venv/bin/python plots.py # needs runs/main, runs/bias_on, runs/bias_off
JAX_PLATFORMS=cpu .venv/bin/python charts.py # the rest of the write-up's chartsBoth write into post/figures/.
To redraw the figures without training, download the logs from the release: metrics-main.jsonl, metrics-bias_on.jsonl and metrics-bias_off.jsonl go into runs/main/, runs/bias_on/ and runs/bias_off/ as metrics.jsonl. charts.py also needs the checkpoint and the data.
Everything runs without jax-metal. It's much slower, so use the smaller config:
JAX_PLATFORMS=cpu .venv/bin/python train.py --config small--config small (4 layers, 8 experts, top-2, 128-token context) is meant for 8–16 GB Macs too. It has only been checked for a couple of steps, not trained to the end.
| File | What it is |
|---|---|
deepseek_v3.py |
The model: config, MLA, MoE, balancing, MTP, loss, train step, cached generation. The 190 lines. |
train.py |
Training loop: AdamW with warmup and cosine decay, evaluation, checkpoints, logging |
prepare_data.py, sources.txt |
Download, clean and tokenize the corpus |
sample.py |
Text generation from a checkpoint |
tests.py |
39 tests: causal masking, cache vs full forward, routing, balancing, MTP targets, NaN handling, bf16 master weights, and more |
smoke_metal.py |
Which JAX operations work on Apple's GPU |
plots.py, charts.py |
The figures, all drawn from the logs, checkpoint and corpus |
post/ |
Diagrams, figures and screenshots from the write-up |
"190 lines" counts the non-blank, non-comment lines of deepseek_v3.py after its docstring:
tail -n +5 deepseek_v3.py | grep -cvE '^\s*(#|$)' # 190- Dense MoE compute. Every expert runs on every token and the unpicked ones are zeroed. The routing math is exact, but there's no sparse dispatch or expert parallelism.
- No weight absorption in MLA. The latent is decompressed explicitly, which is easier to read and slower at inference.
- Not included: FP8 training, DualPipe, node-limited routing, YaRN context extension, and any post-training.
- Scale. 36.2M parameters vs 671B, 6 layers vs 61, a 256-token context vs 128K.
All training text is public domain in the United States: Dashiell Hammett's Red Harvest, The Dain Curse, The Maltese Falcon and Continental Op stories (CC0 editions from Standard Ebooks), plus about 200 detective and crime novels from Project Gutenberg. Every source is listed in sources.txt. Public-domain status outside the US may differ.
DeepSeek-AI, DeepSeek-V3 Technical Report, 2024.
Code: MIT. The training text is public domain in the US (see above).