从文本到聊天模型

本文面向希望系统理解 GPT 类模型训练过程的读者。目标不是只解释 Transformer 的结构，而是沿着 nanochat 仓库的真实执行路径，串起数据准备、分词器训练、基础模型预训练、评估、监督微调、强化学习和推理服务。

如果你原本想学习经典 nanoGPT，需要先注意：当前仓库是 nanochat。它保留了 GPT 预训练的核心思想，同时加入了更完整的端到端流程。因此，它非常适合作为理解“小型 ChatGPT 如何从零训练出来”的源码案例。

目录

1. 项目全景
2. 数据集准备
3. 训练 BPE 分词器
4. DataLoader 如何构造训练样本
5. GPT 模型如何完成前向传播
6. 预训练循环与参数更新
7. 如何评估基础模型
8. SFT 如何将基础模型变成聊天模型
9. RL 如何强化数学能力
10. 推理、KV Cache 与聊天 UI
11. 推荐的源码阅读顺序
12. 常用命令与产物目录

1. 项目全景

端到端入口是 runs/speedrun.sh。它设计为在一个 8×H100 GPU 节点上训练达到 GPT-2 能力水平的模型。

完整链路如下：

ClimbMix 原始文本
  → 下载 Parquet 分片
  → 训练 BPE 分词器
  → 构造预训练 token batch
  → GPT 基础模型预训练
  → BPB / CORE 评估
  → SFT 对话微调
  → ChatCORE 评估
  → CLI / Web UI 推理

可选扩展：
SFT 模型 → GSM8K 强化学习 → RL 模型

在阅读具体实现前，先建立三个基本认识：

1. 预训练、SFT 和 RL 使用的是同一个 GPT 模型。 不同阶段主要改变训练数据、损失位置和优化目标。
2. GPT 的基础任务始终是预测下一个 token。 聊天能力不是另一个网络，而是通过特殊 token 和对话数据训练出来的行为。
3. runs/speedrun.sh 是主线，RL 是可选阶段。 默认竞速脚本覆盖分词器、预训练、SFT、评估和聊天，但不会执行 RL。

2. 数据集准备

2.1 训练入口如何下载数据

runs/speedrun.sh 先下载少量分片训练分词器，同时在后台继续下载预训练所需的数据：

python -m nanochat.dataset -n 8
python -m nanochat.dataset -n 170 &
DATASET_DOWNLOAD_PID=$!
python -m scripts.tok_train
python -m scripts.tok_eval
wait $DATASET_DOWNLOAD_PID

这是一项直接有效的工程优化：

• 前 8 个分片足以提供约 20 亿字符，用于训练分词器。
• 后台继续下载约 170 个训练分片。
• 分词器训练和数据下载并行进行，减少 GPU 节点的等待时间。

2.2 数据从哪里来

运行时下载逻辑位于 nanochat/dataset.py。

当前数据集地址是：

https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle

本地文件类似：

shard_00000.parquet
shard_00001.parquet
...
shard_06542.parquet

每个分片约包含 2.5 亿字符，使用 zstd 压缩后约为 100MB。

2.3 数据最初如何制作

参考脚本是 dev/repackage_data_reference.py。该脚本用于记录数据制作方式，不会在日常训练时执行。

处理链路：

NVIDIA Nemotron-ClimbMix
  → 将原始 token 解码为文本
  → 使用固定随机种子打乱文档
  → 每约 2.5 亿字符生成一个分片
  → 按 row group 写入 Parquet
  → 使用 zstd 压缩
  → 上传到 Hugging Face

分片的意义不仅是节省磁盘空间。它还允许训练过程按需下载数据、多 GPU 并行读取不同 row group，并在中断后近似恢复读取位置。

2.4 训练集与验证集划分

nanochat/dataset.py 使用非常直接的规则：

parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:]

• 前面的分片全部作为训练集。
• 最后一个分片固定作为验证集。
• 即使只下载 8 个训练分片，程序也会额外下载固定验证分片。

