第 2 周 · Day 4：预训练、显存账与分布式训练

时长 8 小时：§1 数据（1.5h）· §2 训练不稳定（1h）· §3 数值精度（1h）· §4 显存进阶（1h）· §5 分布式训练（2h）· §A–§C（1.5h）

对应正课 Lecture 4 的预训练与硬件优化部分。学习计划点名要会的：数据清洗（去重、质量过滤、去污染）、loss spike 的处理、BF16 为什么打败 FP16、7B 全量微调的显存账、ZeRO 三个阶段各切了什么、梯度检查点。

配套代码 w2day4_checks.py（纯 numpy）：MinHash 与 LSH、三种浮点格式与下溢 / 吞没、7B 训练的参数状态与激活、梯度累积的平均陷阱、ZeRO 各阶段显存、流水线气泡、数据并行与张量并行的通信量

§B 参考实现（code/ 目录）：w2d4_memory_dist.py（显存计算器；梯度检查点与梯度累积的正确性与显存对比；用 torchrun 在 CPU 上跑通 DDP 与 FSDP 的最小例子）。

今天的目标：能画出预训练数据管线并说出每一步防的是什么；能解释 BF16 为什么是训练的默认选择；能在白板上算出 7B 全量训练需要多少显存、ZeRO 每个阶段每卡省到多少；能说清数据、张量、流水线三种并行的通信模式和适用范围。

§0 今天的主线

Day 1 告诉我们要训多大、喂多少数据。今天回答怎么真的把它训出来：

数据从哪来、怎么洗                 §1   → 质量是新的缩放维度（Day 1 §4.3）
训练中途炸了怎么办                 §2   → loss spike
用什么数值精度                     §3   → BF16 + fp32 主权重
一张卡放不下怎么办                 §4–§5 → 显存账 → 梯度检查点 → ZeRO / 张量并行 / 流水线并行

今天要用到的旧结论：16 字节/参数的显存账、LoRA 省什么、梯度检查点的 $\sqrt L$（Week 0 Day 1 §9）；fp16 的动态范围与 loss scaling（Week 0 Day 1 §10.2）；激活显存的 Korthikanti 公式（Week 1 Day 3 §4.3）；FlashAttention 不存 $T\times T$ 矩阵（Week 1 Day 6 §3）；训练 $6ND$（Week 1 Day 2 §5.2）；mean 还是 sum（Week 0 Day 1 §6.4）。

§1 预训练数据

1.1 数据从哪来

主力是 Common Crawl：一个公开的网页爬取档案，每月新增数十亿网页。此外还有代码（GitHub）、书籍、学术论文、百科、多语言语料。Llama 3 用了约 15T token；Hugging Face 的 FineWeb 从 Common Crawl 清洗出了约 15T token 的英文网页数据。

原始网页的质量极差：导航栏、广告、模板文字、乱码、机器生成的垃圾、大量重复。一份典型的清洗管线会丢掉绝大部分原始数据。

1.2 清洗管线

网页 HTML
  → 正文抽取（去掉导航、页脚、广告）
  → 语言识别（fastText 语言分类器）
  → 启发式质量过滤（长度、符号比例、重复行比例、停用词比例……）
  → 基于模型的质量过滤（分类器打分）
  → 去重（精确 + 近重复）
  → 去除个人信息、有害内容
  → 去污染（删掉与评测集重叠的文本）
  → 按比例混合不同来源

每一步都在防一件具体的事：

• 启发式过滤：Gopher 论文的一组规则很有代表性，比如”平均词长不在 3–10 之间的丢掉""以省略号结尾的行占比过高的丢掉""不含常见停用词的丢掉”。它们便宜、可解释，能去掉大量明显的垃圾。
• 模型打分：用一个小分类器判断”这像不像高质量文本”。FineWeb-Edu 用 LLM 给一批网页的”教育价值”打分，再训一个小分类器去筛全量数据，报告在知识和推理类基准上有明显提升。这是”用模型筛数据”的典型做法，也是合成数据的前身。
• 去重（§1.3）：重复数据会被模型背下来，浪费算力，还会损害泛化。
• 去污染：如果评测题出现在训练数据里，分数就不再说明能力。GPT-3 用 13-gram 重叠来检测并删除。第 4 周 Lecture 8 讲评测时，这是第一个要问的问题。

