用约 200 行纯 JAX 复现 DeepSeek V3 架构

来自 DeepSeek-V3 技术报告的架构，用 190 行纯 JAX 实现，在 MacBook Air 上从零开始预训练，训练数据是公有领域的侦探小说。

没有 Flax，没有 Equinox，没有 PyTorch。参数就是普通的数组字典，模型的每一部分都是函数，你可以在 deepseek_v3.py 中从上到下读完。

阅读说明文章： 190 行纯 JAX 实现的 DeepSeek V3，包含图表、逐段讲解的代码以及结果。

190 行里有什么

论文中的思路	作用	位置
多头潜在注意力（MLA）	每个 token 只缓存一个小型潜在向量，而不是完整的键和值：此处每层每个 token 从 480 → 144 个数字	`mla`, `init_cache`
解耦 RoPE	所有头共享一个旋转键，保存在压缩潜在向量之外	`rope`, `apply_rope`, `mla`
DeepSeekMoE	1 个始终运行的共享专家 + 16 个路由专家，每个 token 选 top-4	`route`, `moe`
无辅助损失均衡	每个专家的偏置将路由推向空闲专家；它改变的是哪些专家被选中，而不是它们的权重	`update_bias`, `route`
序列级均衡损失	论文中的小型补充损失（α = 1e-4）	`seq_balance_loss`
多 token 预测（MTP）	一个额外模块还预测下下个 token，在损失中权重为 λ = 0.3	`forward`, `loss_fn`
训练与生成	在 f32 主权重上进行 bf16 计算；使用 MLA 缓存进行预填充 + 解码	`make_train_step`, `generate`

结果

在 MacBook Air M5（24 GB）上以默认配置运行一次：


模型	36.2M 参数，每个 token 激活 13.6M
数据	207 本书，21.7M 训练 token（Hammett 重复 3 次）；验证集是留出的犯罪小说
训练	12,000 步，约 98M token（约 4.5 个 epoch），约 2.6 小时
速度	约 10.7k token/s（中位数）
验证损失	5.02（第 250 步）→ 3.51
专家负载，带偏置	第 2000 步时每个专家 5.7–6.9%（均匀为 6.25%）
专家负载，不带偏置	第 2000 步时每个专家 0.09–18.6%

开启和关闭均衡时，第 2000 步的验证损失相同（3.921 对 3.923）。偏置在不付出任何代价的情况下均衡了专家，这正是论文的主张。

它自己的输出，来自提示词 "The fog came in off the"：

The fog came in off the bay, and the great trees were spreading and the stars were shining.

这是一个架构演示，不是一个强大的语言模型。关于哪些有效、哪些无效以及缺少什么，请参阅说明文章。

试用训练好的模型

v1.0 发布版包含已发布运行的确切检查点和数据。在安装（第 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 3

检查点是一个 Python pickle，加载时可能执行代码。只加载来自你信任来源的检查点。

复现它

你需要一台 Apple Silicon Mac、Python 3.12 和 uv。已发布的运行使用了 24 GB 内存。没有 Mac？请参阅在 CPU 上运行。

1. 安装

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.txt

requirements.txt 中的每个版本都已固定。jax-metal 0.1.1 只能与 jax==0.4.34 配合使用，所以不要升级 JAX。

2. 检查你的设置

JAX_PLATFORMS=cpu .venv/bin/pytest tests.py -q    # 39 passed
.venv/bin/python smoke_metal.py                   # optional: what jax-metal supports on your Mac

在 smoke_metal.py 中，"grad through 2-axis contraction" 在 jax-metal 0.1.1 上会失败。这是预期内的，模型编写时已避免它。

3. 构建数据集

.venv/bin/python prepare_data.py

这会下载 sources.txt 中的 207 本书（约 120 MB，约 7 分钟），清理它们，训练一个 8,192 token 的 BPE 分词器，并写入 data/train.bin、data/val.bin 和 data/tokenizer.json。它应该会打印：

train: 211 docs, 83.4M chars, 21.69M tokens
val: 4 docs, 1.1M chars, 0.30M tokens

已发布运行的数据具有以下哈希值：

shasum -a 256 data/train.bin data/val.bin data/tokenizer.json