固定验证集非常重要。不同实验必须在相同数据上评估，验证损失才具有可比性。

3. 训练 BPE 分词器

模型不能直接处理字符串，只能接收整数 token ID。因此需要将文本：

The capital of France is Paris.

转换为类似：

[1234, 5872, 315, 9012, 318, 14882, 13]

3.1 分词器训练入口

入口是 scripts/tok_train.py。

默认参数：

--max-chars=2_000_000_000
--doc-cap=10_000
--vocab-size=32768

含义：

• 最多使用 20 亿字符训练分词器。
• 每篇文档最多使用前 10,000 个字符，防止少数超长文档占比过大。
• 最终词表包含 32,768 个 token。

3.2 BPE 的基本思想

按字符切分会导致序列过长：

training → t r a i n i n g

按单词切分则难以处理词表外的新单词：

training → training

BPE 在二者之间取平衡。它从基础字节开始，反复合并语料中高频出现的相邻片段：

t r a i n i n g
→ t r a in in g
→ t r a ining
→ training

常见片段使用更少 token 表示，不常见文本仍然可以拆成更小字节片段。因此任何 UTF-8 文本都能够编码。

3.3 nanochat 的实现

核心实现位于 nanochat/tokenizer.py 的 RustBPETokenizer。

tokenizer = rustbpe.Tokenizer()
tokenizer.train_from_iterator(
    text_iterator,
    vocab_size_no_special,
    pattern=SPLIT_PATTERN,
)

项目组合了两个库：

• rustbpe：快速训练 BPE 合并规则。
• tiktoken：训练完成后高效执行编码和解码。

3.4 预切分规则

BPE 不会直接在整篇文档上任意合并。它首先使用正则表达式 SPLIT_PATTERN 切分文本，再在各片段内部训练合并规则。

它会分别处理：

• 单词
• 空格
• 标点符号
• 换行
• 数字
• 英文缩写

数字最多每两个字符为一组：

123456 → 12 | 34 | 56

这是针对 32K 小词表的调整，用于避免大量数字组合占据词表空间。

3.5 特殊 token

nanochat/tokenizer.py 预留了以下特殊 token：

<|bos|>
<|user_start|>
<|user_end|>
<|assistant_start|>
<|assistant_end|>
<|python_start|>
<|python_end|>
<|output_start|>
<|output_end|>

其中：

• <|bos|> 在预训练阶段使用，表示新文档开始。
• 其余 token 主要在 SFT 和推理阶段使用，用于表达对话边界和工具调用。

聊天内容最终会表示为：

<|bos|>
<|user_start|>你好<|user_end|>
<|assistant_start|>你好，请问有什么需要帮助？<|assistant_end|>

3.6 输出文件

训练完成后，分词器保存在：

~/.cache/nanochat/tokenizer/tokenizer.pkl
~/.cache/nanochat/tokenizer/token_bytes.pt

• tokenizer.pkl 保存词表和 BPE 合并规则。
• token_bytes.pt 保存每个 token 对应的 UTF-8 字节数。

缓存根目录默认为 ~/.cache/nanochat，也可以通过 NANOCHAT_BASE_DIR 修改。

3.7 为什么记录 token 字节数

不同分词器会将同一段文本拆成不同数量的 token。直接比较平均 token loss 并不公平。

nanochat 使用 BPB（Bits Per Byte，每字节位数）：

BPB = 总负对数似然 / (ln(2) × 总字节数)

BPB 越低，模型预测文本的能力越强。特殊 token 的字节数设置为 0，不会计入 BPB。

scripts/tok_eval.py 还会将当前分词器与 GPT-2、GPT-4 分词器比较，观察新闻、代码、数学、科学文本、非英文文本和数据集文本上的压缩比。

4. DataLoader 如何构造训练样本

分词器将文档转换为长短不同的 token ID 序列。GPU 训练需要规则的矩阵，因此 DataLoader 还要将文档整理为固定长度 batch。

核心代码位于 nanochat/dataloader.py。

