跳到主要内容

数据截至 (上游 commit 92d63d4e8bb4)

nanochat — 架构与原理

30 秒导读： nanochat 是 Karpathy 的「最小全栈 ChatGPT 复刻」：一个仓库、一台 8×H100、约 1.5 小时、约 $48，从原始文本走完 tokenizer → 预训练 → SFT → RL → CLI 对话的全流程，训出一个 GPT-2 级别的聊天模型。它是目前读「一个 chat 模型到底是怎么造出来的」最完整的可读参照。

1. 这是什么（零基础也能懂）​

一句话定义​

nanochat 是一个单机可跑通的 LLM 全管线实验台：从下载原始文本、训分词器开始，到预训练一个 GPT，再做监督微调（SFT）和强化学习（RL），最后用带 KV cache 的推理引擎在命令行和你聊天——每一步都是这个仓库里几百行可读的 PyTorch。

解决什么问题 / 给谁用​

学 LLM 的人面对的通常是一堆「各管一段」的库：训分词器的不管训练，训练的不管对齐，对齐的不管推理。nanochat 把整条链放进同一个仓库，回答三个问题：

  • 从原始文本到 chat 模型，完整管线长什么样，一步不藏。
  • SFT 相对预训练在什么位置，RL 又在 SFT 之上加了什么。
  • 每个阶段实际要多少算力、多少代码。

它面向想端到端读懂并 hack 的研究者与工程师，而不是想要 SOTA 模型或万能框架的用户。README 原话：为「一个人的可读性」优化，不为 SOTA 优化。

它能做什么​

阶段脚本干什么
分词scripts/tok_train.py用 rustbpe 在 ~2B 字符上训 32K BPE 分词器
预训练scripts/base_train.py训 base GPT，--depth 一个旋钮定全部规模超参
基座评测scripts/base_eval.pyCORE 分数（DCLM 论文）+ bits per byte
SFTscripts/chat_sft.py在对话混合数据上教模型聊天、多选、工具调用
RLscripts/chat_rl.py在 GSM8K 上做 on-policy 策略梯度
对话评测scripts/chat_eval.pyARC / MMLU / GSM8K / HumanEval + ChatCORE
推理scripts/chat_cli.pyKV cache 引擎对话，内置计算器工具

用起来什么样​

整个 GPT-2 复现就装在一个 shell 脚本（runs/speedrun.sh）里，核心四步：

python -m scripts.tok_train                                             # 1) 训分词器
torchrun --nproc_per_node=8 -m scripts.base_train -- --depth=24 --fp8   # 2) 预训练
torchrun --nproc_per_node=8 -m scripts.chat_sft                         # 3) SFT
python -m scripts.chat_cli                                              # 4) 聊天

然后就能对话（README 实录，speedrun 模型）：

User: Why is the sky blue?

Assistant: The sky is blue due to an optical illusion called the Rayleigh Scattering of Light, which occurs when sunlight passes through tiny particles in the air...

README 自嘲这个 4e19 FLOPs 的模型「有点像和幼儿园小孩聊天」——它会幻觉，但管线是真的。

一句话直觉 / 类比​

把 nanochat 想成一本可以直接执行的教材：兄弟项目 nanoGPT 只写了「预训练」一章，nanochat 把整本书写完——分词、预训练、后训练、推理，每章只有几百行。

2. 顶层全景（它大概怎么转）​

一张顶层图​

怎么读：从左到右是时间顺序——资产（数据、分词器）在左，训练阶段居中，产物在右。①~⑤ 是五个阶段，与下面的部件表一一对应。

ClimbMix-400B parquet
        │
        ▼
① 分词 tok_train         ② 预训练 base_train        base ckpt
   rustbpe→tiktoken  ──►    GPT + MuonAdamW    ──►  (base_checkpoints/d24)
        │                    DDP, 可选 FP8              │
        ▼                                               ▼
   tokenizer.pkl      ③ SFT chat_sft                 sft ckpt
   (ids+mask 渲染) ──►   任务混合 + loss mask   ──►  (chatsft_checkpoints)
                                                          │
                      ④ RL chat_rl                        ▼
                         采样→优势→策略梯度  ──►       rl ckpt
                                                          │
                      ⑤ Engine + chat_cli                 ▼
                         KVCache + 工具状态机  ──►     命令行对话

五个阶段共享的底座（不在图里）：nanochat/common.py 的设备/精度/DDP 初始化，nanochat/checkpoint_manager.py 的 checkpoint 存取。

部件一句话职责​

