第一次自己动手去跑一个十亿规模参数模型的训练脚本时,我对分布式训练的理解其实很粗糙:显卡不够,那就多加几张。真把多卡环境搭起来以后才发现,分布式训练远不是“卡越多跑得越快”这么简单,显存、通信、负载均衡、精度误差会一起找上你。这篇笔记是把我从实际踩坑出发,摸到的那套大模型分布式训练基础理论整理出来的备忘,适合刚接触大模型的同学,也适合那些开会要跟训练团队对齐进度、但自己又没逐行写过分布式代码的算法工程师。我会尽量按“为什么需要、拆成哪几种并行、显存怎么省、通信怎么花、怎么排查问题”的顺序来梳理,保证你看完以后至少不会再被一堆“分片”“并行”“同步”的术语吓住。
1. 分布式训练到底在解决什么问题
1.1 单卡训练碰到的那堵墙:显存
分布式训练听起来很抽象,但出发点非常朴素:一张显卡放不下模型,只能把模型和计算拆开,放到多张显卡上协作完成。做大模型的第一课往往不是网络结构怎么设计,而是先搞明白参数规模意味着什么。
我见过很多人在单卡上硬跑大模型,结果要么是启动时报显存不足,要么是训练到中途就被操作系统杀掉。原因不是某张卡坏了,而是参数规模把显存吃完了。这里可以做一个非常简单的估算。假设模型参数量是 W,用普通单精度浮点数存储,每个参数占 4 个字节,那么模型权重本身就占据 4W 字节。训练过程中,反向传播需要保存梯度,这部分又要 4W 字节。如果优化器采用带一阶动量和二阶动量的方案,状态占用还会再多出 8W 字节。换句话说,训练状态总占用大约等于 16W 字节。一个 10B 参数的模型,光是权重、梯度和优化器状态就接近 160GB,这还没算前向计算时产生的中间激活值。单卡显存通常只有几十 GB,所以根本装不下。
很多人会问,那能不能不做反向传播,只做推理?推理时确实只需要权重,不需要梯度和优化器状态,所以显存压力小很多。但训练是另一回事,你必须保存梯度,也必须保存优化器状态。这也是为什么大模型训练几乎离不开分布式,不是算法者选择分布式,而是显存上限逼着你这么做。
1.2 算力墙:单卡训练时间等不起
显存只是第一个约束,时间约束更现实。假设一张显卡每秒能完成 100 万亿次浮点运算,也就是 100T FLOPS 左右,而一个 10B 参数模型前向加反向需要差不多 2 倍参数的浮点运算量,即 20T FLOPS 左右。那么理论上一轮迭代也要 0.2 秒。如果训练数据有几亿个 token,每个 token 都要过一遍模型,你还需要几十万次迭代。单卡处理上万次迭代要用很多天,一万次迭代就可能需要几周甚至几个月。
显存不够还能靠换更大的卡解决,但时间不够就只能靠堆更多的算力。分布式训练正是为了同时解决显存和算力这两个问题而存在的。它把一个大任务切分到多台设备上,让每台设备只处理一部分工作,再通过协调机制把这些工作结果拼起来。至于怎么切分,是分布式训练理论里最核心的内容。
2. 并行策略的基本分类:数据、模型、流水线
2.1 数据并行:每张卡都有一份完整模型
数据并行是最好理解的一种并行方式。它的核心思想是:让多张卡各自持有一份完整的模型参数,然后把一个大的批次拆成多个小的批次,分发给不同显卡。每张卡独立完成一次前向和反向计算,得到各自的梯度,最后通过通信把这些梯度做一个全局平均,再让所有卡用平均后的梯度更新模型。
我最初接触分布式训练时,第一个用的就是数据并行。它的优点是代码改动小,训练逻辑和单卡几乎一致,不足也很明显:每张卡都要存完整的参数、梯度和优化器状态,显存浪费比较严重。训练大规模模型时,数据并行通常只用在最外层,因为它本质上没有降低单卡显存的压力。对于参数量超过显卡容量上限的模型,单纯用数据并行没有意义,因为模型根本加载不进任何一张卡。
数据并行里还有一个容易被忽视的细节:同步训练和异步训练的区别。同步训练指所有卡算完当前批次后一起求平均,再更新参数,进度一致,数学上接近单卡批量增大后的梯度;异步训练指各卡按自己的节奏更新,虽然吞吐更高,但可能让模型陷入参数不一致的问题。现在工业界的多数大模型训练都采用同步训练,因为它更好收敛、结果更可控。
2.2 张量并行:把一层切成左右两半
当模型单层因为太大而放不进一张卡时,数据并行就处理不了,得考虑张量并行。所谓张量,通俗说就是权重矩阵。张量并行会在模型的一层内部做切分,把矩阵乘法拆成多块执行。
我习惯用一个例子来理解它。假设有一个全连接层,输入维度是 4096,输出维度是 4096。正常情况下,权重矩阵是 4096 乘 4096,输入和权重相乘得到一个 4096 维输出。张量并行会把这个权重矩阵按列切成两块,每块都是 4096 乘 2048,分别放在两张卡上。输入向量也要复制给两张卡。每张卡用自己那一半权重计算出一个 2048 维中间结果,最后再把两边的结果拼起来。这种切分方式把计算负载分散到了两张卡上,每张卡只需存储一半权重。
张量并行最大的挑战是通信特别频繁。因为前向的每一步都需要把分块的中间结果拼起来,反向还要再把梯度分回去,这导致每次计算几乎都伴随一次全局通信。如果卡与卡之间的带宽不够,通信开销会直接抵消计算加速效果。所以张量并行更适合在高带宽的卡间进行,比如同一台机器上的多个显卡,或者使用高速互连拓扑的小规模集群。它的切分维度比较细,适合超大模型但在工程实现上也是最容易出问题的部分。
2.3 流水线并行:把网络按层切开接力跑
如果说张量并行是横向切分一层,流水线并行就是纵向把整个模型按照层数切开。前若干层放在第一张卡,中间若干层放在第二张卡,后面的层放在第三张卡。数据像流水线一样依次通过每一张卡,就好像工厂车间里分几个工位,每个工位只处理自己负责的工序。
流水线并行理解起来不复杂,难点在于产能不均衡。前向计算是顺序经过各层的,后向计算也要倒序回来。如果某个工位耗时特别长,后复工位就得等待,这种空闲时间在专业术语里叫“气泡”。减少气泡的方法是引入微批次,也就是把一个批次进一步切成更小的微型批次,让多个微型批次轮流填充流水线,而不是整个批次都走完之后才开始下一个批次。
流水线并行减少的显存非常明显,因为每张卡只负责存储部分层。但它的通信频率比数据并行低很多,因为只需要在相邻层的边界处传递激活值和梯度。所以流水线并行适合跨设备网络带宽一般的环境,比如多台服务器拼在一起形成一个训练集群。
3. 显存优化的进阶手段,让并行策略跑得动
3.1 混合精度:一半显存、更高吞吐
并行策略解决的是“放不下”的问题,显存优化解决的是“卡里能省就省”的问题。实战中最重要的显存优化手段是混合精度训练。核心思路很简单:前向和反向计算都用低精度浮点数,比如半精度,而权重更新、某些稳定性要求高的计算仍然保留单精度。
为什么可以这么做?因为深度学习中网络参数存在大量冗余,低精度带来的误差在可接受范围内,有时候还能因为内存带宽占用降低而获得明显的速度提升。半精度浮点数只占 2 字节,相比 4 字节的单精度直接省了一半,动态范围足够覆盖大部分梯度值。不过半精度有一个坑:梯度过小时容易下溢成零,导致参数更新失效。所以实际训练时会给损失乘一个大数,比如 1024 或 4096,让梯度放大到安全范围内,同时损失缩放的效果在很多框架里已经自动化。
还有一种变体是“脑浮点”,它比半精度有更大的指数范围,但尾数精度更低。训练稳定性方面更省心,因为不容易出现梯度下溢,不过在需要高精度推理和数值敏感场景下要额外验证效果。老实说在训练大批次模型时,主流方案普遍偏向这个变体,因为它对超参的敏感度更低。
3.2 梯度检查点:用计算换显存
另一个非常实用的显存优化手段是梯度检查点。正常反向传播需要知道每一层的中间激活值,以便计算梯度。显存里就要保存所有层的激活值。如果模型很深,这些激活值加起来会非常庞大。梯度检查点选择不保存全部激活值,只保存少量关键节点,反向传播需要梯度时,再从前向的检查点重新计算一次激活值。
这是一种典型的“用计算换显存”策略。代价是前向会被多跑一遍,训练速度会有一定下降;收益是显存占用能下降一个量级。在模型大到必须减少显存占用,但又不希望拆分成更好并行度的层次结构时,梯度检查点很有价值。实践中我常把它和张量并行一起用,因为张量并行会让激活通信变多,梯度检查点又会省下不少激活显存,两者是一对不错的组合。
3.3 参数分片:雕梯度和优化器状态也拆开
传统数据并行最大的问题是每张卡都保留完整模型副本。如果能把参数、梯度和优化器状态按卡号分片存放,每张卡只存其中一部分,需要用到全局参数时再通过通信收集,这就能让整体显存占用变成原来的几分之一。
这种思路在很多分布式训练方案里叫参数分片,也常被称为分片策略。它有三个递进阶段:第一阶段只分片优化器状态,第二阶段只分片优化器状态加梯度,第三阶段把参数、梯度、优化器状态全都分片。阶段越靠后,单卡显存省得越多,但通信开销也会增加。因为每个节点在计算时如果缺少某些参数,必须先从持有它的节点拿过来,算完之后还要把梯度再传回去。
我实际跑大模型时,最常用的组合是:外层数据并行,层间流水线并行,内层张量并行,再叠加参数分片和混合精度。不要小看这种多重组合,每多一层并行,代码里的通信策略、显存计划和错误排查难度都会同步变大。初学者应该先从数据并行加混合精度跑通,再逐步引入分片和流水线并行。
4. 通信开销:分布式训练里隐藏的成本大头
4.1 集合通信原语:数据并行依赖的底层件
并行策略能够运转,全靠通信把不同显卡的计算结果合并起来。分布式训练里最常接触的几种通信操作包括:广播、全归约、全收集、归约分散。它们的名字容易混淆,但核心逻辑很明确。
广播把一份数据从源节点发给所有节点。数据并行反向传播得到的梯度要全局平均,就会用到类似全归约的操作:它把所有节点的数据按元素相加或取平均,然后把结果分发给所有参与的节点。全收集则是把各节点的数据收集起来再分发给所有人,常用于张量并行中去拼激活结果或梯度分片。归约分散介于全归约和全收集之间,它把数据归约后分散存储到各节点,而不是把完整结果还原给每个人,参数分片会用到这类操作。
这些通信原语听上去很底层,但它们决定了训练速度的上限。之前有个错误认知是做数据并行时,只要每张卡算自己的梯度,最后同步一下就完事。真要同步全部梯度时,通信量会非常大。一个 10B 模型每轮都要同步梯度,假设梯度存储为 40GB,几十个节点同步一次的通信量就是几百 GB。这样就理解为什么在高性能计算里,网络互连方案、通信库和通信调度必须重点调优。
4.2 环状全归约:把通信量摊平
朴素的同步方式是一张卡把结果发给所有卡,或者把所有卡结果汇总到一张卡再广播回去。这两种方式虽然能正确工作,但通信瓶颈极其明显。工业界更常用的是环状全归约。
环状全归约的逻辑是用一张环状拓扑把所有节点串起来。第一步,每个节点把自己数据切成与节点数数量相同的块;第二步,每个节点只把其中一块传给下一个节点,并接收上一个节传过来的块,一边接收一边做局部累加。最终每个节点会得到一块完整归约结果;第三步通过全收集过程,把自己手里的块发给其他节点,同时从其他节点收集缺失的块。这样总通信量比朴素的星形拓扑少一个量级以上。
这也是为什么大模型训练集群要特别关心拓扑结构。节点之间在物理上如果已经能保证高带宽直连,通信调度就容易得多;如果跨节点走普通网络,环状通信虽然能减少总传输量,也可能受单跳延迟影响。所以设计训练集群时,常见的做法是尽量把张量并行限制在同一台物理机内部,机间只做数据并行和流水线并行。
4.3 通信和计算重叠共用
通信归通信,计算归计算,如果两者串行执行,性能损失很大。实际工程里引入了一种技巧叫通信和计算重叠。最典型的是在反向传播过程中,某一层的梯度计算完成后,立刻启动通信操作把它同步出去,而不用等所有层全部算完。
以前我把所有层的梯度攒到一个桶里,直到整个反向传播结束再统一通信,结果每轮等待通信的时间非常久。后来把梯度按层切分成多个小的通信块,一旦某几层的梯度计算完成就立即通信,整体耗时明显下降。这里需要注意一个矛盾:通信块切得越小,越容易和计算重叠,但通信次数变多,会引入额外的调度开销。所以工程上常见做法是设置一个合适的桶大小,让通信次数、通信数据量和计算重叠达到平衡。你可以通过控制脚本里的梯度聚合参数来调这个值,也可以观察训练日志里的通信占比判断是否设置合理。
5. 三种并行策略怎么搭配,才能不互相打架
5.1 从单层并行到多维并行
前面讲的三种并行并不是互斥的,实际上它们处于不同的切分维度,能组合使用。数据并行解决的是“多机多卡如何并行处理多个批次”,张量并行解决的是“单层参数太大怎么切”,流水线并行解决的是“模型很深怎么按层分工”。
典型的大模型训练组织形式是:先按流水线把整个模型切成若干大段,每段分配给一组设备;组内再用张量并行把每一层横向展开;最外层配合数据并行复制多份完整流水线执行相同的训练流程。业界的叫法可能是“3D 并行”,也就是维度上的叠加。
这样设计的原因是显存约束和并行效率天然有层级关系。张量并行通信最密集,只适合放同一台设备内;流水线并行通信频率较低,可以跨设备;数据并行对通信带宽的敏感性居中,经过优化后可以大规模扩张。训练超大模型时,可以考虑把流水线作为主切分维度,各设备组之间用数据并行叠加,卡内部的关键大层再上张量并行。这里有个简单的工作流程:先估算单卡能承载多少层,再决定切成几个流水线阶段,最后看每个阶段的张量是否大到还需要张量并行。
5.2 选型的四个判断依据
面对一个具体大模型时,怎么选并行策略?我总结了几条比较实用的判断依据。
第一,先看模型能不能单卡完全加载。如果加载不了,优先考虑参数分片和流水线并行。第二,看单层计算量和显存占比。单层矩阵规模极大,就用张量并行;单层不大但层数很多,就优先流水线。第三,看设备间带宽。同一台物理机内部带宽高,张量并行放心用;跨物理机带宽低,尽量把张量并行限制在单机内,多用数据并行和流水线并行。第四,看训练脚本的可维护性。张量并行和流水线并行对代码侵入性更强,改动面更大,数据并行和参数分片则相对容易切换到现有代码。
我把常见的并行策略粗略画成下表,供参考:
| 策略 | 显存压力 | 通信频率 | 代码复杂度 | 适用场景 |
|---|---|---|---|---|
| 数据并行 | 高,每卡完整副本 | 中高,每步同步梯度 | 低 | 模型单卡放得下,想提升吞吐 |
| 张量并行 | 中,单卡只存一层切片 | 极高 | 高 | 超大单层,高速卡间互联 |
| 流水线并行 | 低,每卡只存部分层 | 低,边界传递激活值 | 中 | 层数多,显存不够且带宽一般 |
| 参数分片 | 低,权重梯度状态都被分片 | 中高,参数使用前恢复 | 中高 | 单卡放不下,又不想引入高复杂并行 |
5.3 混合并行需要警惕的隐性负载
混合并行最怕出现“某个环节的计算负载明显偏大,其他环节空转”的情况。比如流水线并行中,如果前几层计算复杂后几层计算简单,那么后几层会大量空闲。张量并行中,如果每一层切分方式不一致,切到某层时通信量特别大,整卡都有可能被通信卡住。所以使用混合并行时,不能只看总参数能否放下,还要关注每一层的算子特性。
实际操作时,我发现一个比较有效的做法:在正式训练前,先用小规模数据跑一个简短的“基准测试”,不断调整并行维度、微批次个数和参数分片策略,观察每一张卡的利用率和显存曲线。只要有一张卡的显存接近上限,或者利用率持续比别的卡低,就说明并行策略还有优化空间。很多人直接把大模型训练脚本部署上去就跑,出了性能瓶颈后却不知道问题出在计算、通信还是拓扑,这种痛苦我也经历过。
6. 实操过程中常见的故障与排查心得
6.1 初始化阶段:最常栽在拓扑与端口设置上
分布式训练很容易在启动时就报错,原因通常不在模型,而在进程之间如何识别彼此。多机训练时,每个进程要有全局编号、机器编号和本地编号。如果这些编号对应关系配错,某些进程接收到的数据来自错误的节点,导致训练结果错乱或直接卡死。
另一个经典问题是没有设好通信所需的地址和端口。有些集群网络环境下机器之间默认不通,容易造成超时或连接拒绝。我常用的排查办法是:先在两台机器上直接测试端口连通性,再启动训练脚本。如果端口有问题,而不是先怀疑模型代码,能省下大量时间。
初始化阶段还容易出现批次大小的分配问题。数据并行要求各卡分到的批次大小尽量一致,否则梯度平均会出现偏差。部分初学者会把全局批次大小除以卡数得到每卡批次大小,却忘了在脚本里同步修改数据采样器,最后出现部分卡数据重复、部分卡数据缺失的诡异现象。解决方案是仔细检查数据加载器是否设置了分布式采样,同时打印每卡实际处理的数据量对比一下。
6.2 训练过程中的数值不稳定:损率突然变大,甚至变成 NaN
损失函数突然变成 NaN 是分布式训练里最让人头疼的问题之一。很多人第一反应是学习率过大,但我实践中发现,混合精度和通信叠加时,数值不稳定更容易出现。
混合精度训练中,梯度可能因为规模较小而丢失精度,特殊精度类型的动态范围又有限。如果损失缩放设置不合理,反向传播得到的梯度不是溢出就是被压低到几乎不可用。解决思路是先检查损失缩放值,或者直接切换为对精度要求更低的训练再观察是否仍出 NaN。如果是分布式训练,还要检查是否把 fp32 主权重保存和更新正确,不要在低精度参数上执行优化器更新。
另一个容易被忽视的原因是梯度累计和分布式平均的相互作用。梯度累计时得到的是局部梯度总和,而分布式同步又做了平均。如果这两个步骤的缩放系数不一致,实际更新步长就会莫名其妙变大或变小。调试时可以把梯度累计步数先设为 1,保证每步的梯度解释是清晰的,再做其他调整。
6.3 性能分析与资源利用率的判断
当训练速度远低于预期时,第一步不是调代码,而是正确判断瓶颈在哪。先观察计算资源利用率,如果很多显卡的利用率都很低,而网络传输量很高,大概率是通信瓶颈。如果显卡利用率很高,但显存接近上限波动,那就是显存调度有瓶颈。
通信瓶颈可以尝试增大每卡批次、减少通信次数、调整梯度通信块大小、或者把张量并行从跨机改为机内。显存瓶颈则可以考虑开启梯度检查点、调整激活重计算策略、换用更低精度的参数类型。计算瓶颈往往能通过优化模型结构或者算子融合解决,这就和并行策略无关了。
我自己习惯在训练脚本里定时输出每张卡的显存占用、通信等待时长和训练吞吐。如果某个 rank 显存占用特别高,而其他 rank 空闲,就需要检查是不是存在负载不均衡。比如流水线并行时各阶段 micro 批次数目不对,就会出现前面卡忙死、后面卡忙闲不均的现象。
6.4 分布式训练里那些容易被忽略的小细节
有几个小问题排查成本低,但非常影响训练稳定性。
第一个是随机种子。每张卡需要设置不同的随机种子,才能保证数据增强和模型初始化结果不同而导致的行为不一致被及时控制。但全流程又需要整体可复现,所以通常主节点有固定种子,各卡的局部随机源再根据机器编号派生。
第二个是日志同步。多个进程同时写日志到同一个文件时,日志会相互覆盖,严重时影响性能。最好让每个 global rank 写独立文件,或者只在 rank 0 输出主要信息。
第三个是容错机制。训练到一半网络抖动导致某节点断连,如果没有检查点和重启机制,几天时间就白费了。分布式训练正确配置并不是一劳永逸,还要做好滚动保存和断点恢复。至少保证每几百步保存一次模型权重和优化器状态,并定期验证恢复流程能正常工作。
7. 给初学者的建议:从最小集群开始验证
在动手搭大规模分布式训练之前,我强烈建议先用两台机器或者一台机器内的两张卡跑通最小实验。很多人一上来就上几十张卡,一旦报错,问题维度太多,根本定位不了。最小集群能让你先验证通信、拓扑和数据加载是否正确,再逐渐扩大规模。
我在实际学习时有几个经验可以分享一下。第一,一定要先跑一个几百步的稳定性测试,观察损失是否正常下降,再开始正式训练。第二,并行策略从最简单数据并行开始,不要直接上三维混合并行,因为混合并行出问题时你很难判断是某一层写错还是通信时序问题。第三,训练中遇到阈值变换,先只调整一个变量,不要同时改批次大小、并行策略和学习率,否则即使改对了也不知道是哪个变量起的作用。第四,每次实验后把关键配置和现象记录到笔记里。分布式训练里很多“看起来一样”的配置,在真实集群下会因为拓扑、驱动版本和显存型号差异,展现出完全不同的性能表现。
参数分片、混合精度和通信库的组合使用,每一步都需要吃透原理。小规模验证虽然不是最酷的部分,但它能保住你后面跑大型实验时真正需要的那份确定性。从项目管理的角度,过早投入大规模集群,不仅资源消耗高,而且会让人把大量时间浪费在排查琐碎环境问题上,而不是思考模型本身和训练策略。
对刚起步的人来说,分布式训练最该记住的是一句话:先跑通,再跑快,最后把容量和规模讲出来。这句话听起来朴素,但我个人体会里,它比任何高级调优技巧都管用。