4.1 输入与标签错开一位

假设 token 序列为：

<|bos|> The sky is blue

对应 token ID：

[0, 41, 72, 19, 96]

DataLoader 构造：

inputs  = [0, 41, 72, 19]
targets = [41, 72, 19, 96]

模型在每个位置根据前文预测下一个 token：

实现非常简单：

cpu_inputs.copy_(row_buffer[:, :-1])
cpu_targets.copy_(row_buffer[:, 1:])

如果上下文长度 T=2048，DataLoader 需要先构造长度为 T+1=2049 的序列。

4.2 每篇文档前添加 BOS

每篇文档编码时都会添加 <|bos|>：

token_lists = tokenizer.encode(doc_batch, prepend=bos_token)

例如：

<|bos|> Paris is in France.
<|bos|> Python is a language.

BOS 告诉模型新文档已经开始，避免模型错误地将两篇无关文档理解为一段连续文本。

4.3 Best-fit 装箱

预训练 DataLoader 使用 best-fit 算法，将不同长度文档尽量完整地装入固定容量序列。

假设每行容量为 10，缓冲区中有：

文档 A：6 tokens
文档 B：4 tokens
文档 C：3 tokens

算法优先选取可以完整放入剩余空间的最长文档：

行 1：[文档 A 6 tokens][文档 B 4 tokens]

4.4 剩余空间不足时裁剪

如果剩余空间为 2，但最短文档也有 3 个 token，预训练 DataLoader 会裁剪最短文档：

写入：文档 C 的前 2 个 token
丢弃：文档 C 的剩余 token

这种设计有明确取舍：

• 优点：没有 padding，每个位置都参与训练，GPU 利用率高。
• 缺点：默认 T=2048 时，大约 35% 的 token 会因裁剪丢弃。

预训练语料极其充足，因此项目选择吞吐量优先。

4.5 多 GPU 数据分配

运行：

torchrun --standalone --nproc_per_node=8 -m scripts.base_train

会启动 8 个进程。每块 GPU 对应一个 rank，按步长读取不同 Parquet row group：

GPU 0：0, 8, 16, ...
GPU 1：1, 9, 17, ...
GPU 2：2, 10, 18, ...
...
GPU 7：7, 15, 23, ...

每块 GPU 处理不同数据，避免重复训练。

4.6 Batch 与梯度累积

预训练脚本计算：

tokens_per_fwdbwd = device_batch_size * max_seq_len
world_tokens_per_fwdbwd = tokens_per_fwdbwd * ddp_world_size
grad_accum_steps = total_batch_size // world_tokens_per_fwdbwd

以竞速配置为例：

device_batch_size = 16
max_seq_len       = 2048
GPU 数量          = 8

每块 GPU 一次前向和反向传播处理：

16 × 2048 = 32,768 tokens

8 块 GPU 合计：

32,768 × 8 = 262,144 tokens

如果总 batch size 是 524,288 tokens，则需要累计两次梯度：

524,288 / 262,144 = 2

梯度累积允许显存较小的设备模拟更大的 batch。

5. GPT 模型如何完成前向传播

模型核心位于 nanochat/gpt.py。

首先掌握主干：

token ID
  → Embedding
  → 多层 Transformer Block
  → LM Head
  → 每个位置的下一个 token 概率
  → 与 targets 计算交叉熵

5.1 模型配置

竞速脚本使用 --depth=24。模型宽度自动计算：

model_dim = depth * aspect_ratio
num_heads = model_dim // head_dim

默认情况下：

depth        = 24
aspect_ratio = 64
head_dim     = 128

model_dim = 24 × 64 = 1536
num_heads = 1536 / 128 = 12

nanochat 的核心设计是：用户主要调整一个复杂度旋钮 --depth，其他参数尽量自动推导。

5.2 Token Embedding

DataLoader 输入：

idx.shape = (B, T)

Embedding 查询：

x = self.transformer.wte(idx)

输出：

x.shape = (B, T, C)

对于 d24：

