190 行纯 JAX：DeepSeek V3 架构被压进一个文件

2026 年 11 月 14 日，Hacker News 上出现一个 Show HN：「DeepSeek V3 in ~200 lines of pure Jax」。点进去是 Vaibhaav Tiwari 的仓库 deepseek-v3-jax。标题里的 200 是个约数，README 自己写的是 190 行。我按 README 末尾给的命令数了一遍：deepseek_v3.py 去掉开头文档字符串后，非空非注释行正好 190。

这个差别值得先说清楚。190 行不是营销数字，它指的是模型本身——MLA、MoE、负载均衡、多 token 预测、损失函数、训练步和带缓存的生成，全在一个文件里。数据准备、训练循环、测试和画图都在别的文件，但那些本来也不算「架构」的一部分。

它凭什么算纯 JAX

requirements.txt 里和模型直接相关的是 jax==0.4.34、jaxlib、jax-metal 0.1.1、optax、numpy；没有 Flax，没有 Equinox，没有 PyTorch。参数是嵌套的普通 dict，数组就是 jax.numpy 数组。每个部件都是普通函数，从上往下读就行，不需要先学某个框架怎么组织状态。

这对读代码的人是好事。框架省掉的是样板，代价是你要先知道它怎么藏状态。这里没有那层，forward 就是把 embedding 取出来，一层层调 block，最后过 norm 和 head。想改结构，改的是显式的形状和 Config 字段，不用先找框架的注册表。

读代码的顺序

如果只想看架构，不用从第一行啃。先看 Config：所有形状都在这里，d_model、头数、MLA 的 q/kv 秩、专家数、top-k、序列长度一目了然。然后按 rms_norm → rope/apply_rope → mla → route/moe → block → forward → loss_fn → make_train_step → generate 的顺序走，基本就是一次前向、一次反向、一次解码。函数之间没有隐式状态，形状写在注释里；出错时先对形状，比翻框架源码快得多。

init_params 值得单独看。它把这个模型的参数布局讲清楚了：embedding 和 head 分开，每层各有一个 attention 子 dict；MoE 层的专家权重是 [E, d, h] 的三维数组，共享专家则是一个独立的 SwiGLU；MTP 模块各自带一个 block 和一个投影。看懂了这棵树，后面所有 p["..."] 都不会迷路。

190 行里的四个关键设计

MLA。 它不缓存每个头的完整 K/V，而是缓存一个压缩后的 latent。README 给的数字是每个 token 每层从 480 个数降到 144 个。解码时再把 latent 解压成各头的 K 和 V。decoupled RoPE 把位置那一小段单独留着，所有头共用一个 rotary key，不塞进压缩 latent。这个设计直接决定缓存大小，也是它能在 24 GB 机器上跑的原因之一。代码里 init_cache 只分配这两块，mla 在预填充阶段写入、解码阶段做 dynamic_update_slice，读起来比论文里的公式直白。

DeepSeekMoE。 1 个共享专家常开，16 个路由专家，每个 token 取 top-4。路由用 sigmoid 分数，选谁由分数加一个 per-expert bias 决定；bias 只改「选谁」，不改 gate 权重。训练中按负载更新 bias：过载的往下压，闲着的大往上抬。这样不用辅助损失也能把专家负载拉平。代码里 update_bias 只有一行 sign，但正是这一行把「均衡」从损失函数挪到了路由决策里，梯度不再为了均衡去改专家权重。

一个很小的平衡损失。 代码里仍保留 sequence-wise balance loss，α 取 1e-4，防止单条序列把专家用得太偏。它和前面的 bias 是两套机制，一个改路由选择，一个在损失里加软约束。

多 token 预测。 额外一个模块预测再下一个 token，损失里权重 λ 取 0.3。实现上把上一层的 hidden 和下一个位置的 embedding 拼起来，过一层投影和一个小 block，再共用输出头。训练和生成都支持：bf16 计算、f32 主权重，预填充后逐 token 解码。

跑出来是什么样

默认配置是 36.2M 参数、每个 token 激活 13.6M、6 层、256 token 上下文。作者拿它和 671B 参数、61 层、128K 上下文的 DeepSeek V3 对比，差距是数量级的——这是架构演示，不是复现。

数据是 207 本书、21.7M 训练 token。12,000 步在 M5、24 GB 的 MacBook Air 上约 2.6 小时，中位速度 10.7k token/s。验证损失从 5.02（第 250 步）降到 3.51。

