Vaibhaav-Tiwari/deepseek-v3-jax

项目介绍

这个仓库用 190 行非空非注释的纯 JAX 代码复现 DeepSeek-V3 技术报告中的核心架构，不依赖 Flax、Equinox 或 PyTorch。参数是普通数组字典，模型全部写成可自上而下阅读的函数，集中在 deepseek_v3.py。覆盖的内容包括：带解耦 RoPE 的多头潜在注意力（MLA）、DeepSeekMoE（1 个共享专家 + 16 个路由专家，每 token 选 top-4）、无辅助损失的偏置均衡、序列级均衡损失、多 token 预测（MTP），以及带缓存的 prefill/decode 生成。

已发布的训练是一个小规模预训练实验，不是可用的语言模型：36.2M 参数（每 token 激活 13.6M），6 层，256 token 上下文，在约 2170 万 token 的公有领域侦探小说上训练。作者称 12,000 步后验证损失从 5.02 降到 3.51，并称在第 2000 步时路由偏置让专家负载更均匀且不改变验证损失。配套的长文逐段解释了代码。

项目适合想阅读或修改一份紧凑、忠实还原 DeepSeek-V3 架构实现的人，也适合想找一个小型 JAX 训练示例的人。它明确是架构演示，不是聊天模型。

快速使用

环境要求：Apple Silicon Mac（已发布运行使用 24 GB 内存）、Python 3.12 和 uv。克隆并安装固定版本的依赖：

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

所有依赖版本都已固定；jax-metal 0.1.1 只能配 jax==0.4.34，不要升级 JAX。用 JAX_PLATFORMS=cpu .venv/bin/pytest tests.py -q 验证环境（39 个测试）。如果只想运行预训练检查点而不训练，从 v1.0 发布页下载检查点和分词器，再采样：

mkdir -p data runs/main
curl -L -o runs/main/ckpt.pkl  <发布页 ckpt.pkl 地址>
curl -L -o data/tokenizer.json <发布页 tokenizer.json 地址>
.venv/bin/python sample.py --prompt "The fog came in off the" --num 3

检查点是 Python pickle，加载时可能执行代码，只加载可信来源的文件。完整训练需先运行 prepare_data.py 下载约 120 MB 书籍并训练分词器，再运行 train.py（主训练约 2.6 小时）。具体的发布地址、数据哈希和消融命令见复现步骤。

使用约束

• 完整训练针对 Apple Silicon 上的 jax-metal 调优；CPU 回退可用 JAX_PLATFORMS=cpu，但 --config small 小配置只验证了几步，没有完整训练过。
• 架构做了刻意简化：MoE 采用稠密计算（每个专家都跑，未选中的置零），MLA 显式解压潜在向量而非权重吸收，不包含 FP8、DualPipe、节点受限路由、YaRN 上下文扩展和后训练。规模是 36.2M 参数，而真实模型是 671B。
• 重新运行 prepare_data.py 很可能得到不同的数据集哈希，因为 Project Gutenberg 文件无法固定版本；README 建议需要字节级复现时直接下载已发布数据。
• 训练文本仅在美国属于公有领域，其他地区可能不同。代码采用 MIT 许可证。

HotRepo 点评

项目在 2026-11-15 登上 GitHub Python 新项目榜，一小时内新增 9 星、12 个 fork。前一天的 Show HN 讨论获得 180 分和 9 条评论，有人称赞这种从第一性原理重建模型的做法（armcat），也有读者询问如何建立直觉（mattfrommars），回复推荐了 3blue1brown、Karpathy 教程、JAXformer 和 nanochat。还有评论指出使用别人用过的数据集更容易直接比较结果（nico），这对只有一次小规模训练的本项目同样适用。它的技术价值在于可读的参考实现和可复现的消融实验，而不是基准测试结果。