第一次把训练脚本从一张卡改成多卡的时候我天真地以为只要把 batch 切一半放到两张卡上就行。结果代码没改完显存先不够了改完显存速度又上不去了最后发现真正的瓶颈是通信。大模型的分布式训练说白了就是在这三件事之间反复权衡显存放不放得下、算力够不够快、通信跟不跟得上。这篇文章算是我在做大规模预训练时整理的基础理论笔记面向的读者是已经会训单卡模型、但还没系统接触过分布式并行的人。我会把并行策略、显存优化、通信原语、工程落地这些内容按我实际踩坑的顺序写下来尽量不讲废话。1. 单卡到大模型之间隔着两道墙很多人以为分布式训练是为了“加速”但真正的原因比加速更硬核单卡根本装不下也根本算不完。这两件事像两道墙一样横在中间第一道叫显存墙第二道叫算力墙等翻过去之后才发现后面还有一道通信墙。1.1 显存墙为什么参数放不进一张卡先说一组非常直观的数字。假设一个 10B100亿参数的大模型在 BF16 精度下仅权重参数就需要 10B × 2 字节 20GB。听着不多一张 80GB 的卡好像能放进去但这只是权重。训练过程中还需要维护梯度同样是 20GB 左右。而真正夸张的是优化器状态混合精度训练常用的 Adam 优化器每个参数要维护 FP32 主权重、一阶动量、二阶动量总共 12 字节。10B 参数就是 120GB。三样加一起20GB 权重 20GB 梯度 120GB 优化器状态 160GB一张 80GB 的卡连存都存不完更别提前向计算时产生的激活值了。所以当你听到类似“某某 7B 模型单卡就能推理”时那只是推理只需要权重和 KV Cache。训练完全是另一回事优化器状态是模型参数的好几倍这是最容易忽略的一点。显存墙的本质就是优化器状态 梯度 激活值让训练时的显存需求远大于模型文件大小。我们后面讲的 ZeRO、激活重计算、混合精度全都指向这一道墙。1.2 算力墙单卡跑完需要多少年显存可以靠优化硬塞但算力不行。训练一个大模型的总计算量业界有一个近似公式总 FLOPs ≈ 6 × N × D其中 N 是模型参数量D 是训练 token 总数。这个 6 的来源是前向每个 token 每个参数约 2 次浮点运算反向约 4 次。拿一个 10B 模型、训练 300B token 来算总 FLOPs 6 × 10¹⁰ × 3 × 10¹¹ 1.8 × 10²² FLOPs假设一张卡能达到 300 TFLOPS 的 BF16 稠密算力这个值已经很乐观了单卡秒数 1.8 × 10²² ÷ (3 × 10¹⁴) 6 × 10⁷ 秒大约是 694 天。也就是说单卡训一个 10B 模型要将近两年。如果换成 100B 模型二十年都不够看。所以分布式训练根本不是为了“快一点”而是为了让一个本来不可能完成的任务变得可以完成。1.3 通信墙规模化之后真正卡住你的瓶颈显存不够可以加卡算力不够也可以加卡但卡一多新的问题来了卡与卡之间要互相等数据和梯度通信开销会随着并行规模的扩大而急剧增长。我刚开始做多卡训练时犯过一个错误以为 8 卡比 1 卡快 8 倍。实际跑出来只有 4 倍多。原因很简单每一步训练结束所有卡都要把梯度汇总一遍这个汇总动作在 8 张卡之间有大量数据需要通过网络传输。卡越多迁移次数越多通信时间就越长而计算时间被切得越来越短两相对比通信占比自然越来越大。可以打个比方一个团队合作写一本书分工很容易每个人负责几章。但如果你每写完一页都要等所有人的批注汇总那写作速度就被汇总流程彻底锁死了。分布式训练里的“批注汇总”就是通信。通信墙不像显存墙那么显而易见但它才是真正决定你集群能跑到多快的关键。2. 并行策略的三种基本形态DP、TP、PP2.1 数据并行DP复制模型切分数据数据并行是最直观、最容易理解的一种并行方式每一张卡上都存一份完整的模型副本训练数据按 batch 维切成多份分发给各卡每张卡独立完成前向和反向算出梯度。然后通过 all-reduce 操作把所有卡的梯度求和再用这个总梯度去更新每一张卡上的模型。它的优点非常明显实现简单和单卡训练的逻辑几乎一样模型结构完全不用改。只要显存能放下完整模型数据并行几乎是最省事的起步方案。所以很多大模型框架在你刚开始并行训练时默认路径就是数据并行。但数据并行有两个致命缺点。第一模型复制多少份显存就翻多少倍。前面算了10B 模型训练状态大概 160GB哪怕你有 16 张卡每张卡依然要各存 160GB这显然不现实。第二通信量随着模型参数量的增加线性增长每一步都要把全套梯度同步一遍。模型越大这个 all-reduce 的代价就越大。所以数据并行适合“模型能装进一张卡”的情况但大模型恰恰装不进去。这促使我们考虑第二种方式把模型本身切开分到不同的卡上。2.2 张量并行TP切开一层算子的矩阵张量并行是把一个层内部的矩阵运算拆到多张卡上。以 Transformer 里的 FFN前馈网络为例它的核心计算是Y W₂ × activation(W₁ × X)我们可以把 W₁ 按列切开让卡 1 和卡 2 各算一半的隐藏维度W₂ 则按行切开两卡分别做矩阵乘最后通过 all-reduce 把两者结果相加得到和单卡完全一致的输出。注意这里每一步都会有大规模的矩阵计算在卡间“拼接”TP 的通信非常密集。举个例子一个隐藏层 4096 维的 FFN按 2 路张量并行切分每卡只需要计算 2048 维的部分。但每层的前向和反向都要做通信层数越多通信次数越多。这也是为什么张量并行一般只被用在同一个物理节点内——节点内的高速互联带宽高、延迟低才经得住 TP 这种“每层都通信”的模式。如果把 TP 跨到不同节点网络延迟立刻成为灾难。2.3 流水线并行PP把 Transformer 层竖着切开张量并行是横向切一层流水线并行则是纵向切层把模型的第 110 层放在卡 1第 1120 层放在卡 2以此类推。数据像流水线一样从第一阶段流到最后阶段。纯流水线并行最大的问题是“气泡”当第一个样本还在第一张卡上计算时后面的卡完全空闲。等到样本逐层传到最后的卡前面阶段可能已经算完了白白空等。怎么解决引入 micro-batch。把一个大的 batch 切成很多份小 batch依次注入流水线。第一阶段算完第一个 micro-batch 后马上算第二个不用等后面第二阶段收到第一个 micro-batch 时就开始这样各个阶段都能保持忙碌。气泡率的近似公式是气泡比例 ≈ (PP - 1) / (PP micro-batch 数 - 1)micro-batch 数量越大气泡越小。但 micro-batch 也不能无限大它受显存和梯度累计的逻辑限制。实践中PP 的切分数一般不会太大常见的是 48 段再多就可能出现调度复杂、气泡难压的问题。3. 混合并行为什么是 DPTPPP 而不是只用一种我最早有一个疑惑既然三种并行各有玩法那大模型训练到底用哪种后来和各种框架打交道多了才发现现实世界里的超大模型从来不是只用一种并行而是按“三层嵌套”的思路把三种方式组合起来。理由不复杂三种并行的通信特征完全不同分开使用会各踩各的坑组合在一起反而能扬长避短。3.1 三种并行策略的通信特征对比我把三种并行的关键特征列了个表这个表我每次做集群方案时都会拿出来看并行方式通信频率单次通信粒度通信范围主要代价数据并行 DP每步一次全量梯度全局所有卡显存复制、通信量随参数量线性增长张量并行 TP每层多次中间激活同一节点内通信极其密集无法跨节点流水线并行 PP每个 micro-batch 一次相邻阶段激活和梯度相邻阶段流水线气泡、调度复杂度只看这张表就能明白为什么实际方案都是混合的。只用 DP模型大一点显存直接爆掉梯度同步量还大只用 TP每层都要通信节点一多延迟完全扛不住只用 PP虽然跨节点可行但单卡上的计算量和显存压力没法有效分摊到更多维度。于是主流的思路就是在节点内做 TP在节点间做 PP在最外层做 DP。3.2 并行维度怎么排先 TP再 PP最后 DP很多框架在超大规模训练时使用的并行维度顺序可以概括为DDP → PP → TP也就是说最靠近单卡的是 TP然后往外是 PP最外层是 DP。为什么这样排因为 TP 的通信最频繁、最紧急必须放在速度最快的节点内互联上PP 的通信频率低可以接受跨节点延迟DP 虽然要全局同步但每个训练步只有一次 all-reduce频率最低所以放最外层哪怕走较慢的跨节点网络也划算。我自己做集群排布时的经验是先画出物理拓扑标出节点内带宽和跨节点带宽然后把 TP 放在节点内能塞多少塞多少PP 顺着节点之间排DP 在全局范围分配。这个顺序不绝对但绝大多数情况下它都是合理的。3.3 一个可参考的配置512 卡训 10B 模型纸上谈兵容易我们看一个实际配置。假设一个节点内有 8 张卡用 512 张卡训练一个 10B 模型。节点内做 8 路 TP单层计算被切成 8 份模型层间做 4 路 PP把 Transformer 层切成 4 段一个“模型副本”占 8×4 32 张卡512 张卡总共就是 512 ÷ 32 16 路 DP。最终配置是 16 路 DP × 4 路 PP × 8 路 TP。这个配置的好处是TP 完全不出节点PP 只在 4 个节点之间流动DP 则在 16 个模型副本之间做主梯度同步。通信压力被分散到不同维度不会出现某一条链路被打爆的情况。4. 显存优化ZeRO、激活重计算与混合精度并行策略解决的是“模型怎么拆”但拆完之后还有一个更底层的问题每张卡上的显存到底怎么省。这一节讲的是训练大模型几乎逃不掉的三板斧ZeRO 分片、激活重计算、混合精度。4.1 ZeRO让每张卡只存一部分模型状态前面算过一个 10B 模型权重 20GB、梯度 20GB、Adam 优化器状态 120GB总共 160GB。如果在数据并行下每张卡都存 160GB那 8 卡就是 1280GB巨大的浪费。ZeRO 的核心思路是把这些状态在数据并行的卡之间做分片每张卡只存储自己负责的那部分。它分几个级别ZeRO-1只分片优化器状态显存大幅下降通信量基本不变ZeRO-2分片优化器状态 梯度反向传播时梯度不再全量复制ZeRO-3再分片参数前向计算前通过 all-gather 把参数收集回来。代价是通信模式从“一次全量 all-reduce”变成“reduce-scatter all-gather”通信量大约变成原来的 1.5 倍左右。但换来的是显存可以被多卡线性分摊让超大模型训练成为可能。实际工程里7B、13B 这类模型用 ZeRO-2 搭配数据并行就很常见100B 以上则需要 ZeRO-3甚至配合 CPU offload把部分状态放到内存里。4.2 激活重计算拿算力换显存模型参数和优化器解决了但训练时还有一块隐形显存大户——激活值。Transformer 前向过程中每一层都要保存中间结果用于反向传播保存的大小大约是层数 × 序列长度 × 隐藏层维度 × batch size。当序列长度到 8K、32K 的时候激活值的显存会以恐怖的速度膨胀甚至超过参数本身。这时最常用的手段是激活重计算activation checkpointing / recomputation前向时不保存所有中间激活反向传播需要用激活时临时再算一次前向来重放。显存代价从 O(层数) 降到 O(1)代价是增加一次前向计算整体算力开销通常多 30% 左右。这笔账在多数情况下是划算的因为显存不够直接无法训练花 30% 算力换来能训比完全训不了强得多。我的经验是只要是大模型训练激活重计算基本一定要开但要配合没有重计算的基线先做对比避免过度重算导致 MFU 雪崩。4.3 混合精度BF16 为什么成了事实标准混合精度训练的核心逻辑是用 FP16 或 BF16 做前向和反向计算用 FP32 保存优化器状态和主权重。好处很明显显存减半、通信减半、GPU 上对应的计算单元吞吐更高。FP16 的问题是指数位只有 5 位数值范围很小梯度本来就小前向稍微大一点就容易上溢反向的梯度稍微小一点就容易下溢成 0。下溢意味着梯度信息直接消失模型根本学不动。所以 FP16 时代要配合 loss scaling 把梯度整体放大再在更新前缩小回来。BF16 解决了这个问题它有 8 位指数位和 FP32 的范围一样基本不会下溢代价是尾数只有 7 位精度差。但对大模型训练来说保持范围远重要于保持精度BF16 因此成为大模型训练的事实标准。如果你看到训练日志里 loss 突然出现 NAN先怀疑混合精度相关的数值问题这个习惯能省很多排查时间。5. 通信是隐形天花板从原语到通信计算重叠我对分布式训练认知升级最快的一段经历就是从“只看显存”变成“看通信”。这本质上是因为模型越大、并行规模越大通信占整个训练 step 的比例就越高。你不可能优化掉通信但你必须搞清楚通信发生在哪、规模多大、能不能和计算重叠。5.1 常用通信原语all-reduce、reduce-scatter、all-gather分布式训练里所有网络传输最终都落到几个集合通信原语上all-reduce所有节点把自己的数据发送出去汇总后再广播回所有节点数据并行里的梯度同步就是它。reduce-scatter数据先做 reduction再按 rank 分发分片每个 rank 只拿到结果的一小片ZeRO 的梯度分片依赖它。all-gather把各 rank 持有的分片收集齐最后每个 rank 都能拿到完整的张量ZeRO-3 在每次前向时要收集参数就用它。point-to-point只在两个指定 rank 之间点对点传数据流水线并行里相邻阶段传激活和梯度用它。理解这些原语后再去看训练日志里的通信时间就容易多了DP 的 all-reduce 是“每步一次全量”TP 的 all-reduce 是“每个算子来一次”PP 的点对点是“每个 micro-batch 一次”。你看到的任何通信瓶颈都能定位到具体是哪个环节在等哪类原语。5.2 集群拓扑决定并行编排通信不只是“带宽大就行”拓扑结构同样重要。一个节点内的 8 张卡之间通常有极高的互联带宽和很低的延迟但跨节点需要通过交换网络带宽再高也有物理上限延迟更是成倍增长。所以拿到一个集群时我第一件事就是画拓扑图哪些卡在同一个节点节点间几个上联口交换机带宽多少。然后对照并行策略把 TP 尽量锁在节点内PP 放在节点间DP 放全局。这个动作做不好后面调优会非常痛苦。我见过不少人拿着一个理论最优的并行配置跑到某个跨节点网络不好的集群上速度反而比小规模还慢根源就是拓扑和并行维度没对齐。5.3 通信与计算重叠一个真正的调优选重点很多新手看通信占比高第一反应是压缩通信数据但更高效的思路是让通信和计算并行起来。比如数据并行的反向传播传统方式是等到整张卡的反向全部结束再一次性 all-reduce 梯度这会导致每个 step 末尾出现一个明显的通信空洞。优化方式是把梯度分成很多个“桶”哪个桶算完哪个桶先发不用等全部反向结束。这样通信和后续的反向计算天然重叠整体 step 时间能显著缩短。流水线并行里的 micro-batch 调度本质上也是让不同阶段的通信和计算重叠。实践里我通常会看 profiler 输出的通信 gap如果 gap 占比超过 10%就说明通信还没被有效隐藏优先去优化调度和梯度分桶往往立竿见影。6. 工程落地最容易踩的坑理论层面了解并行策略和显存优化之后真正拉开差距的是工程细节。分布式训练的项目代码跑起来容易跑得快、跑得稳很难。我把自己踩过和帮别人排查过的坑整理成以下几个方面。6.1 慢节点木桶效应分布式训练基本是同步训练的每一个 step 都要等所有卡完成。只要有一张卡慢一点点整个集群的吞吐就被它拖到那个水平。我遇到过一次 64 卡训练吞吐比预期低了 30%。查了很久才发现是 rank 7 的数据加载线程被 CPU 别的进程抢占每 step 慢了几百毫秒。卡数越多这个木桶效应越严重。排查慢节点的方法很简单给每张卡记录每 step 耗时拉出分布找出离群值再去看是数据加载问题、显存频率降频还是网络故障。记住分布式训练里“大部分卡快”没有意义最慢的那张卡决定一切。6.2 checkpoint 与容错不能从头再来大模型训练一个任务动不动就是几天甚至几周进程崩掉、节点断网、显存 OOM 都太常见了。如果 checkpoint 设计得不好一次故障就可能让你损失好几十个小时。这里的核心原则是分布式保存分片 checkpoint每张卡只保存自己持有的那部分参数分片、优化器状态分片以及当前 step 序号不要每个 rank 都 dump 一份全量模型。恢复时按照相同的并行配置直接加载。同时checkpoint 写入本身也有开销建议用异步方式避免保存时整条流水线阻塞。我吃过一次大亏把 checkpoint 都写在同一台机器的本地盘上节点一挂备份全丢。所以 checkpoint 存储至少要有冗余最好写到独立存储。6.3 调试与观测从日志到 profiling分布式训练的调试难度和卡数成正比。一个 128 卡的任务出了问题根本没法像单卡一样直接加断点看。我的调试顺序基本固定先用极小的 batch 和极短的序列把链路跑通确认 loss 能正常下降看 loss 曲线有没有 NAN、有没有突然抖动重点查学习率、混合精度、数据顺序看每张卡的 GPU 利用率有没有某张卡一直是 0% 或者明显偏低用 profiler 看每个 step 的 breakdown找出通信占比、算子耗时、空闲 gap打印 grad norm监控梯度是否爆炸或消失。这些小技巧看着简单但能帮你快速把“不知道哪出了问题”缩小到“计算、通信、数据、数值”这四类问题里。特别是 grad norm很多人都不看它往往是判断训练是否健康的最早信号。6.4 数据加载管线最容易被忽视的瓶颈分布式训练时GPU 算力很大但 CPU 负责的 tokenizer、数据读取很容易跟不上。很多项目把 tokenizer 放在训练进程里实时处理一旦数据量大起来CPU 直接被拖垮GPU 就开始空转。我现在的基本做法是所有文本预处理尽量离线完成提前转成二进制数据格式训练时只做读入和随机索引不做重度计算。还要注意 num_workers 的设置以及随机打乱时会不会破坏样本顺序导致缓存失效。数据管线的目标是“GPU 要什么数据就在那里”而不是“GPU 等数据算完”。7. 怎么判断分布式训练做得好不好配置一堆并行策略最终要看效果。我觉得评估一个分布式训练方案就盯三个数字MFU、单卡实际算力、扩展效率。7.1 三个核心数字MFU、单卡实际算力、扩展效率MFUModel FLOPs Utilization是大模型训练里最重要的效率指标之一含义是实际完成的有效计算量除以 GPU 的理论峰值计算量。大模型训练在合理优化下 MFU 一般能到 40%55%再往上需要很强的调优。如果你看到 MFU 只有 15%大概率是通信没重叠好、激活重算开太多或者数据加载在卡脖子。单卡实际算力可以看作 MFU 的另一种表达单位是 TFLOPS。比如一张卡的峰值是 300 TFLOPS实际跑到 150 TFLOPSMFU 就是 50%。它比 MFU 更直观能直接看出“这张卡到底干了多少活”。扩展效率则反映加卡后的吞吐增益。理想情况下 64 卡应该比 32 卡快 2 倍如果只快了 1.6 倍说明扩展效率只有 80%那 20% 去哪了要么是通信太慢要么是气泡太多要么是负载不均。看到扩展效率低下先按这个思路往下拆。7.2 用 6ND 心算资源需求最后分享一个我经常用的估算方法基于前面的总 FLOPs 6×N×D。要估算一个任务需要多少卡时可以用卡时 总 FLOPs ÷单卡峰值 TFLOPS × MFU × 卡数举个例子1B 模型训 100B token总 FLOPs 6×10⁹×10¹¹ 6×10²⁰。假设 256 卡单卡峰值 300 TFLOPSMFU 0.4卡时 6×10²⁰ ÷3×10¹⁴ × 256 × 0.4≈ 1.95×10⁴ 秒 ≈ 5.4 小时这个计算能让你在设计实验之前快速判断可行性一个方案要跑多久、需要多少块卡、划不划算。虽然实际会有各种 overhead但它给了一个很稳的底线预期比凭空感觉靠谱得多。写到这里基本把我做分布式训练前的笔记倒了一遍。回想起来最想强调的还是那句话先跑通小规模再想大规模。很多并行配置在理论上是最优的落到真实集群上会因为拓扑、数据管线、慢节点完全变样。我第一次做多卡训练时以为卡数翻倍吞吐就能翻倍结果只提升了四成后来一步步定位才发现是通信没重叠好外加数据加载拖后腿。这种挫败感每一个搞分布式训练的人都经历过。希望这些笔记能帮你少走一点弯路。