1.3 近重复去重：MinHash + LSH

问题：几十亿网页里找”几乎一样”的页面（同一篇新闻被上百个站转载，只改了页眉页脚）。两两比较是 $O(n^2)$，不可能。

第一步：把文档变成集合。 取所有连续 5 个词（5-gram shingle），文档就是一个 shingle 集合。两篇文档的相似度用 Jaccard $J(A,B)=|A\cap B|/|A\cup B|$。

第二步：MinHash 把大集合压成短签名。 取一个随机哈希函数 $h$，记录集合里所有元素的最小哈希值 $\min_{x\in A}h(x)$。关键性质：

$P\Big(\min_{x\in A}h(x)=\min_{x\in B}h(x)\Big)=J(A,B)$

推导：把 $A\cup B$ 的所有元素按 $h$ 值排序，哈希值最小的那个元素等可能是其中任何一个。两个最小值相等，当且仅当这个元素同时在 $A$ 和 $B$ 里，概率是 $|A\cap B|/|A\cup B|$。$\blacksquare$

用 $k$ 个独立的哈希函数，得到 $k$ 维签名，签名中相同位置的比例就是 $J$ 的无偏估计，标准差 $\sqrt{J(1-J)/k}$。w2day4_checks.py 第 1 组（真实 $J=0.571$，每种 $k$ 用 20 组不同哈希函数）：

第三步：LSH 避免两两比较。 把签名切成 $b$ 段、每段 $r$ 个哈希，任一段完全相同就放进同一个桶，只比较同桶的文档。相似度为 $s$ 的两篇文档成为候选对的概率是 $1-(1-s^r)^b$。第 1 组（$b=20,r=5$）：相似度 0.3 时 0.047，0.5 时 0.470，0.7 时 0.975。在阈值附近是陡峭的 S 形：相似的几乎一定被抓到，不相似的几乎一定被放过。

1.4 数据混合与退火

不同来源按比例混合（网页、代码、数学、多语言……），比例本身是一个重要的超参，通常用小模型实验来调（又是 Day 1 的缩放律思路）。

退火：训练最后阶段，学习率衰减的同时，提高高质量数据的比例。Llama 3 报告这样做能在关键基准上获得明显提升。直觉：学习率小的时候模型”定型”，这时候喂什么，它最终就更像什么。

§2 训练不稳定：loss spike

2.1 现象

大规模预训练的 loss 曲线上经常出现突然的尖峰：loss 在几步内飙升，有时自己恢复，有时直接发散。OPT-175B 公开的训练日志记录了多次因发散而回退、调整后重启的过程。

2.2 常见原因与对策

工程纪律：把梯度范数画出来。它通常比 loss 更早预警（Week 0 Day 4 §A.6）。

§3 数值精度：为什么是 BF16

3.1 三种格式

w2day4_checks.py 第 2 组：

BF16 就是把 FP32 的尾数砍掉 16 位：指数位不变，所以范围和 FP32 一样；精度只有约 3 位十进制有效数字。FP16 反过来：精度比 BF16 高，范围却窄得多。

3.2 FP16 的问题：范围

梯度往往很小。第 2 组模拟 10 万个以 $10^{-6}$ 为中心的梯度：7.9% 在 FP16 里直接下溢为 0。这就是 Week 0 Day 1 §10.2 讲的 loss scaling：先把 loss 乘一个大数（比如 1024），梯度随之放大，下溢比例降到 0.0%，更新前再除回去。能用，但要动态调 scale，出现 inf 时要跳过这一步，工程上很烦。

3.3 BF16 的问题：精度，以及为什么需要 fp32 主权重

BF16 的范围够，不需要 loss scaling。但它的精度粗：1 附近相邻两个可表示数的间隔是 $2^{-7}=0.0078$。第 2 组：权重 1.0 加上 $10^{-3}$ 的更新，在 BF16 里舍入回 1.0，更新被吞掉了。

所以混合精度训练的标准做法（Week 0 Day 1 §9.1 那张 16 字节的表）：前向和反向用 BF16，优化器在 FP32 的主权重上累加更新，再把结果转成 BF16 给下一步用。

