跳到主要内容

数据截至 (上游 commit f1e2ace65149)

01 · train_gpt2.c：CPU 参考实现

这一章讲什么： 全仓库最该先读的文件。train_gpt2.c（1182 行）用纯 C + 少量 OpenMP 写完 GPT-2 的训练：模型怎么住进内存、前向逐层算什么、反向怎么手工推、AdamW 怎么更新、main 循环长 什么样。读完它，后面 CUDA 版的一切优化你都知道在优化什么。

1. 它要解决的小问题​

不用任何框架，「训练 GPT-2 一步」到底是哪几步数学？

PyTorch 里这是 loss = model(x, y); loss.backward(); opt.step() 三行。llm.c 把这三行展开成 1182 行——但没有一行是浪费的：每个层的前向和反向各一个函数，加模型定义、checkpoint 读取、 主循环，仅此而已。

2. 模型在 C 里怎么住：一块大数组 + 一把指针​

2.1 16 个参数张量​

GPT-2 的全部参数被列在一个 struct 里（ParameterTensors，train_gpt2.c:536-553）：

张量形状含义
wte(V, C)token 嵌入（同时兼任最后的分类器，见 §3.3）
wpe(maxT, C)位置嵌入
ln1w/b、ln2w/b、lnfw/b(L, C) / (C)各 LayerNorm 的缩放与平移
qkvw/b(L, 3C, C)QKV 投影（Q、K、V 合成一个矩阵）
attprojw/b(L, C, C)注意力输出投影
fcw/b(L, 4C, C)MLP 升维
fcprojw/b(L, C, 4C)MLP 降维

每个张量的元素数由 fill_in_parameter_sizes（train_gpt2.c:556-577）算出。注意两个细节：

  • 填充词表：logits/嵌入按 padded_vocab_size（50304 = 50257 向上取到 128 的倍数）分配，为了 让 CUDA kernel 对齐（train_gpt2.c:529 的注释；CPU 版沿用了同一约定）。
  • NUM_PARAMETER_TENSORS = 16：这个数在后面测试、checkpoint、ZeRO 分片里反复出现。

2.2 一次 malloc，指针依次排进去​

分配方式是全仓库最有「C 味」的设计之一（malloc_and_point_parameters，train_gpt2.c:580-598）：

// 示意，非源码（真实代码见 train_gpt2.c:580-598）
float* params_memory = mallocCheck(num_parameters * sizeof(float));  // 一整块
float* it = params_memory;
for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) {
    *(ptrs[i]) = it;              // 把 struct 里第 i 个指针指到当前位置
    it += param_sizes[i];         // 游标前进一个张量
}

要点：没有逐张量 malloc，没有智能指针。grads（梯度）用同一个函数再切一块同样布局的内存； gpt2_free（train_gpt2.c:1035-1044）只需 free 掉几个大指针。23 个激活张量 （ActivationTensors，train_gpt2.c:601-625）照此办理，大小由 fill_in_activation_sizes（train_gpt2.c:628-655）按 B、T 算出。

2.3 激活是「为反向缓存的中间量」​

23 个激活张量值得多看一眼，因为它精确回答了「反向传播需要记住什么」：

  • 每个 LayerNorm 存 mean 和 rstd（反向要用，见 §4.3）。
  • 注意力存 preatt（softmax 前的分数）和 att（softmax 后的概率），形状 (L, B, NH, T, T)。
  • fch（GELU 前）和 fch_gelu（GELU 后）都存——GELU 反向需要输入值。
  • 残差流的每一站（encoded、residual2、residual3）都在，因为反向要按原路把梯度加回去。

3. 前向：gpt2_forward 一次走完​

3.1 层内九连​

gpt2_forward（train_gpt2.c:765-890）主体是一个 for 循环，每层依次调 9 个函数 （train_gpt2.c:863-872）：

residual ─► layernorm ─► matmul(QKV) ─► attention ─► matmul(proj) ─► (+residual)
        ─► layernorm ─► matmul(FC 4C) ─► gelu ─► matmul(proj C) ─► (+residual)

对应代码就是把这串调用原样写出，例如前三行：

layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C);
matmul_forward(l_qkv, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C);
attention_forward(l_atty, l_preatt, l_att, l_qkv, B, T, C, NH);

（train_gpt2.c:863-865）每层开始前先算一串 l_ 前缀的层内指针（params.qkvw + l * 3*C * C 这 种），这正是 §2 内存布局的用法：层 l 的张量 = 大数组 + l × 单层大小。

