这篇笔记探讨了大规模预训练失败的常见原因及并行训练策略的演进。训练失败的核心问题在于打破因果性和引入累积偏差,例如专家路由和数值精度问题都可能导致模型悄然失效。为了应对大规模计算,并行策略从数据并行演进到 FSDP,再到管道并行,每一步都解决了显存或通信瓶颈,但又带来了新的挑战,如流水线气泡和架构限制。这表明,随着模型和集群规模的扩大,训练的工程难题会不断以新的形式出现,而不会简单消失。
预训练失败的常见原因
训练大规模模型是一项极其不稳定的工作,失败的根源通常可以归结为两个核心问题:打破因果性和引入偏差。
破坏因果性
在训练中获取了推理时无法获得的信息,会导致模型学到错误的依赖关系。
专家路由 (Expert Routing): 在 MoE 模型中,为了平衡每个专家的计算负载,训练时可能会采用“专家选择”策略,即每个专家从所有 token 中挑选最适合自己的部分。这打破了因果性,因为一个 token 是否被处理,可能取决于它后面的某个 token。
据传,这种因果性破坏是导致某些模型表现不及预期的原因之一。
Token 丢弃: 为了优化计算,专家可能会忽略掉那些匹配度不高的 token。这种做法同样会破坏因果性,因为一个较早的 token 是否被处理,可能取决于一个较晚的 token 与该专家的匹配度。
一个核心疑问是:为什么这种微小的因果性破坏会产生如此严重的影响?尽管在推理时无法“看到未来”,但在训练中接触到这类信息似乎会导致模型产生难以修复的偏差。
引入偏差
与可以被平均掉的方差不同,系统性的偏差会不断累积,最终导致训练结果严重失真。
数值精度问题: 这是一个典型且难以发现的偏差来源。例如,在进行梯度累加时使用 FP16 精度。FP16 在数值较大时,其表示精度会急剧下降。
- 当累加器的值超过一个临界点(如 1024)后,每次加上一个很小的值(如 1),可能会因为四舍五入而被忽略,导致最终的累加结果与真实值相差巨大。
- 这种微小的、持续的计算偏差会像滚雪球一样越滚越大,最终毁掉整个训练过程。
推理引擎与训练引擎的差异: 在强化学习(RL)中,如果用于生成数据的推理引擎和训练引擎在数值计算上存在微小差异,这种差异会引入策略外偏差 (off-policy biases),对最终模型的质量产生重大影响。
对 AI 训练的启示
有人认为,训练失败的原因是有限的,只要解决了几个关键问题(如数值精度),后续的训练就会一帆风顺。但更可能的现实是:
随着模型和集群规模的每一次扩大,都会涌现出全新的、更微妙的工程难题。问题的种类不会收敛,只会不断演变。
这表明,自动化编写底层核心计算代码(Kernel)可能比想象中更难,它更接近一个 AGI-complete 的问题。即使是拥有全球顶尖工程师的英伟达,也需要花费大量时间来优化其硬件,这说明了问题的复杂性。
预训练中的并行策略
预训练所需算力巨大,其基本公式为 6 * N * D(N 为参数量,D 为 token 数)。因此,必须将计算任务拆分到多个 GPU 上。
从数据并行到 FSDP
数据并行 (Data Parallel): 最简单的策略。在每个 GPU 上复制完整的模型权重,并将数据分批处理。
- 瓶颈: GPU 的显存有限,无法容纳越来越大的模型权重和激活值。
完全分片数据并行 (FSDP): 解决显存瓶颈的方案,也是目前的默认首选策略。
- 工作方式: 每个 GPU 只存储模型参数的一小部分(一个分片)。在计算每一层时,通过 All-gather 操作从其他 GPU 获取完整的层参数,计算完成后立即丢弃。
- 通信开销: FSDP 的通信成本看起来很高,但相较于传统数据并行(DP)在反向传播时需要 All-reduce 同步梯度,FSDP 通过将 All-reduce 优化为 Reduce-scatter,总通信开销仅比 DP 增加了约 50%。
- 优势: FSDP 可以将计算与通信重叠。因为下一层的权重与当前层的计算无关,所以可以在计算当前层时,提前开始为下一层收集权重。
FSDP 的瓶颈
当 GPU 数量增加到一定程度时,纯 FSDP 会遇到瓶颈。
通信交叉点 (Comms Crossover):
- 随着 GPU 数量增加,每个 GPU 的计算时间会减少。
- 但跨节点的总通信时间基本不变。
- 当通信时间超过计算时间时,GPU 将会等待数据,导致计算效率(MFU)急剧下降。此时,就需要引入新的并行维度。
批处理大小下限 (Batch Size Floor):
- FSDP 本质上仍是数据并行,每个 GPU 至少要处理一个完整的序列(sequence)。
- 如果你的总批处理大小是 1000 万 token,序列长度是 1 万,那么你只有 1000 个序列。这意味着,使用纯 FSDP 最多只能扩展到 1000 个 GPU。
管道并行及其问题
为了解决 FSDP 的扩展瓶颈,通常会引入管道并行 (Pipeline Parallelism),即将模型的不同层分配到不同的 GPU 组上。
主要问题:流水线气泡 (Pipeline Bubbles)
- 在每个批次开始时,负责模型后半部分的 GPU 处于空闲状态;在批次结束时,负责前半部分的 GPU 会先空闲下来。这种空闲时间就是“气泡”,会降低整体的计算效率。
- 在训练中,由于需要等待整个批次的梯度计算完成并更新模型,无法通过简单重叠不同批次来完全消除气泡。
架构限制: 管道并行给模型架构创新带来了巨大限制。例如,当模型的不同层分布在不同的计算节点上时,实现跨层连接(如 Kimi 的注意力机制)会变得非常困难和低效。这会严重拖慢研究迭代的速度。