结论（学习计划的自测题）：BF16 打败 FP16，因为训练中范围比精度更重要：范围不够会让梯度下溢、激活溢出，需要 loss scaling 来补；精度不够可以靠 FP32 主权重补，而且神经网络对舍入噪声本身相当鲁棒。

更进一步：FP8（E4M3 最大约 448，E5M2 最大约 57344）已经用于大规模训练的矩阵乘（DeepSeek-V3 报告了 FP8 混合精度训练），需要细粒度的缩放因子。这是 Day 5 量化的”训练版”。

§4 显存账：进阶

4.1 7B 全量训练

Week 0 Day 1 §9.1 算过参数相关的部分：16 字节/参数。现在把激活也算上。w2day4_checks.py 第 3 组（7B，序列 4096，每卡 batch 1，BF16 混合精度，激活用 Korthikanti et al. 2022 的公式）：

光是参数状态就放不进一张 80GB 的卡。 这是 §5 一切切分方法的出发点。

4.2 梯度检查点

Week 0 Day 1 §9.4 给过原理：不存中间激活，反向时重算。实践中常见的是逐层检查点：只保存每层的输入（$2sbd$ 字节），反向到某层时重新做这一层的前向。代价是多一次前向，由训练 = 前向 + 2 倍前向（Week 1 Day 2 §5.2），总计算从 3 份变 4 份，约 +33%。换来的是第 3 组那一行：激活从 17 GiB 降到 1.5 GiB。

什么时候开：显存不够时的第一选择，通常比减小 batch 更划算（小 batch 降低 GPU 利用率）。

4.3 梯度累积，以及一个经典 bug

显存只够放小 batch，但想要大 batch 的训练效果：跑 $k$ 个微批次，梯度累加，再更新一次。等价于一个 $k$ 倍大的 batch，但有一个陷阱。

如果每个微批次的有效 token 数不同（变长数据、padding 不计入损失），用 reduction="mean" 算每个微批次的 loss，再除以 $k$，得到的是”平均的平均”，不等于全局按 token 平均。第 4 组（4 个微批次，有效 token 数 500 / 100 / 400 / 50）：

50 个 token 的微批次和 500 个的权重一样大，短序列被过度加权。梯度同理。这正是 Week 0 Day 1 §6.4”mean 还是 sum”的延伸，2024 年 HF transformers 里修过这个 bug。正确做法：用 reduction="sum"，最后除以整个累积步的总 token 数。

§5 分布式训练

5.1 数据并行（DP / DDP）

做法：每张卡放一份完整的模型，各自处理不同的数据，反向后用 all-reduce 把梯度求平均，每张卡做同样的更新。

通信：每步同步一次全部梯度。ring all-reduce 每卡收发约 $2\frac{N-1}N\times$ 梯度大小。第 7 组：7B 模型 8 卡，每步每卡约 22.8 GiB。可以和反向计算重叠（算完一层的梯度就开始同步这一层），所以能跨机。

局限：每张卡都要放下完整的模型和优化器状态（§4.1 的 100 GiB），大模型根本放不下。

5.2 ZeRO 与 FSDP：把冗余切掉

数据并行里，$N$ 张卡存着 $N$ 份一模一样的优化器状态、梯度和参数。ZeRO（Rajbhandari et al. 2019）的思路：每张卡只负责 $1/N$，需要的时候再通信取回。设参数量 $\Psi$，混合精度 AdamW（参数 2 字节、梯度 2 字节、优化器状态 $K=12$ 字节：FP32 主权重 + $m$ + $v$）：

第 5 组（按 $\Psi=7\times10^9$ 算，所以 DDP 一列是 104.3 GiB，而 §4.1 用的是 Llama-2-7B 的实际参数量 $6.74\times10^9$，得 100.4 GiB）：

读这张表：优化器状态占 16 字节中的 12 字节，所以 ZeRO-1 收益最大；只有 ZeRO-3 能让每卡显存随卡数线性下降。代价：ZeRO-3 在前向和反向时都要临时 all-gather 需要的参数，通信量约为 DDP 的 1.5 倍（ZeRO 论文的分析）。ZeRO-1、ZeRO-2 的通信量与 DDP 相同。

PyTorch 的 FSDP（Fully Sharded Data Parallel）就是 ZeRO-3 的实现，是今天 PyTorch 生态里最常用的方案。

5.3 张量并行（TP）