3.2 每个层的数学（各 20~70 行）​

层函数（行号）一句话
嵌入encoder_forward（train_gpt2.c:35）out[b,t] = wte[token] + wpe[t]，两次查表相加
LayerNormlayernorm_forward（train_gpt2.c:78）每个 (b,t) 对 C 维求均值/方差，归一化后缩放平移；缓存 mean/rstd 给反向（train_gpt2.c:114-115）
矩阵乘matmul_forward（train_gpt2.c:184）朴素三重循环的 tiling 版：8 个 (b,t) 一组，权重行复用 8 次（train_gpt2.c:196-225）
注意力attention_forward（train_gpt2.c:271）4 趟循环：QK 点积 → 减 max 求 exp → 归一化（含因果掩码）→ 加权求和 V
GELUgelu_forward（train_gpt2.c:408）tanh 近似的逐元素非线性
残差residual_forward（train_gpt2.c:436）逐元素相加
softmaxsoftmax_forward（train_gpt2.c:449）减 max 的数值稳定版；只遍历到真实 V，padding 区写 0
交叉熵crossentropy_forward（train_gpt2.c:486）loss = -log(probs[target])

注意力是唯一跨时间步混合信息的层（train_gpt2.c:279-281 的注释专门点出这一点）；其余所有运算 都在每个 (b,t) 位置独立进行。

3.3 权重共享：分类器就是 wte​

循环结束后还有三步：最终 LayerNorm、logits、softmax。其中 logits 的矩阵乘用的权重是 params.wte（train_gpt2.c:869）：

matmul_forward(acts.logits, acts.lnf, params.wte, NULL, B, T, C, Vp);

这就是 GPT-2 的 weight tying——嵌入矩阵转置后兼任输出投影。反向时两份梯度都累加进同一个 grads.wte：一份来自 logits 的 matmul_backward（train_gpt2.c:935），一份来自 encoder_backward（train_gpt2.c:1004）。

4. 反向：没有 autograd 的手工链式法则​

4.1 总体结构​

gpt2_backward（train_gpt2.c:898-1005）是前向的镜像：先把 dlosses 全部填成 1/(B*T)（均值 loss 对各位置 loss 的导数，train_gpt2.c:928-930），然后按前向的逆序逐层调对应的 _backward。

一个贯穿始终的约定（CUDA 版的头注释把它写成了明文，llmc/layernorm.cuh:1-10）：

  • 参数梯度用 +=——这样梯度累积（多个 micro-batch）天然成立；
  • 激活梯度用 =——每个激活的梯度只有一个来源，直接写更快；
  • 残差流例外——它是两路梯度的汇合点，所以 residual_backward 和 layernorm_backward 对它 用 +=。

4.2 两个好懂的例子​

softmax + 交叉熵的融合反向（crossentropy_softmax_backward，train_gpt2.c:502-522）。前向 是两步（softmax → -log），反向推完链式法则后收敛成教科书里最美的一行：

float indicator = i == ix ? 1.0f : 0.0f;
dlogits_bt[i] += (p - indicator) * dloss;

（train_gpt2.c:516-517）「概率减 one-hot」——softmax+CE 的导数不需要任何中间量，这也是 CUDA 版 fused classifier 能把整个反向第一步塞进一个 kernel 的数学基础（第 2 章 §6）。

matmul 反向 = 两个矩阵乘（matmul_backward，train_gpt2.c:231-268）。前向 out = inp @ w^T ，对输入的梯度是 dinp = dout @ w，对权重的梯度是 dw = dout^T @ inp。代码把两个方向拆成两 个独立循环：dinp 按 (b,t) 并行、dw 按输出通道并行——注释解释了为什么不合成一趟：「那样没有好 的并行化策略」（train_gpt2.c:233-236）。

4.3 LayerNorm 反向的三项公式​

layernorm_backward（train_gpt2.c:120-160）是手工求导的样板。前向是 out = w * (x - mean) * rstd + b，反向对 x 的梯度化简成三项：

dval += dnorm_i;                    // term 1: 直接通路
dval -= dnorm_mean;                 // term 2: 均值通路
dval -= norm_bti * dnorm_norm_mean; // term 3: 方差通路
dval *= rstd_bt;                    // 最后乘回 rstd