(16, 2048) → (16, 2048, 1536)

Embedding 将离散 token 转换为可训练的连续向量。

5.3 Transformer Block

每个 Block 包含：

x = x + self.attn(norm(x), ...)
x = x + self.mlp(norm(x))

结构：

输入 x
  → RMSNorm
  → Causal Self-Attention
  → 残差连接
  → RMSNorm
  → MLP
  → 残差连接

残差连接让信息和梯度更容易穿过深层网络。

5.4 Attention

Attention 首先映射出 Q、K、V：

q = self.c_q(x)
k = self.c_k(x)
v = self.c_v(x)

形状：

Q: (B, T, H, D)
K: (B, T, Hkv, D)
V: (B, T, Hkv, D)

直观理解：

• Q：当前位置想寻找什么信息。
• K：每个历史位置提供什么索引。
• V：每个历史位置真正携带什么内容。

核心计算可以简化为：

score = Q × Kᵀ
weight = softmax(score)
output = weight × V

5.5 Causal Attention

语言模型不能偷看未来。预测 is 时，模型只能看到：

The sky

不能提前看到后面的 blue。

nanochat 通过因果掩码实现：

flash_attn.flash_attn_func(q, k, v, causal=True, ...)

5.6 RoPE 位置信息

模型还需要区分：

dog bites man
man bites dog

nanochat 使用 RoPE（Rotary Position Embedding）。它根据位置旋转 Q 和 K 向量，让 Attention 分数能够感知相对位置。

5.7 MLP

Attention 负责 token 之间的信息交换。MLP 负责处理当前位置内部的信息：

x = self.c_fc(x)       # C → 4C
x = F.relu(x).square()
x = self.c_proj(x)     # 4C → C

可以概括为：

Attention：从其他 token 获取信息
MLP：处理当前位置已经获得的信息

5.8 LM Head 与损失

经过全部 Transformer Block 后：

logits = self.lm_head(x)

形状变化：

(B, T, C) → (B, T, vocab_size)

每个位置都会输出整个词表的分数，再与正确 target 计算交叉熵：

loss = F.cross_entropy(
    logits.view(-1, logits.size(-1)),
    targets.view(-1),
    ignore_index=-1,
)

5.9 经典 GPT 与 nanochat 竞速优化

第一次阅读时，优先理解：

Embedding
Transformer Block
Causal Self-Attention
RoPE
MLP
Residual Connection
LM Head
Cross Entropy

nanochat 还加入了多项竞速优化：

后几项属于实验优化，不是理解 GPT 原理的前置条件。

6. 预训练循环与参数更新

入口位于 scripts/base_train.py。

6.1 自动决定训练 token 数

nanochat 默认不会直接写死训练步数，而是根据模型规模估算训练 token 数：

target_tokens = target_param_data_ratio * num_scaling_params
num_iterations = target_tokens // total_batch_size

当前代码默认：

--target-param-data-ratio=12

竞速脚本为了更快越过 GPT-2 阈值，使用：

--target-param-data-ratio=8

模型越大，参数越多，需要看到的训练 token 通常也越多。

6.2 自动选择总 batch size

nanochat 根据经验公式估算适合的 batch size：

Bopt ∝ D^0.383

其中：

• Bopt：最优 batch size。
• D：训练 token 数。

最终 batch size 会取最接近的 2 的幂，便于高效计算。

6.3 学习率调度

学习率经历三个阶段：

warmup → constant → warmdown

默认配置：

warmup_steps   = 40
warmdown_ratio = 0.65
final_lr_frac  = 0.05

含义：

• 初期逐步提高学习率，避免随机初始化阶段更新过猛。
• 中期保持较大学习率，快速学习。
• 后期降低学习率，细化参数并稳定收敛。

6.4 AdamW 与 Muon 分工

nanochat 同时使用两类优化器：

AdamW 的核心过程：

记录梯度的一阶动量
记录梯度平方的二阶动量
自适应调整更新幅度
应用 weight decay
更新参数

Muon 的核心过程：