做法：把每一层的矩阵切开分到多张卡上。Megatron-LM 的经典切法：MLP 的第一个矩阵按列切（每张卡算一部分中间维度），第二个矩阵按行切（每张卡算一部分和），最后一次 all-reduce 求和。注意力按头切。每个 Transformer 块的前向只需要 2 次 all-reduce（注意力后、MLP 后）。

通信：每层都要同步，而且在关键路径上（下一层要等这一层的结果）。第 7 组：每次 all-reduce 32 MiB（$b=1,s=4096,d=4096$，bf16），32 层前向加反向约 4 GiB。在 NVLink（约 600 GB/s）上传 1 GiB 约 1.8 ms，在 InfiniBand（每卡约 25 GB/s）上约 43 ms。所以张量并行只放在机内。

5.4 流水线并行（PP）

做法：按层切，第 1–8 层在卡 1，第 9–16 层在卡 2……一个 batch 切成 $m$ 个微批次，像流水线一样依次流过各个阶段。

代价：气泡。 流水线填满和排空时，总有卡在空等。$p$ 个阶段、$m$ 个微批次，空闲比例是 $(p-1)/(m+p-1)$（GPipe）。第 6 组：

微批次越多气泡越小，但要同时存的激活也越多。1F1B 调度（前向一个、反向一个交替）不改变气泡大小，但把同时存的激活从 $m$ 份降到约 $p$ 份。

通信：只在相邻阶段之间传激活，量小，适合跨机。

5.5 三维并行

大规模训练把三种并行组合起来：

机内（NVLink，高带宽）：        张量并行
跨机（InfiniBand，低带宽）：     流水线并行
最外层：                        数据并行（通常是 ZeRO-1 / FSDP）

选择原则：通信越频繁、越在关键路径上的并行，放在带宽越高的地方。 超长序列还会加上序列 / 上下文并行（沿序列维度切激活），这是长上下文训练的必需品。

你需要知道到什么程度：能说清每种并行切的是什么、通信什么、代价是什么，能根据模型大小和硬件做粗略的选择。第 2 周的 LoRA 微调（Day 6）通常一张卡就够，用不到这些；但读论文和面试时这是必考内容。

§A 训练工程：API 用法与边界

A.1 梯度检查点

from torch.utils.checkpoint import checkpoint
y = checkpoint(block, x, use_reentrant=False)       # 推荐显式写 use_reentrant=False
model.gradient_checkpointing_enable()                # HF 模型一行开启
model.config.use_cache = False                       # 训练时关掉 KV Cache（Week 1 Day 4 §A.1）

• 新版本 PyTorch 要求显式传 use_reentrant，不传会警告。False 是推荐的新实现。
• 被检查点包住的模块里如果有 dropout，重算时要用同样的随机数，PyTorch 默认会保存并恢复 RNG 状态（preserve_rng_state=True），别关掉它。
• HF 模型开启梯度检查点又用 LoRA 时，如果输入嵌入被冻结，要调用 model.enable_input_require_grads()，否则梯度可能流不到检查点区域里（Day 6 §A）。

A.2 混合精度

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    loss = model(x, y)[1]
loss.backward()                                      # bf16 不需要 GradScaler

FP16 要配 torch.amp.GradScaler，顺序是 scaler.scale(loss).backward() → scaler.unscale_(opt) → 裁剪 → scaler.step(opt) → scaler.update()（Week 0 Day 2 §A.5）。

A.3 分布式的最小骨架

# 启动：torchrun --nproc_per_node=2 train.py
import torch.distributed as dist
dist.init_process_group("nccl")                      # GPU 用 nccl，CPU 调试用 gloo
rank, world = dist.get_rank(), dist.get_world_size()
model = torch.nn.parallel.DistributedDataParallel(model.to(rank), device_ids=[rank])
sampler = torch.utils.data.DistributedSampler(dataset)   # 每个进程拿不同的数据
for epoch in range(E):
    sampler.set_epoch(epoch)                         # ★ 不调用的话，每个 epoch 的打乱顺序都一样