最能说明问题的是负载均衡消融：开启 bias 时，2000 步每个专家负载 5.7–6.9%（完全均匀是 6.25%）；关掉后是 0.09–18.6%。而验证损失在 2000 步几乎一样：3.921 对 3.923。也就是说，均衡几乎是免费的，这和论文的说法一致。但生成质量别期待太高，README 自己说这是 architecture demo，不是强语言模型。

复现路径与坑

需要 Apple Silicon Mac、Python 3.12 和 uv，依赖版本全部钉死。jax-metal 0.1.1 只配 jax 0.4.34，别顺手升级。先跑测试：

JAX_PLATFORMS=cpu .venv/bin/pytest tests.py -q

README 写的是 39 passed；我在 Linux 的 CPU 上照跑，39 项同样全过。tests.py 覆盖因果掩码、缓存与全量前向一致、路由、均衡、MTP 目标、NaN 和 bf16 主权重，大多跑 TINY 配置，所以不依赖 Mac 的 GPU。smoke_metal.py 里有一项预期失败（grad through 2-axis contraction），模型刻意绕开了它。

数据准备会下载 207 本书，约 120 MB、7 分钟，训一个 8,192 词的 BPE tokenizer。它清洗 Gutenberg 文本、按书切分训练和验证，并把 Hammett 的作品重复 3 次，让模型多见一点硬汉派的对白。官方发布还给了 138 MB 的 checkpoint；README 提醒它是 Python pickle，加载等于执行代码，只该用来路可信的文件。没有 Mac 也能读和改，CPU 上跑 --config small（4 层、8 专家、top-2、128 token 上下文）先验证流程，只是慢。

训练脚本本身不神秘：AdamW、warmup 加 cosine 衰减、梯度裁剪，每次评估保存 checkpoint，指标写进 metrics.jsonl。种子固定，但 Metal 上的浮点运算不保证逐位可复现，作者也建议把结果当作「接近」而不是「一致」。

改一处结构：top-4 变 top-2

我选的是路由稀疏度，不动模型代码，只改 Config：

top_k: int = 4          # 改成 2
n_routed: int = 16

总参数仍是 36.2M，因为专家权重没变；变的是每个 token 激活几个路由专家，从 4 个降到 2 个，激活计算量随之下降。测试也仍然 39 项全过，因为路由测试跟着 TINY 配置走，形状和归一化断言不依赖默认值。

改完该看什么：每层专家负载是不是更不均、验证损失在第 2000 步差多少、以及它和 bias 开关的交互。README 本来就提供了 2000 步的对照跑法，照那个格式做一轮就行。这个旋钮不影响参数量，却直接改变每 token 的实际计算量和专家利用率；机器越小，越值得先试。反过来，把 top_k 调大不会自动让模型更聪明，只会更贵，除非你同步扩专家数并重新训练。另一个容易踩的点是：如果你改 MLA 的 kv_rank，测试里 test_mla_cache_size 把 144 和 480 写死了，记得同步改断言。这不是 bug，是把设计意图钉在测试里。

边界

• MoE 是 dense 计算：每个专家都跑，没选中的乘 0。路由数学精确，但没有稀疏 dispatch 和专家并行，所以「激活参数」只是账面数字，实际算力不一定省。
• MLA 没做 weight absorption，latent 显式解压，读起来清楚，推理更慢。
• 没实现 FP8 训练、DualPipe、node-limited routing、YaRN 和任何后训练。
• 190 行只覆盖模型；训练循环、数据清洗、tokenizer、测试和图表都在它之外。把它当教材，不要当可用的 LLM 底座。

如果你只想找一份能读懂的 DeepSeek V3 架构实现，这个仓库目前是成本最低的入口之一；如果你想拿它当训练框架，先接受它主动砍掉的那些东西。它不是「小号 DeepSeek」，而是一张能跑起来的结构图——知道每个数字从哪来、每个开关动哪里，比背论文里的公式有用。

本文基于以下公开报道整理。

参考来源

• Vaibhaav-Tiwari/deepseek-v3-jax — GitHub（2026-11-14）
• Show HN: DeepSeek V3 in ~200 lines of pure Jax — Hacker News（2026-11-14）

本文由 AI 辅助整理，事实来自项目 README、requirements.txt 与本地测试运行；训练结果引用作者公开记录，未在本机复现完整训练。