梯度
  → 动量
  → 对更新矩阵近似正交化
  → 方差归一化
  → 更新参数

第一次学习时，只需掌握：

优化器读取梯度，并小幅修改参数，使下一次 loss 尽可能降低。

6.5 一个完整训练 step

核心循环：

for micro_step in range(grad_accum_steps):
    loss = model(x, y)
    loss = loss / grad_accum_steps
    loss.backward()
    x, y, dataloader_state_dict = next(train_loader)

optimizer.step()
model.zero_grad(set_to_none=True)

完整过程：

1. DataLoader 提供 inputs 和 targets
2. GPT 前向传播，得到 loss
3. loss.backward() 计算每个参数的梯度
4. 累积多个 micro-batch 的梯度
5. optimizer.step() 更新参数
6. 清空梯度
7. 进入下一轮

需要注意：

• backward() 只计算梯度。
• 真正修改参数的是 optimizer.step()。

6.6 梯度是什么

梯度表示：

某个参数变化一点点，会让 loss 如何变化？

简化更新公式：

新参数 = 旧参数 - 学习率 × 梯度

AdamW 和 Muon 会对该公式进行更复杂的调整，但基本方向一致。

6.7 训练期间的检查

预训练会周期性执行：

日志还会显示：

loss
step
tok/sec
MFU
训练耗时
预计剩余时间

MFU 是 Model FLOPS Utilization，表示 GPU 理论算力的利用程度。

6.8 Checkpoint

默认保存目录：

~/.cache/nanochat/base_checkpoints/d24/

文件类似：

model_001234.pt
meta_001234.json
optim_001234_rank0.pt
optim_001234_rank1.pt
...

多 GPU 场景使用 DistMuonAdamW。它在优化器内部同步梯度和参数，采用类似 ZeRO-2 的方式切分优化器状态：

各 GPU 完成 backward
  → reduce_scatter 聚合并切分梯度
  → 每块 GPU 更新自己负责的部分
  → all_gather 收集更新后的参数
  → 所有 GPU 获得一致模型

7. 如何评估基础模型

训练完成后，默认执行：

torchrun --standalone --nproc_per_node=8 \
  -m scripts.base_eval -- --device-batch-size=16

入口位于 scripts/base_eval.py。

默认执行：

sample → bpb → core

7.1 文本采样

模型会续写固定提示词：

The capital of France is
The chemical symbol of gold is
The opposite of hot is
If 5*x + 3 = 13, then x is

固定提示词使用：

temperature=0

即每一步选择概率最高的 token。

采样适合发现明显异常：

• 输出完全乱码。
• 输出陷入重复。
• 模型只会输出极少数 token。
• 训练过程中能力没有改善。

但采样具有主观性，不适合作为唯一指标。

7.2 BPB

BPB 是 Bits Per Byte：

BPB = 总负对数似然 / (ln(2) × 文本总字节数)

项目分别在训练集和验证集计算 BPB：

train BPB 下降：模型更会拟合训练文本
val BPB 下降：模型对未见文本的预测能力改善

验证集 BPB 通常平滑稳定，适合比较局部实验效果。

但跨数据集比较 BPB 要谨慎。数据分布变化后，BPB 数字不再具有直接可比性。

7.3 CORE

CORE 来自 DCLM 论文，是排行榜的主要能力指标。GPT-2 阈值为：

CORE = 0.256525

项目目标是：

在 8×H100 上，用尽可能短的时间使 CORE > 0.256525

7.4 CORE 如何评估选择题

假设问题是：

法国的首都是？
A. Paris
B. London
C. Tokyo

基础模型不是聊天模型，因此评估器分别构造候选续写：

法国的首都是 Paris
法国的首都是 London
法国的首都是 Tokyo

然后计算模型对每个候选答案的平均 loss：

Paris  loss = 0.8  ← 选择
London loss = 2.4
Tokyo  loss = 3.1

loss 最低的候选续写最符合模型学到的语言分布，因此视为模型答案。