• 只在 rank 0 打印日志、保存 checkpoint，否则 N 个进程同时写同一个文件。
• 梯度累积时，前 $k-1$ 个微批次用 with model.no_sync(): 跳过梯度同步，只在最后一个同步，省 $k-1$ 次 all-reduce。
• FSDP 用 torch.distributed.fsdp.FullyShardedDataParallel 包装模型，关键参数是按哪一层切分（auto_wrap_policy，通常按 Transformer 块）和混合精度策略。
• DeepSpeed 用一个 JSON 配置 ZeRO 阶段（"zero_optimization": {"stage": 2}）；Accelerate 在两者之上提供统一接口（accelerate config 后 accelerate launch）。

A.4 保存与恢复

• ZeRO-3 / FSDP 下每张卡只有一部分参数，保存时要么各存各的分片（恢复时要同样的卡数和切分方式），要么先汇总成完整的 state_dict（显存和时间开销大）。要能跨配置恢复，就存完整的。
• 恢复训练要同时恢复优化器状态、学习率调度器、数据加载器的位置、随机数状态，缺一个都会让恢复后的曲线和原来的对不上。

§B 实操

1. 跑通 w2day4_checks.py，7 组结论都看懂。第 3、5 组自己用纸笔重算一遍。
2. 推导（1h）：(a) 证明 $P(\min h(A)=\min h(B))=J(A,B)$；(b) 推出 ZeRO 各阶段的每卡显存公式；(c) 推出 GPipe 的气泡比例 $(p-1)/(m+p-1)$。
3. 显存计算器（0.5h）：python code/w2d4_memory_dist.py --calc --params 7e9 --gpus 8，对比全量 / LoRA / QLoRA / ZeRO 各阶段。
4. 梯度检查点与梯度累积（1.5h）：在 Week 1 的 mini-GPT 上开关梯度检查点，测峰值显存和每步时间（GPU 上）；验证梯度检查点前后的梯度逐元素一致；复现 §4.3 的平均陷阱。
5. 分布式最小例子（1.5h）：用 torchrun --nproc_per_node=2 在 CPU（gloo 后端）上跑通 DDP 和 FSDP，验证两进程训练的梯度与单进程大 batch 一致。
6. 读论文：ZeRO 论文的第 5 节（三个阶段的显存与通信分析），对照今天 §5.2。

参考实现：code/w2d4_memory_dist.py（--calc 只需 numpy，已实际运行）。

§C 自测

动机与直觉层

1. 预训练数据清洗的每一步各在防什么？为什么要去污染？
2. 为什么要做近重复去重？为什么不能两两比较？
3. 退火阶段提高高质量数据的比例，直觉是什么？
4. loss spike 有哪些常见原因和对策？为什么要监控梯度范数？
5. BF16 为什么在训练里打败了 FP16？BF16 的弱点是什么，怎么补？
6. 数据并行、张量并行、流水线并行各切的是什么？各自适合放在哪一级互联上？

推导与计算层

7. 证明 MinHash 的碰撞概率等于 Jaccard 相似度。
8. 写出 LSH 的候选概率 $1-(1-s^r)^b$，说明它为什么是 S 形。
9. 7B 模型全量训练需要多少显存？激活有多少？FlashAttention 和梯度检查点各省多少？
10. 写出 ZeRO 各阶段的每卡显存。7B、8 卡时各是多少？
11. 推出 GPipe 的气泡比例。$p=8$、$m=16$ 时是多少？
12. 梯度检查点的时间代价为什么是约 +33%？

工程层

13. 梯度累积时，变长数据上用 reduction="mean" 有什么问题？怎么改？
14. DistributedSampler 为什么每个 epoch 都要 set_epoch？
15. 梯度累积时为什么要用 no_sync()？
16. FSDP / ZeRO-3 下保存 checkpoint 有什么要注意的？

出关标准

• 能画出预训练数据管线，说清每一步防什么，包括去污染
• 能推出 MinHash 的碰撞概率等于 Jaccard，说清 LSH 的作用
• 能列出 loss spike 的原因与对策
• 能解释 BF16 为什么是训练默认精度，以及为什么还需要 fp32 主权重
• 能算出 7B 全量训练的显存（含激活），以及梯度检查点、FlashAttention 的节省
• 能指出梯度累积的平均陷阱
• 能写出 ZeRO 三个阶段的每卡显存，说出 FSDP 与 ZeRO-3 的关系
• 能说清三种并行的切法、通信与适用范围，算出流水线气泡
• 用 torchrun 跑通一个 DDP / FSDP 的最小例子

疑问清单