开发者用 190 行纯 JAX 复现 DeepSeek V3
开发者发布开源项目，用约 190 行纯 JAX 代码从零实现 DeepSeek V3 核心架构，包含多头潜在注意力（MLA）、解耦 RoPE、DeepSeekMoE 与无辅助损失负载均衡等关键设计，并在 MacBook Air 上完成预训练。项目不依赖 Flax、Equinox 或 PyTorch，参数以纯字典数组表示，代码可逐行阅读，是理解前沿 MoE 架构的高价值教学与实验范本。
报道正文
本文目录4 个章节
DeepSeek V3 in pure JAX
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.
What's in the 190 lines
| 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 |
Results
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.
Try the trained model
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.
Reproduce it
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 tokensThe 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.jsonYou 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 exac
（注：更多技术实现细节与完整 API 文档请查阅原项目 README）