7.5 Few-shot 与中心化

CORE 支持 few-shot 提示，即在正式问题前提供少量示例，让模型仅通过上下文理解任务格式。

不同任务的随机准确率不同：

二选一：约 50%
四选一：约 25%

项目会减去随机基线并归一化，再对所有任务取平均值，得到 CORE。

7.6 三类指标如何配合

实际研究中通常：

先观察 val BPB 是否改善
  → 再确认 CORE 是否提升
  → 最后查看生成样本是否正常

8. SFT 如何将基础模型变成聊天模型

SFT 是 Supervised Fine-Tuning，中文通常称为监督微调。

入口位于 scripts/chat_sft.py。

8.1 SFT 没有更换模型结构

SFT 直接加载预训练模型：

model, tokenizer, meta = load_model("base", device, phase="train")

GPT 架构、分词器和损失函数都不变：

inputs → GPT → logits → cross entropy → backward → optimizer.step()

变化的是训练数据和参与 loss 计算的位置。

8.2 只训练助手输出

SFT 对话：

<|bos|>
<|user_start|>What is the capital of France?<|user_end|>
<|assistant_start|>Paris.<|assistant_end|>

模型仍然执行 next-token prediction，但只在助手回复部分计算 loss：

token                                    mask
<|bos|>                                  0
<|user_start|>                           0
What is the capital of France?           0
<|user_end|>                             0
<|assistant_start|>                      0
Paris.                                   1
<|assistant_end|>                        1

SFT DataLoader 将 mask=0 的 target 修改为 -1：

targets[mask_targets == 0] = -1

GPT 的交叉熵使用：

ignore_index=-1

最终效果：

用户消息：模型可以看到，但不要求模型模仿
助手回复：模型可以看到，并要求模型学会生成

8.3 SFT 数据混合

训练数据混合定义在 scripts/chat_sft.py。

TaskMixture 会将数据混合并使用固定随机种子打乱。如果希望提高某个任务的占比，可以将同一个 Task 多次加入列表。

8.4 身份数据

竞速脚本会下载：

identity_conversations.jsonl

其格式类似：

[
  {"role": "user", "content": "Who created you?"},
  {"role": "assistant", "content": "I am nanochat..."}
]

模型身份不是硬编码在程序中，而是通过训练数据注入。仓库提供了参考生成脚本 dev/gen_synthetic_data.py。

8.5 数学工具调用

GSM8K 数据包含：

<<12/60=0.2>>

tasks/gsm8k.py 将其转换为：

[
    {"type": "python", "text": "12/60"},
    {"type": "python_output", "text": "0.2"},
]

再渲染为：

<|python_start|>12/60<|python_end|>
<|output_start|>0.2<|output_end|>

mask 规则：

工具输出由程序提供，不应该要求模型自己预测。

8.6 SFT DataLoader 与预训练的差异

SFT 同样使用 best-fit 装箱，但剩余空间不足时使用 padding，而不是裁剪：

预训练：裁剪文档，吞吐量优先
SFT：padding，保留完整对话优先

padding 位置也会设置为 target=-1，不参与 loss。

8.7 SFT 停止条件与 checkpoint

SFT 默认完整遍历一次混合对话数据集。最终模型保存到：

~/.cache/nanochat/chatsft_checkpoints/d24/

此时模型已经可以作为聊天助手运行。

9. RL 如何强化数学能力

RL 是可选阶段，入口位于 scripts/chat_rl.py。

当前实现只针对 GSM8K 数学题。

9.1 SFT 与 RL 的差异

SFT：

问题 → 标准解答 → 模仿标准解答

RL：

问题
  → 模型自己采样多个回答
  → 检查最终答案是否正确
  → 正确回答获得奖励
  → 提高优秀轨迹的生成概率

9.2 从 SFT checkpoint 开始

RL 加载 SFT 模型：

model, tokenizer, meta = load_model("sft", device, phase="eval")

模型必须先通过 SFT 掌握对话格式、回答结束 token、数学解题格式和工具调用格式。

