预训练并行化与失败训练运行笔记（12分钟阅读）

# 关于预训练并行与失败训练运行的笔记 来源：https://www.dwarkesh.com/p/notes-on-pretraining-parallelisms *我写了一些闪卡（https://flashcards.dwarkesh.com/pretraining/）来帮助自己记住下面这些内容。* 我和某位人士进行了一次有趣的对话，讨论了为什么预训练运行常常会失败。了解事情会以各种具体方式搞砸，以及为什么训练是一项如此不稳定的操作，这非常有趣。高层面上，破坏因果关系和引入偏差似乎是关键罪魁祸首。 **破坏因果关系：** - 执行专家路由时，首先通过路由器，它会给出每个token想要每个专家的分数。从这里有两条路可以走：1. Token路由，即从token的角度读取分数，并将每个token分配给其top k的专家。问题在于，这可能导致各专家之间的分配极不平衡，这对性能非常不利。或者，你可以在（且仅在）训练中采用专家选择（Expert Choice），即根据每个专家相对更偏好哪些token来拆分token。这样你可以强制每个专家得到大致相同数量的token。但最大的问题是这破坏了因果关系，因为token n被分配给哪个专家，可能依赖于token n+k会被路由到哪个专家。而破坏因果关系非常糟糕，因为你在训练中获取了（并据此更新）在部署时不会看到的信息。- 传闻这是Llama 4表现平平的原因。- 我猜你可以在prefill推理时做专家选择？但实际中，将token分配给那些在真实训练中本不会收到该token的专家，可能效果不佳。- 老实说，我不完全理解为什么破坏因果关系如此糟糕。我明白在真实推理中你无法超越因果关系。但为什么这个微小的偏差会变成一个大问题？ - 另一个可能破坏因果关系的事情是丢弃token（Token Dropping）。即专家忽略批次中它们本应处理的、但排名不高的token，而舍弃它们可以避免超出padding范围。这破坏了因果关系，因为后面一个与该专家匹配度更高的token可能导致前面的token被忽略。- 显然这是Gemini 2 Pro的一个问题。 **引入偏差：** - 偏差比方差糟糕得多 —— 方差可以平均掉，但偏差会累积放大。 - 据说最初的GPT-4训练很慢，并且一开始就搞砸了，原因是以下bug：他们在all-reduce这样的集合通信中使用了FP16。FP16的精度按对数密度分布 —— 在1到2之间，尾数位将区间划分成大约0.001的间隔。但在1024及以上，尾数可能划分成多个整数值的间隔。假设某个集合操作需要将1+1...加10000次 —— 你可能遇到这样的情况：一旦加到1024，再加1，变成1025，然后向下舍入到最近的间隔1024，再加1，以此类推。最终计算出的值可能比真实值差10倍。如果你试图将许多小梯度累加到一个大累加器中，这就是个大问题。而且可以想象这个bug有多难找！ **对AI训练的影响：** - 有些认为我们可以治愈衰老的人认为，人类因年老而死基本上有5种不同的方式（心脏病、癌症等），如果我们治愈了这5种疾病，那么基本上就解决了衰老问题。你可以对这些失败的预训练运行提出类似的问题 —— 训练运行失败是否有5种不同的方式，以至于一旦某个实验室搞定了数值问题，后续就会一帆风顺？还是说在每个新的规模级别上，你都会不断遇到新的、特殊的问题？与我交谈的人似乎认为是后者 —— 他指出，即使在数值问题内部，也有太多搞砸的方式。而且随着规模扩大，新的问题还会不断出现。 - 不看好AI很快就会完全自动化内核编写。大概是因为他认为这比一些人认为的更接近AGI完整问题。另一种思想流派认为，“嘿，哪个内核能让attention或MLP在这个scaleup上运行最快是一个极其可验证的领域，因此我们可以轻松地通过RL达到超人性能。”但他说，拥有世界上最好内核工程师的Nvidia，也花了很长时间来优化Blackwell，这实际上表明这相当困难，可能不是那么容易形成闭环的。 - 有时人们会认为，用于RL生成的推理和用于终端用户生成的推理基本相同。但这个人指出，在RL推理中，推理引擎和训练引擎之间的数值漂移会导致这些微妙的off-policy偏差，这对最高质量的训练影响很大。但如果只是面向用户服务，就不是问题。 - 强调了拥有一个有纪律的流程来整合计算倍增器（compute multipliers）是多么重要，因为存在累积具有微妙偏差的bug的风险。 --- **来自Horace He（https://horace.io/）给我和朋友们做的一次精彩讲座的笔记。** 这次讲座之所以精彩，是因为Horace将整个主题构建成了一系列问题和解决方案的链条：这是我们要做的，这是它为什么会坏，这是我们的修复方法，这是为什么这个修复最终也会坏掉。大多数解释只是罗列了一堆不同的策略，从未将它们与所解决的问题联系起来，也从未解释你为什么会选择其中一个而非另一个。 - 预训练FLOPs的公式 = 6ND。前向传播每参数每token需要2个FLOPs（乘法+加法）。反向传播是前向的2倍，因为你需要计算关于两个输入矩阵的梯度。所以2+4=6。 - 好吧，我们不能把所有事都放在一块GPU上。那么如何拆分这个问题？显而易见的解决方案是数据并行 —— 你将模型权重复制到每块GPU上，每块GPU只处理批次的一部分。- 显而易见的问题是，每块GPU的HBM是有限的 —— B300是288GB —— 随着模型越来越大，这不足以存储权重，更不用说它们的激活值了。 - 好吧，接下来我们尝试的是全分片数据并行（FSDP） —— 每块GPU只存储每一层参数的1/N —— 在处理每一层之前，你从所有GPU上all-gather该层的完整参数（每块GPU只存储每层的1/N）。处理之后，每块GPU丢弃收集到的参数。- 需要强调的是，这是默认的优先选择。你只有在拥有太多GPU、被迫前进时才会放弃它（原因稍后解释）。之所以这是默认选项，是因为很容易让计算和通信时间重叠 —— 这是因为通信的只是权重，它们不依赖于前一层发生的事，所以你可以在还在计算这一层的同时就开始all-gather下一层。相比之下，张量并行或专家并行确实需要在处理下一层之前共享前一层的激活值。流水线并行的问题是气泡（bubbles），如下所述。- 从通信量角度来看，FSDP起初看起来极其昂贵 —— 你从所有GPU上all-gather每一层的完整权重，用它们做一次矩阵乘法，然后就丢弃了。但这忽略了常规数据并行已经让你付出的代价 —— 在常规DP中，你仍然需要在反向传播的每一层之后做一次all-reduce，以便在所有的GPU上同步这批的梯度。那个all-reduce的通信量是参数量的2倍。FSDP增加了all-gather —— 前向传播每层一次，反向传播每层一次。但一次all-gather的通信量只有all-reduce的一半。所以朴素的FSDP通信量最终是参数量×4（前向和反向的all-gather，加上反向上的all-reduce）。你还可以做得更好：因为每个梯度分片只需要最终到达拥有它的那块GPU，所以用reduce-scatter替换all-reduce（跳过了最后的广播步骤）。这样你就得到了参数量×3的总通信量 —— 比普通DP多了50%的开销。 - 那么为什么你不能一直只用FSDP呢？- 通信量交叉点：你希望计算时间大于通信时间 —— 你不希望通信成为瓶颈。但由于FSDP的计算时间会随着GPU数量的增加而减少，而通信时间不会，因此当你扩展FSDP的GPU数量时，你的MFU可能会完全崩溃。当发生这种情况时，你需要同时使用流水线并行。- 计算时间 = (6 × token数 × 激活参数量) / (每GPU算力 × GPU数量) - 随着GPU数量增加而减少 - 通信时间 = (总参数量 × 3) / (NVLink域大小 × InfiniBand带宽) - 通信时间不会随着你增加更多域而增加。这让我很困惑。每个域共同拥有全部参数，你需要在反向传播的每一层之后跨域同步梯度。你可能会认为增加更多域意味着环中的跳数更多，所以all-reduce会变慢。但标准的环形算法将消息分成与参与者数量相等的块。更多的跳数意味着更多的域，但每跳的块大小按比例变小。（当块变得过小以至于每跳的延迟占据主导时，这种情况就会被打破，此时你会切换到树形算法。）- 技术上讲，你可以比简单地对所有域之间的梯度做一个单一的all-reduce做得更好。你采用分层集合通信来优化跨多个NVLink域的通信时间。关键要记住的是，域内的每块GPU都有自己访问InfiniBand的带宽。所以你想用尽它，因为互连带宽是瓶颈。你通过尽量在scaleup内部完成尽可能多的操作，然后再向外传输来实现这一点。所以你在scaleup内部做reduce-scatter，让每块GPU获得该层一个分片的域级归约梯度，然后跨域在对应的GPU之间对这些分片做all-reduce，最后在域内做all-gather。这降低了通信时间线，从而将交叉点向右移动。- 我用Cursor和Composer 2做了一个动画来说明：- 如果你看这些方程，你会看到如果增加批次大小，交叉点右移；如果让模型变得更稀疏，交叉点左移。- 这也是为什么TPU更擅长FSDP —— 因为一个域内有更多的加速器。 - 批次大小下限：FSDP是数据并行的，所以每块GPU至少处理一个序列。Attention是在序列内计算的，并且不能（轻易地）跨GPU拆分。如果你的关键批次大小是1000万个token，序列长度是1万个，那么你只有1000个序列 —— 所以即使有充足的通信带宽，纯FSDP也无法扩展到超过1000块GPU。 - 流水线并行（为解决这些问题你会在FSDP基础上添加的下一个方法）的问题：- 流水线并行的问题不同 —— 这里产生气泡的原因是，在批次开始时，专用于最后几层的GPU没有被使用，反之在批次结束时，专用于最开始几层的GPU没有被使用。你不能通过在训练中重叠批次来解决流水线气泡，因为你需要先合并梯度并更新模型，然后才能处理下一个批次。- 但你也增加了架构约束 —— 像Kimi的attention-to-residuals（每个block关注所有之前层的残差）这样的东西，当这些残差位于不同的流水线阶段时就变得非常困难。类似地，交错滑动窗口和全局attention层可能会导致阶段间的负载不平衡。处理所有这些会拖慢研究迭代，而这是你能犯的最大的罪过。 #### 关于此帖子的讨论 ### 准备了解更多？