d37aca97e707bb81b6e4d266cc1adfd1f3e3556457c6eb724e495118a734e0bb  data/train.bin
b7ac285b115e6793b5d36aa5b4bc6e81a5fccfc7af0908a306d3828ebca9c31c  data/val.bin
c5a142b4006a967c880dda44ad792719d53c16d947da3b6c2332d6132d9967cc  data/tokenizer.json

你很可能会得到不同的哈希值。Hammett 的书固定到 Standard Ebooks 的确切提交，因此输出完全一致，但 Project Gutenberg 的文件无法固定，而且 Gutenberg 此后重新发布了 203 本书中的 28 本，做了小幅修改（主要是去掉重音符号，比如 "José" → "Jose"）。token 数量相同，所以训练行为应该一致。

要在完全相同的字节上训练，跳过 prepare_data.py，改为下载已发布的数据：

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 above

4. 训练

.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 it

每次运行都会把 metrics.jsonl（loss、验证 loss、每个专家的负载、偏置、tokens/s）和 ckpt.pkl（每次评估时保存）写入其 --out 文件夹，该文件夹默认为 runs/main。随机种子是固定的，但 Metal 上的 GPU 数学运算并非逐位确定，所以你的数值应该接近上面的结果，而不是完全一致。

5. 生成

.venv/bin/python sample.py --prompt "The fog came in off the" --num 3

默认值：temperature 0.8、top-k 50、每个样本 300 个新 token（--temp、--top-k、--n）。

6. 生成图表

.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 charts

两者都写入 post/figures/。

要在不训练的情况下重绘图表，请从 release 下载日志：metrics-main.jsonl、metrics-bias_on.jsonl 和 metrics-bias_off.jsonl 分别以 metrics.jsonl 的形式放入 runs/main/、runs/bias_on/ 和 runs/bias_off/。charts.py 还需要 checkpoint 和数据。

在 CPU 上运行

一切都可以在没有 jax-metal 的情况下运行。速度会慢得多，所以请使用较小的配置：

JAX_PLATFORMS=cpu .venv/bin/python train.py --config small

--config small（4 层、8 个专家、top-2、128-token 上下文）也适用于 8–16 GB 的 Mac。它只检查了几步，并未训练到结束。

文件

文件	内容
`deepseek_v3.py`	模型：config、MLA、MoE、均衡、MTP、loss、训练步、缓存生成。那 190 行。
`train.py`	训练循环：带 warmup 和 cosine 衰减的 AdamW、评估、checkpoint、日志
`prepare_data.py`、`sources.txt`	下载、清洗并 tokenize 语料库
`sample.py`	从 checkpoint 生成文本
`tests.py`	39 个测试：因果掩码、缓存与完整前向、路由、均衡、MTP 目标、NaN 处理、bf16 主权重等
`smoke_metal.py`	哪些 JAX 操作可在 Apple 的 GPU 上运行
`plots.py`、`charts.py`	图表，全部来自日志、checkpoint 和语料库
`post/`	文章中的示意图、图表和截图

"190 行"统计的是 deepseek_v3.py 在其 docstring 之后的非空、非注释行：

tail -n +5 deepseek_v3.py | grep -cvE '^\s*(#|$)'    # 190

简化之处

稠密 MoE 计算。每个专家都会对每个 token 运行，未被选中的专家输出置零。路由数学是精确的，但没有稀疏分发或专家并行。
MLA 中没有权重吸收。潜在表示被显式解压，这更易读，但推理更慢。
未包含：FP8 训练、DualPipe、节点受限路由、YaRN 上下文扩展，以及任何后训练。
规模。36.2M 参数 vs 671B，6 层 vs 61 层，256-token 上下文 vs 128K。

数据与许可

所有训练文本在美国均属公有领域：Dashiell Hammett 的 Red Harvest、The Dain Curse、The Maltese Falcon 和 Continental Op 故事（来自 Standard Ebooks 的 CC0 版本），以及来自 Project Gutenberg 的约 200 部侦探和犯罪小说。所有来源都列在 sources.txt 中。美国以外的公有领域状态可能不同。

参考

DeepSeek-AI，DeepSeek-V3 Technical Report，2024。

许可

代码：MIT。训练文本在美国属公有领域（见上文）。