9.3 每道题生成多个 rollout

默认配置：

--num-samples=16
--device-batch-size=8
--temperature=1.0
--top-k=50

模型对同一道题生成 16 个回答。每个完整生成序列称为 rollout。

评分器只提取最终数字：

正确 → reward = 1
错误 → reward = 0

9.4 Advantage：相对奖励

项目不会直接使用 reward，而是计算：

mu = rewards.mean()
advantages = rewards - mu

假设 16 个回答中有 4 个正确：

平均奖励 mu = 4 / 16 = 0.25

正确回答 advantage = 1 - 0.25 =  0.75
错误回答 advantage = 0 - 0.25 = -0.25

含义：

• 正 advantage：提高该回答轨迹的概率。
• 负 advantage：降低该回答轨迹的概率。
• 如果所有回答都同样正确或同样错误，则不会产生更新。

9.5 Policy Gradient

训练目标：

logp = -model(inputs, targets, loss_reduction="none")
pg_obj = (logp * advantages.unsqueeze(-1)).sum()
loss = -pg_obj
loss.backward()

直观上：

正确轨迹 → 提高其中 token 的概率
错误轨迹 → 降低其中 token 的概率

9.6 这是简化版 GRPO

文件注释将其称为带引号的 “GRPO”。它实际更接近简化版 on-policy REINFORCE。

它移除了：

• 参考模型
• KL 正则
• trust region
• PPO ratio
• PPO clip
• z-score 标准化中的除以标准差

保留的核心是：

同一道题采样多个回答
  → 计算相对奖励
  → 使用 policy gradient 更新模型

9.7 RL checkpoint

RL 模型保存到：

~/.cache/nanochat/chatrl_checkpoints/d24/

可以通过以下命令使用：

python -m scripts.chat_eval -i rl
python -m scripts.chat_cli -i rl
python -m scripts.chat_web -i rl

10. 推理、KV Cache 与聊天 UI

训练时，模型一次处理完整序列：

(B, T) → GPT → (B, T, vocab_size)

推理时，模型逐 token 生成：

提示词 → token 1 → token 2 → token 3 → ...

核心实现位于 nanochat/engine.py。

10.1 没有 KV Cache 的问题

假设提示词有 100 个 token，需要生成 3 个 token。

朴素做法：

第 1 次：计算 100 个 token
第 2 次：重新计算 101 个 token
第 3 次：重新计算 102 个 token

历史 token 被反复计算，浪费大量算力。

10.2 哪些内容可以缓存

Attention 中：

Q = 当前输入的查询
K = 可供查询的索引
V = 可供读取的内容

历史 token 的 K 和 V 在生成过程中不会变化，因此可以保存：

计算新 token 的 Q、K、V
  → 将新 K、V 追加到缓存
  → 使用新 Q 查询全部历史 K、V

这就是 KV Cache。

10.3 KV Cache 的形状

缓存形状：

(
    num_layers,
    batch_size,
    seq_len,
    num_heads,
    head_dim,
)

每个 Transformer 层都有独立缓存。

10.4 Prefill 与 Decode

Engine 将推理分为两个阶段。

Prefill

对完整提示词执行一次前向传播：

<|bos|>
<|user_start|>Why is the sky blue?<|user_end|>
<|assistant_start|>

这一步将所有提示词 token 的 K/V 写入缓存。

Decode

之后每次只传入一个新 token：

ids.shape = (B, 1)

模型复用缓存中的历史 K/V，减少重复计算。

10.5 采样参数

Engine 支持：

temperature=0  → 每次选择概率最高的 token
temperature>0  → 按概率随机采样
top_k=50       → 只从概率最高的 50 个 token 中采样

温度越高，输出通常越随机。

10.6 多样本生成

RL 阶段需要对同一道题生成多个回答。Engine 只对提示词执行一次 prefill，然后复制 KV Cache：

一个提示词
  → 只计算一次提示词 KV
  → 复制缓存
  → 并行生成多个不同回答

