数据截至 (上游 commit f1e2ace65149)
04 · 多 GPU 与混合精度
这一章讲什么： 把单卡 bf16 训练做「对」、把多卡训练做「快」的全部工程细节。两条线： 数值线（bf16 + fp32 master + 随机舍入）和通信线（NCCL、DDP、ZeRO-1），最后落在训练一步的完 整编排与断点续训上。
1. 它要解决的小问题
- bf16 只有 8 位尾数，直接拿它做
w -= lr * update：当 update 远小于 w 时，更新量会被舍入 吞掉，训练悄悄停滞。怎么既享受 bf16 的算力又不丢精度？ - 多卡各算各的梯度，怎么合？通信怎么不堵在反向传播后面？
2. 数值线：bf16 + fp32 master + 随机舍入
2.1 三种精度的分工
| 数据 | 精度 | 在哪 |
|---|---|---|
| 权重、激活、梯度 | bf16（floatX，llmc/cuda_common.h:90-91） | 前向/反向的全部 kernel |
| master 权重 | fp32（model->master_weights，train_gpt2.cu:300） | 只被优化器读写 |
| 优化器状态 m/v | fp32 | gpt2_allocate_state（train_gpt2.cu:404-409） |
| LayerNorm 的 mean/rstd、losses | fp32 | 激活 struct 里这几个字段是 float*（train_gpt2.cu:185-191） |
| matmul 累加 | fp32（cuBLAS compute type） | llmc/matmul.cuh:126 |
即「存储和传输 bf16，累加和状态 fp32」——bf16 的动态范围够大（指数位与 fp32 相同），但尾 数短，所以凡是「很多小数累加成大数」的地方都必须回到 fp32。
2.2 AdamW kernel：在 fp32 上更新，随机舍入写回
adamw_update（llmc/adamw.cuh:19-47）的单参数逻辑：
// 示意，非源码（真实代码见 llmc/adamw.cuh:26-46）
float grad = grad_scale * (float)grads[i]; // bf16 梯度转 fp32
m = lerp(grad, m, beta1); // fp32 动量
v = lerp(grad*grad, v, beta2);
float old = master ? master[i] : (float)params[i]; // 从 fp32 master 取旧值
float p = old - lr * (m_hat / (sqrt(v_hat) + eps) + wd * old);
stochastic_rounding(p, ¶ms[i], seed); // 写回 bf16：随机舍入
if (master) master[i] = p; // fp32 master 存全精度
两个细节：
- lerp 用两次 FMA 而不是三次运算（
llmc/adamw.cuh:12-15），数值更稳。 - 权重衰减只施加于 2D 张量（嵌入和各矩阵），LayerNorm/bias 不衰减；选择写在一个三元表达
式里（
train_gpt2.cu:1076-1081）。
2.3 随机舍入：为什么不是 round-to-nearest
普通舍入是「四舍五入到最近的 bf16」——当真实更新量恒小于半个 bf16 步长时，每次都被舍掉， 小更新系统性消失。随机舍入改为按距离比例掷签：
# 示意，非源码
def stochastic_round(x, rand_u16): # x 是 fp32
lo16 = f32_bits(x) & 0xFFFF # 被丢掉的小数部分（占步长的比例）
if lo16 > rand_u16: # 以 lo16/65536 的概率向上
return bf16_ceil(x)
else:
return bf16_floor(x)
期望上等于真值：小更新以正比于自身大小的概率被保留，长期不再系统性丢失。真实实现
（llmc/cuda_utils.cuh:269-278）用 SquirrelNoise5 从 (线程、block、种子) 生成随机数，保证每
个参数每个 step 的舍入确定可复现。这套机制同时用在三处：AdamW 写回、encoder 反向写
dwte/dwpe、从 master 重建权重（llmc/adamw.cuh:43、llmc/encoder.cuh:103-112）。
2.4 TF32：fp32 模式的隐藏加速
即使编译成 FP32，只要 GPU 是 Ampere+，matmul 也默认开 TF32 tensor core
（common_start，train_gpt2.cu:1191-1192）——等价于 PyTorch 的
torch.set_float32_matmul_precision('high')（同处注释）。-f 0 可关。
3. 通信线：NCCL、DDP、ZeRO-1
3.1 进程与通信初始化
multi_gpu_config_init（llmc/zero.cuh:411-463）按 -pm 三选一交换 NCCL unique ID：
| 方法 | 机制 | 适用 |
|---|---|---|
mpi | MPI_Bcast 广播；本地 GPU 选择用 hostname 哈希数同机进程（llmc/zero.cuh:372-408） | 有 MPI 的环境 |
tcp | rank 0 起 TCP server 分发 ID（llmc/zero.cuh:229-333） | 无共享 FS、无 MPI |
fs | rank 0 把 ID 写进共享文件，其余 sleep 轮询读（llmc/zero.cuh:336-369） | 有共享文件系统 |
README 解释了为什么要三种：新版 slurm 常不带 PMIx，MPI 未必可用（README.md 的 multi-node
一节）。
3.2 梯度归约：DDP vs ZeRO-1
multi_gpu_async_reduce_gradient（llmc/zero.cuh:514-556）是唯一的梯度通信入口，按
zero_stage 分两路：
zero_stage = 0 (DDP):
每 rank 持有全量梯度 ─ ─ ncclAllReduce(avg) ──► 每 rank 得到全量平均梯度
每 rank 重复更新全部参数（通信量 = 全参数量）
zero_stage = 1 (ZeRO-1):
每 rank 持有全量梯度 ── ncclReduceScatter(avg) ──► 每 rank 只得到自己那 1/N 分片的平均梯度
每 rank 只更新自己分片的 m/v/master ── ncclAllGather ──► 拼回全量新权重
（优化器状态内存 ÷ N，通信量不变）
分片规则就是「参数量 ÷ 进程数」，必须整除，否则 multi_gpu_get_shard_offset 直接报错退出
（llmc/zero.cuh:496-511）；set_zero_configs（llmc/zero.cuh:558-580）在不能整除时自动退
回 stage 0。stage 2/3 尚未实现，选了会打印提示并退回。
3.3 异步：通信藏在反向后面
梯度归约不是反向全部结束才发，而是逐层、且只在最后一个 micro-step 发
（train_gpt2.cu:923-948）：反向算到第 l 层结束，立刻把这一层 12 个张量的梯度指针打包成一组
NCCL 调用。两条 stream 的衔接靠 event（llmc/zero.cuh:526-529）：
cudaEventRecord(config->compute_nccl_sync, compute_stream); // 在计算流上打个点
cudaStreamWaitEvent(config->nccl_stream, ...); // 通信流等到这个点
ncclGroupStart(); ... ncclGroupEnd(); // 一组归约排进通信流
Host 不用同步，排完队继续给 GPU 派活——反向还在往浅层算，深层的梯度已经在网卡上跑了。非块
内参数（wte/wpe/lnf）在循环结束后单独再归约一组（train_gpt2.cu:957-961）。
3.4 loss 也要跨卡平均
每个 micro-step 的 loss 累加在 acts.losses 里，最后一步先卡内确定性求和
（global_sum_deterministic，单 block 归约保确定性，llmc/cuda_utils.cuh:190-206），再
ncclAllReduce(ncclAvg) 跨卡平均（train_gpt2.cu:952-956）。
4. 训练一步 的完整编排
main 循环的训练段（train_gpt2.cu:1826-1864）把前面所有零件串起来：
for micro_step in 0..grad_accum_steps: # 梯度累积内循环
dataloader_next_batch
gpt2_forward # dloss 已按 1/(B*T*grad_accum) 缩放（train_gpt2.cu:819）
gpt2_backward_and_reduce # 最后一个 micro-step 触发 NCCL
z-loss / z-grad = OutlierDetector 更新 # 窗口 128 的 z-score（llmc/outlier_detector.h:36）
grad_norm = gpt2_calculate_grad_norm # 全局范数
若 z 超阈 → 打印并跳过本步更新（train_gpt2.cu:1844-1848）
否则 grad_scale = min(1, 1.0/grad_norm) # 裁剪到范数 1.0（train_gpt2.cu:1850-1851）
gpt2_update(lr, β1=0.9, β2=0.95, eps, wd, grad_scale, step+1)
两处值得展开：
- 梯度累积的缩放藏在 dloss 里。 每个 micro-step 的 loss 权重是
1/(B*T*grad_accum_steps)（train_gpt2.cu:819），所以 N 个 micro-step 的梯度累加结果天然 是「平均梯度」，主循环不需要任何除法。 - ZeRO-1 下范数只算本地分片。
gpt2_calculate_grad_norm（train_gpt2.cu:992-1033）按 分片偏移算局部平方和，再ncclAllReduce合成全局（train_gpt2.cu:1000-1024）；范数计算 本身用两级归约 kernel（llmc/global_norm.cuh:27-49）。
学习率调度（cosine/linear/constant/wsd 四种）在 llmc/schedulers.h:83-98；每步的吞吐按
EMA 平滑并做偏差修正（train_gpt2.cu:1859-1869）。
4.1 MFU：训练快不快，看这一行
gpt2_estimate_mfu（train_gpt2.cu:1126-1150）用 PaLM 论文的近似公式
flops_per_token = 6N + 6·L·C·T 算出本步 FLOPs，除以耗时得实际算力，再除以按 GPU 型号查表
得到的峰值（llmc/mfu.h:56 的 gpu_db，收 Volta/Ampere/Ada/Hopper 各型号的 tensor core 峰值
与频率）——这就是日志里的「%.1f%% bf16 MFU」。llm.c 复现 GPT-2 时报告 ~50%+ 的 MFU，是「手
写栈也能打」的关键证据。
5. 断点续 训：连舍入随机性都存下来
5.1 checkpoint 里有什么
write_checkpoint（train_gpt2.cu:1316-1334）写三类文件：
| 文件 | 谁写 | 内容 |
|---|---|---|
model_XXXXXXXX.bin | 仅 rank 0 | 模型权重（bf16/fp32，各 rank 相同） |
state_XXXXXXXX_RRRRR.bin | 每个 rank | 本 rank 的 m/v、master weights、RNG 状态、dataloader 位置与 shuffle 状态（save_state，train_gpt2.cu:1209-1248） |
DONE_XXXXXXXX | rank 0 | 空文件，标记该 checkpoint 完整（resume 只认有 DONE 的） |
5.2 bit-perfect 的关键一招
恢复时（load_state，train_gpt2.cu:1250-1314）先读回 fp32 master weights，然后做一件很漂
亮的事（train_gpt2.cu:1283-1287）：
model->rng_state = model->rng_state_last_update; // 回到上次更新前的 RNG 状态
gpt2_update(..., /* init_from_master_only */ true); // 用同一串随机数重新舍入
bf16 权重不直接存（它由 master + 舍入随机数唯一决定），而是用保存的 RNG 状态把舍入过程重放
一遍——恢复后的 bf16 权重与崩溃前逐位相同。init_from_master_kernel
（llmc/adamw.cuh:60-71）的 block size 都和 AdamW 保持一致，只为让随机数序列对齐
（llmc/adamw.cuh:93 注释）。恢复后还能换精度：只存了 fp32 master，所以同一份 state 可以
按新精度重建权重。
续训入口是 -y 1 + find_max_step 找最新 DONE（llmc/utils.h:192），scripts/ 下的复现脚
本把启动命令包在 while true 里，崩了自动接力（scripts/run_gpt2_124M.sh:12-30）。
6. 关键细节与坑
- 整除是硬约束。 参数量、各张量大小都要能被进程数整除，否则报错或自动退回 DDP
（
llmc/zero.cuh:500-511、:565-575）。 - 显存不够时静默降级。
cudaMallocConditionallyManaged（llmc/cuda_utils.cuh:213-232）在 OOM 时改用 managed memory（host 兜底），并打印「可能因 host↔device 换页而明显变慢」的警告 （train_gpt2.cu:415-419）。 - fp16 路径是半成品。 编译可选，但没有梯度缩放（
llmc/cuda_common.h:86-88注释）， checkpoint 加载直接拒绝 fp16（train_gpt2.cu:459-463）；fp16 的 stochastic_rounding 还是 todo（llmc/cuda_utils.cuh:279-281）。 - 推理与训练的精度坑。 采样时用较小的 T 前向，非 cuDNN 路径要把 att 缓冲区未用区域清零
（
train_gpt2.cu:722-725），cuDNN 路径则注释坦承「T≥256 时逐位一致，但 cuDNN 可能换算法」 （train_gpt2.cu:1765-1772）。 - MFU 表是手维护的。 不在
gpu_db里的新卡返回 -1（train_gpt2.cu:1143-1146），日志里就 没有 MFU 一栏。
7. 代码地图
| 主题 | 文件路径 | 符号名 |
|---|---|---|
| 精度类型 | llmc/cuda_common.h、llmc/cublas_common.h | floatX、PrecisionMode、cublas_compute |
| 随机舍入 | llmc/cuda_utils.cuh | stochastic_rounding、SquirrelNoise5、Get2dNoiseUint |
| AdamW | llmc/adamw.cuh | adamw_update、adamw_kernel3、init_from_master |
| 多卡配置与初始化 | llmc/zero.cuh | MultiGpuConfig、multi_gpu_config_init、multi_gpu_get_local_device_idx |
| 梯度归约 | llmc/zero.cuh | multi_gpu_async_reduce_gradient、multi_gpu_get_shard_offset、set_zero_configs |
| 范数与裁剪 | llmc/global_norm.cuh、train_gpt2.cu | global_norm_squared、gpt2_calculate_grad_norm |
| 异常检测 | llmc/outlier_detector.h | OutlierDetector、update_detector |
| 学习率 | llmc/schedulers.h | get_learning_rate、lr_scheduler_init |
| MFU | train_gpt2.cu、llmc/mfu.h | gpt2_estimate_mfu、get_flops_promised、gpu_db |
| 续训 | train_gpt2.cu、llmc/utils.h | save_state、load_state、write_checkpoint、find_max_step |
| 主循环编排 | train_gpt2.cu | main、gpt2_backward_and_reduce、gpt2_update |