（train_gpt2.c:149-154）注意它只需要前向缓存的 mean、rstd 和输入 x——这正是 §2.3 里那 些 (L, B, T) 小缓冲存在的原因。仓库还配了一篇从零推这个公式的教程 （doc/layernorm/layernorm.md），先 PyTorch 再 C，是官方指定的入门读物。

4.4 注意力反向​

attention_backward（train_gpt2.c:347-405）把前向的 4 趟倒着重放：先过 V 的加权求和（得到 datt、dvalue），再过 softmax（用 att[t2] * (I - att[t3]) 的雅可比，train_gpt2.c:378-385 ），最后过 QK 点积（dquery、dkey）。每行旁边都写着「前向这里是……所以现在……」的对照注释， 是学手工反向最好的 60 行。

5. 优化器：AdamW 三十行​

gpt2_update（train_gpt2.c:1007-1033）就是 PyTorch 文档里 AdamW 的逐行翻译：

float m = beta1 * model->m_memory[i] + (1.0f - beta1) * grad;   // 一阶动量
float v = beta2 * model->v_memory[i] + (1.0f - beta2) * grad * grad; // 二阶动量
float m_hat = m / (1.0f - powf(beta1, t));                      // 偏差修正
float v_hat = v / (1.0f - powf(beta2, t));
model->params_memory[i] -= learning_rate * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * param);

（train_gpt2.c:1020-1031）m/v 两个缓冲在首次调用时 calloc（train_gpt2.c:1011-1014）——优 化器状态是参数量的两倍，这在 CPU 版是 fp32，GPU 版同样如此（见第 4 章）。

6. 主循环：main 的 100 行​

main（train_gpt2.c:1077-1180）把一个最小训练流程串了起来：

  1. gpt2_build_from_checkpoint 读 gpt2_124M.bin——256 个 int 的文件头（magic 20240326 + 版 本 + 超参），后面紧跟全部 fp32 参数（train_gpt2.c:707-763）。
  2. 两个 DataLoader（train/val，B=4、T=64，train_gpt2.c:1085-1097），一个 Tokenizer。
  3. 40 步循环：每 10 步估一次 val loss；每 20 步采样 64 个 token 打印；每步 dataloader_next_batch → gpt2_forward → gpt2_zero_grad → gpt2_backward → gpt2_update （train_gpt2.c:1163-1168）。
  4. 采样用 sample_mult（train_gpt2.c:1062-1073）：掷一枚 [0,1) 的硬币，在累积分布上找落点。

7. 关键细节与坑​

  • GELU 和 -Ofast 打架。 仓库想全局开 -Ofast，但 GELU 在 fast-math 下会算错，于是单独给 gelu_backward 关了 finite-math 优化（train_gpt2.c:417-421，注释里挂着 issue #168）。这是 「编译器优化改变浮点语义」的实物教材。
  • maxval 初始值是 -10000。 注意力和 softmax 的减 max 都用这个魔数起步，旁边留着 // TODO something better（train_gpt2.c:294、train_gpt2.c:462）。参考实现的可读性优先于 边角严谨。
  • B、T 首次前向后锁死。 激活内存在第一次 forward 时惰性分配，之后 B/T 不许变，否则直接退出 （train_gpt2.c:806-813）。简单但够安全。
  • 推理是暴力重算。 采样循环每生成一个 token 就对全部 (B,T) 重新前向，注释自曝 「very wasteful」（train_gpt2.c:1134-1135）；而且 B 路并行采样只用了第 0 路。
  • padding 词表只分配不训练。 softmax/交叉熵都只遍历到真实 V（50257），padding 区的概率恒为 0（train_gpt2.c:466-476）；wte 的 padding 行在 Python 侧初始化时就是 0。

8. 代码地图​

主题文件路径符号名
层前向train_gpt2.cencoder_forward、layernorm_forward、matmul_forward、attention_forward、gelu_forward、residual_forward、softmax_forward、crossentropy_forward
层反向train_gpt2.cencoder_backward、layernorm_backward、matmul_backward、attention_backward、gelu_backward、residual_backward、crossentropy_softmax_backward
参数/激活布局train_gpt2.cParameterTensors、ActivationTensors、fill_in_parameter_sizes、malloc_and_point_parameters、malloc_and_point_activations
模型生命周期train_gpt2.cgpt2_build_from_checkpoint、gpt2_forward、gpt2_zero_grad、gpt2_backward、gpt2_update、gpt2_free
主循环与采样train_gpt2.cmain、sample_mult、random_f32
对照教程doc/layernorm/layernorm.md（文档，从 PyTorch 推到 C）