部件干什么在哪个文件
分词器rustbpe 训练、tiktoken 推理；渲染对话成 (ids, mask)nanochat/tokenizer.py
GPT 模型decoder-only Transformer（rotary / QK norm / GQA / 滑窗 / value embeds）nanochat/gpt.py
优化器Muon（矩阵）+ AdamW（其余），ZeRO-2 式状态分片nanochat/optim.py
注意力后端FA3 优先、SDPA 兜底，同一 APInanochat/flash_attention.py
FP8~150 行 Float8Linear，替代 torchaonanochat/fp8.py
预训练数据ClimbMix parquet 下载与读取nanochat/dataset.py
数据装载BOS 对齐 best-fit packing，100% 利用nanochat/dataloader.py
推理引擎KVCache + 采样 + 工具调用状态机nanochat/engine.py
代码沙箱子进程执行 LLM 生成的 Pythonnanochat/execution.py
checkpoint按阶段分目录存取模型/优化器/元数据nanochat/checkpoint_manager.py
任务ARC/MMLU/GSM8K/HumanEval/SmolTalk 统一接口tasks/*.py
指标bpb、CORE、ChatCOREnanochat/loss_eval.py、nanochat/core_eval.py、scripts/chat_eval.py

主线走一遍（高层）​

一条数据从硬盘到对话的旅程：

  1. 下载：nanochat/dataset.py 把 ClimbMix-400B 的 parquet shard 拉到 ~/.cache/nanochat，最后一个 shard 固定作验证集。
  2. 分词：scripts/tok_train.py 用 rustbpe 训出合并表，包成 tiktoken Encoding 存盘（nanochat/tokenizer.py:42 train_from_iterator）。
  3. 预训练：scripts/base_train.py 按 --depth 推出宽度、批大小、学习率、步数，用 MuonAdamW 训 base GPT，产出 base_checkpoints/d{depth}。
  4. SFT：scripts/chat_sft.py 加载 base checkpoint，在 SmolTalk + MMLU + GSM8K 混合对话上训练，只对 assistant 的 token 算 loss。
  5. RL：scripts/chat_rl.py 加载 sft checkpoint，在 GSM8K 上每题采 16 个答案，用「奖励减均值」当优势做策略梯度。
  6. 对话：scripts/chat_cli.py 加载 sft/rl checkpoint，用 Engine（KV cache）流式生成；模型写出 <|python_start|>...<|python_end|> 时，引擎本地算出结果塞回生成流。

目标：先看懂大盘。细节下钻到各章。

3. 阅读地图​

按由浅入深排序，建议通读；只想看某一段可直接跳：

章讲什么什么时候读
01 · 分词器与对话渲染rustbpe/tiktoken 双拼；(ids, mask) 是怎么来的想懂 SFT 数据协议之前必读
02 · GPT 模型与预训练模型现代配方；--depth 单旋钮；MuonAdamW；BOS 对齐装载；FP8全库最重的一章
03 · 后训练链SFT 混合与 mask；退化成 REINFORCE 的 RL；ChatCORE关心对齐怎么做
04 · Engine 推理KVCache；一次 prefill 克隆 N 份；工具状态机；沙箱与基准关心推理 / serving

前置知识：知道「Transformer 是什么」即可；各章自带直觉，不预设你读过 nanoGPT。

4. 巧妙之处（可借鉴的技术）​

每条先白话点出妙处，再上锚点；细节在对应章节展开。

  • --depth 单旋钮：宽度 = depth × 64 再对齐 head_dim，训练 token 数 = 12 × 参数量，批大小按 D^0.383 幂律外推、学习率按 √B 缩放、weight decay 按 T_epoch 恒定外推——用户只选「要多大的模型」（scripts/base_train.py:129-140、263-302）。
  • MuonAdamW 一个类通吃单机与分布式：无进程组时通信全部退化为本地操作；有进程组时优化器在 step() 内部自己做梯度 reduce_scatter + 参数 all_gather（不包 DDP），优化器状态按 ZeRO-2 分片（nanochat/optim.py:184-259 MuonAdamW）。
  • BOS 对齐 best-fit packing：每行必以 BOS 开头，「能整个放下的最大文档」优先装入，装不下才裁剪——用约 35% 的裁剪浪费换「每个 token 都能 attend 回文档开头」（nanochat/dataloader.py:74-161）。
  • 训练用 rustbpe、推理用 tiktoken：训练完把合并表直接倒进 tiktoken.Encoding，白嫖 tiktoken 的 Rust 高速编码；特殊 token 不参与训练、按偏移量续号（nanochat/tokenizer.py:42-58）。
  • loss mask 与对话渲染一体：render_conversation 直接吐出 mask（user/工具输出 = 0，assistant = 1），SFT 把它平移一位写成 -1 的 ignore_index——监督信号的定义就在分词器里（nanochat/tokenizer.py:140、scripts/chat_sft.py:288-290）。
  • 一次 prefill、克隆 KV cache 采 N 份：Engine 先 batch=1 预填 prompt，再把 cache 复制成 num_samples 份并行采样，RL 的「一题 16 答」不必重复 prefill（nanochat/engine.py:176-218）。
  • FA3 / SDPA 同一接口：用 SimpleNamespace 导出 flash_attn_func / flash_attn_with_kvcache，有 FA3 用 FA3，没有就 SDPA 手工维护 cache——上层零分支（nanochat/flash_attention.py:191-195）。
  • 150 行 FP8：不用 torchao 的 tensor subclass 派发表，直接一个 autograd.Function 包住 torch._scaled_mm，量化/反量化全在里面（nanochat/fp8.py:125 _Float8Matmul）。

5. 边界与局限​

诚实清单——它刻意不做什么、会在哪崩：

  • 单机边界：为单节点 8 卡设计；没有张量并行/流水并行，多机要自己改（优化器只做数据并行 + ZeRO-2）。
  • 质量非目标：README 明言为「一个人的可读性」优化而非 SOTA；产出是「GPT-2 级别」的模型，会幻觉、会一本正经地胡说。
  • 沙箱非安全边界：execution.py 文件头自己列出「不防」清单——不挡网络、不防 ctypes、无内核级隔离，只防意外不防恶意（nanochat/execution.py:14-21）。
  • 工具只有计算器：生成期只识别 <|python_start|> 块，且 use_calculator 白名单只放行纯算术和字符串 .count()（nanochat/engine.py:46-79）。
  • RL 只示范在 GSM8K 上，且 fp16 不支持 RL（README 的 dtype 节：「SFT supports this too but RL currently does not」）。
  • 无 FA3 时性能打折：SDPA 兜底路径不支持滑窗注意力，长上下文利用率会很差，训练脚本会刷屏警告（scripts/base_train.py:107-117）；Blackwell（sm100）目前也走兜底（nanochat/flash_attention.py:29-31）。
  • 数据集换代有迁移成本：2026-03-04 起预训练集从 FinewebEdu-100B 换成 ClimbMix-400B，旧缓存会触发警告并回退（nanochat/dataset.py:38-58）。
  • 快动仓库：源卡片注明 fast-moving——本库所有引用 as-of frontmatter 的 sourceCommit，行号可能随上游漂移。

6. 横向对比​

同书架兄弟项目的不同取舍：

  • nanoGPT：同一作者的前作，只覆盖「预训练」一章（model.py + train.py 各约 300 行）。nanochat 是它的「续完全书」：模型配方现代化（rotary / QK norm / GQA / 滑窗 / value embeds，对 nanoGPT 的学习式位置编码 + GELU + 权重绑定），并补齐分词、SFT、RL、KV-cache 推理。读法：先 nanoGPT 建立基准线，再看 nanochat 每一处「为什么变复杂了」。
  • minbpe：同一作者的纯 Python 教学版 BPE。nanochat 的分词器是它的「生产化」：训练换成 Rust（rustbpe）、推理换成 tiktoken，外加 9 个对话特殊 token 和 (ids, mask) 渲染协议。
  • verl：工业级 RL 后训练框架（单控制器 + 混合引擎，管 4 个模型 × 上百卡）。nanochat 的 RL 是它的极简对偶：一个模型、一份脚本、on-policy、无 PPO clip、无 KL。先看 nanochat 理解「RL 训练的本质三步：采样 → 算优势 → 加权梯度」，再看 verl 理解规模化要付出多少基础设施。

7. 代码地图（导航索引）​

主题文件路径符号名
分词器训练入口scripts/tok_train.pyRustBPETokenizer.train_from_iterator
双拼分词器nanochat/tokenizer.pyRustBPETokenizer、SPECIAL_TOKENS、SPLIT_PATTERN
对话渲染 + masknanochat/tokenizer.py:140render_conversation
RL/评测提示渲染nanochat/tokenizer.py:241render_for_completion
GPT 配置与模型nanochat/gpt.pyGPTConfig、GPT、CausalSelfAttention
初始化方案nanochat/gpt.py:204GPT.init_weights
优化器组装nanochat/gpt.py:419GPT.setup_optimizer
Muon / AdamW 融合核nanochat/optim.pymuon_step_fused、adamw_step_fused、MuonAdamW
FA3 / SDPA 切换nanochat/flash_attention.pyflash_attn_func、flash_attn_with_kvcache
FP8nanochat/fp8.py_Float8Matmul、Float8Linear
预训练主循环scripts/base_train.pybuild_model_meta、get_lr_multiplier
BOS 对齐装载nanochat/dataloader.py:74tokenizing_distributed_data_loader_with_state_bos_bestfit
SFT 主循环scripts/chat_sft.py:180sft_data_generator_bos_bestfit
RL 主循环scripts/chat_rl.pyget_batch、run_gsm8k_eval
推理引擎nanochat/engine.pyEngine、KVCache、RowState
计算器工具nanochat/engine.py:46use_calculator
代码沙箱nanochat/execution.pyexecute_code、GUARD
checkpoint 存取nanochat/checkpoint_manager.pysave_checkpoint、load_model、build_model
任务接口tasks/common.pyTask、TaskMixture、TaskSequence、load_hub_dataset
CORE / bpb 指标nanochat/core_eval.py、nanochat/loss_eval.pyevaluate_task、evaluate_bpb
对话评测scripts/chat_eval.pyrun_chat_eval