AI Radar — 每天的 AI 热点、模型与教程

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

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

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

190 行里有什么

论文中的思路	它做什么	在哪里
多头潜在注意力（MLA）	每个 token 只缓存一个小型潜在表示，而不是完整的键和值：480 → 144此处为每 token 每层的数字	`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`
训练与生成	bf16 计算，f32 主权重；使用 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 版本拥有已发布运行的精确 checkpoint 和数据。在安装（步骤 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 "雾从" --num 3

该 checkpoint 是一个 Python pickle，加载时可能执行代码。只加载你信任来源的 checkpoint。

复现它

你需要一台 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 项通过
.venv/bin/python smoke_metal.py                   #可选：jax-metal 在你的 Mac 上支持什么

在smoke_metal.py中，“通过 2 轴收缩求梯度”在 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    #应与上面的哈希匹配

4. 训练

.venv/bin/python train.py                                         #主运行：12k 步，约 2.6 小时
.venv/bin/python train.py --steps 2000 --out runs/bias_on         #平衡消融，带偏置
.venv/bin/python train.py --steps 2000 --out runs/bias_off --no-bias   #……以及不带偏置

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

5. 生成

.venv/bin/python sample.py --prompt "雾从" --num 3

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

6. 制作图表

.venv/bin/python plots.py                    #需要 runs/main、runs/bias_on、runs/bias_off
JAX_PLATFORMS=cpu .venv/bin/python charts.py #文章其余部分的图表

两者都写入post/figures/。

要在不训练的情况下重绘图表，请从发布版下载日志：metrics-main.jsonl、metrics-bias_on.jsonl和metrics-bias_off.jsonl放入runs/main/，runs/bias_on/和runs/bias_off/作为metrics.jsonl。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`	模型：配置、MLA、MoE、均衡、MTP、损失、训练步、带缓存的生成。那 190 行。
`train.py`	训练循环：带 warmup 和余弦衰减的 AdamW、评估、检查点、日志
`prepare_data.py`，`sources.txt`	下载、清洗并对语料分词
`sample.py`	从检查点生成文本
`tests.py`	39 个测试：因果掩码、缓存与完整前向对比、路由、均衡、MTP 目标、NaN 处理、bf16 主权重等
`smoke_metal.py`	哪些 JAX 操作能在 Apple 的 GPU 上运行
`plots.py`，`charts.py`	图表，全部取自日志、检查点和语料
`post/`	文章中的示意图、图表和截图

“190 行”统计的是以下文件去掉 docstring 后的非空、非注释行：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 系列故事（CC0 版本来自Standard Ebooks），外加约 200 部来自Project Gutenberg的侦探和犯罪小说。所有来源都列在sources.txt中。美国以外的公有领域状态可能不同。

参考

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

许可

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

Show HN：190 行纯 JAX 复现 DeepSeek V3 架构

190 行里有什么

结果

试用训练好的模型

复现它

在 CPU 上运行

文件

简化了什么

数据与许可

参考

许可

1 个独立信源

相关条目