10.7 工具调用状态机

模型生成：

<|python_start|>12/60<|python_end|>

Engine 检测到 <|python_end|> 后：

1. 解码表达式。
2. 使用受限计算器执行。
3. 将结果编码为 token。
4. 强制注入工具输出。

最终序列：

<|python_start|>12/60<|python_end|>
<|output_start|>0.2<|output_end|>

模型随后继续生成解释文本。

10.8 CLI 与 Web UI

命令行聊天：

python -m scripts.chat_cli
python -m scripts.chat_cli -p "Why is the sky blue?"

Web UI：

python -m scripts.chat_web

默认地址：

http://localhost:8000

Web 服务基于 FastAPI，提供：

10.9 训练与推理对比

11. 推荐的源码阅读顺序

第一次阅读时，建议沿着执行顺序逐步下钻：

1. runs/speedrun.sh：完整流程编排。
2. nanochat/dataset.py：数据分片下载与划分。
3. scripts/tok_train.py：分词器训练入口。
4. nanochat/tokenizer.py：BPE、特殊 token、对话渲染。
5. nanochat/dataloader.py：预训练 batch 构造。
6. nanochat/gpt.py：模型结构与前向传播。
7. scripts/base_train.py：预训练循环。
8. nanochat/optim.py：AdamW、Muon 和分布式优化器。
9. scripts/base_eval.py：BPB、CORE 和样本生成。
10. scripts/chat_sft.py：SFT 数据混合和 loss mask。
11. scripts/chat_rl.py：可选 RL 阶段。
12. nanochat/engine.py：KV Cache 与工具调用。
13. scripts/chat_cli.py：终端聊天。
14. scripts/chat_web.py：Web 服务与流式输出。

如果重点是理解 GPT 原理，可以先暂时跳过：

FP8
Flash Attention 3 内部细节
Muon 正交化数学推导
分布式通信优化
Web UI

先掌握：

文本
  → token
  → 固定长度 batch
  → Embedding
  → Transformer
  → logits
  → cross entropy
  → backward
  → optimizer.step()

12. 常用命令与产物目录

12.1 CPU / Apple Silicon 教学运行

runs/runcpu.sh 提供了适合本地体验的缩小版流程：

bash runs/runcpu.sh

它会训练较小模型，适合理解代码路径，但不会得到强模型。

12.2 8×H100 完整竞速流程

bash runs/speedrun.sh

12.3 单独执行各阶段

# 下载分片
python -m nanochat.dataset -n 8

# 训练与评估分词器
python -m scripts.tok_train
python -m scripts.tok_eval

# 预训练与基础模型评估
python -m scripts.base_train
python -m scripts.base_eval

# SFT 与聊天模型评估
python -m scripts.chat_sft
python -m scripts.chat_eval -i sft

# 可选 RL 与 RL 模型评估
python -m scripts.chat_rl
python -m scripts.chat_eval -i rl

# 聊天
python -m scripts.chat_cli
python -m scripts.chat_web

12.4 主要产物目录

默认根目录：

~/.cache/nanochat/

主要内容：

~/.cache/nanochat/
├── base_data_climbmix/    # 预训练文本分片
├── tokenizer/             # tokenizer.pkl、token_bytes.pt
├── base_checkpoints/      # 基础模型 checkpoint
├── chatsft_checkpoints/   # SFT 模型 checkpoint
├── chatrl_checkpoints/    # RL 模型 checkpoint
├── eval_bundle/           # CORE 评估数据
└── report/                # 训练报告片段

总结

nanochat 将一个小型聊天模型的生命周期完整地放在同一个仓库中：

文本数据
  → BPE 分词
  → GPT 预训练
  → 基础能力评估
  → 对话监督微调
  → 可选强化学习
  → KV Cache 推理
  → CLI / Web UI

理解这条主线后，再深入 FP8、Flash Attention、Muon、滑动窗口和分布式通信等优化，会更加清晰。它们改变的是训练效率和模型效果，而不是 GPT 的基本学习闭环。