★ FEATURED ARTICLE 把多个 GPU 拼起来训练大模型比想象中难而且难在“通信”和“显存”这两个字上。我练过的模型越大越觉得分布式训练不是“加卡就提速”而是要在数据、模型、序列三个维度上做取舍。这篇文章是大模型训练系列的第四篇主线很明确从 DDP 这个绝大多数人都会先接触的数据并行方案出发讲清它为什么能跑通、为什么显存瓶颈还在接着看 ZeRO 怎么把冗余的优化器状态和梯度分摊掉然后再进到张量并行、流水线并行把单层计算和按层切分的并行方式拆开看最后聊上下文并行专门解决长序列训练里矩阵算不动、显存放不下的一类问题。读完你应该能形成一张比较清晰的图什么场景该用什么并行组合以及为什么一个 70B 模型训练任务通常会同时叠加好几种并行策略而不是只靠某一种。1. 并行到底在拆什么先建立三个坐标轴1.1 单卡训练的两个瓶颈不是同一个问题先说显存。以 7B 参数模型为例如果用 bf16 做前向和反向参数本身占 14GB梯度又是 14GB。这还只是暂时性的状态真正的大头在优化器AdamW 要维护一份 fp32 的 master 权重再加一阶动量 m 和二阶动量 v每一项都是 4 字节算下来 7B 乘以 4 字节再乘以 3 份大概是 84GB。哪怕先不看激活值单卡完整训练一个 step 就需要 141484112GB 的显存。这就是为什么很多人发现 7B 模型在 40GB 的卡上做全参微调直接 OOM不是算子写得不好而是单卡纯跑完整训练状态本来就很勉强。算力是另一个维度。就算显存够前向反向一个 step 要跑 10 分钟和跑 1 秒训练效率完全是两个世界。所以分布式训练实际上同时在两条战线上干活一是把模型状态参数、梯度、优化器分散或者复制到多张卡二是把计算量分散到多卡并行执行。1.2 数据的并行、模型的并行、序列的并行围绕大模型训练业界常见的并行方式可以归纳成三类数据并行每张卡持有一份完整模型喂不同的数据批次。DDP 就是这个思路ZeRO 是对它的改进。模型并行把模型本身切开每张卡只负责一部分。按切法又分两种把某一层的矩阵切成多个块是张量并行按层的前后顺序切成多段是流水线并行。序列/上下文并行针对序列长度这个维度把长序列切成多段每张卡负责一段序列上的计算。上下文并行就是这类思路的代表。这三个坐标轴并不冲突。真正的大规模训练往往是数据并行、模型并行、序列并行同时叠加。理解每个方案的切分对象和通信频率比背一堆名词重要得多。2. DDP数据并行的“标准答案”和它的隐藏成本2.1 一次 backward 里到底发生了什么DDPDistributed Data Parallel是 PyTorch 里最常用的分布式训练接口。它的逻辑很直观假设有 N 张卡每张卡分别加载一份完整模型权重然后各自处理一个 local batch。前向互不干扰反向算出各自的梯度后执行一次全局 AllReduce把 N 份梯度求平均最后每个进程用完全相同的梯度去更新模型。关键点在于AllReduce 不只是“等梯度全部算完再做通信”而是把梯度划分成一个个 bucket。PyTorch DDP 默认以一定的桶大小bucket_cap_mb常见 25MB为单位反向传播算完一个 bucket 的梯度就立刻后台发出这个 bucket 的 AllReduce不用等所有梯度都算完。这个机制叫通信计算重叠。所以 DDP 的性能好不好很大程度上取决于你能不能把每个 batch 的样本量调大让计算时间足够覆盖通信时间。2.2 Ring-AllReduce 为什么比朴素 AllReduce 好直接在 N 张卡之间做梯度汇总最笨的方式是每张卡把梯度发给一个主卡主卡聚合后再广播回去。这个方案在主卡侧带宽压力是 O(N) 的卡一多就没法用。实际 DDP 使用的是 Ring-AllReduce思路是把所有卡的梯度按设备数切成 N 块卡与卡之间只和相邻节点交换一块数据绕成一个环。整个过程分成 reduce-scatter 和 all-gather 两段。每一段都有 N-1 次点对点通信但每张卡每次传输的数据只是总梯度的 1/N。整体算下来每张卡的有效传输量大约是模型梯度总量的 2 倍而不是 N 倍。这就是为什么 DDP 能把梯度同步做得看起来不太耗时通信量跟卡数没有线性放大只跟模型大小有关。实际经验里如果模型是 7B、梯度是 14GB那么每个 step 的梯度同步总通信量大概在 28GB 左右。在单机 8 卡、NVLink 互联的场景下这个量级很快就能跑完一旦跨机走 RDMA 网络就要看节点间带宽够不够了。2.3 DDP 的隐藏成本显存冗余DDP 最大的隐藏成本不是通信而是显存冗余。N 张卡各自保存一份完整的参数、梯度、优化器状态相当于所有状态重复保存 N 份。回到 7B 模型的例子单卡完整状态大约 112GB即便分成 64 张卡每张卡仍然要 112GB而不是变成 112/64 GB。所以 DDP 适合那种“单卡还能勉强放下完整模型”的场景。一旦模型大到单卡放不下哪怕你有 1000 张卡也没有用因为每张卡都必须有一份完整拷贝。这时候就要上 ZeRO它的核心工作就是把 DDP 的显存冗余一层层减掉。3. ZeRO把优化器状态、梯度、参数三层冗余逐一干掉3.1 三个阶段的显存账本ZeROZero Redundancy Optimizer来自 DeepSpeed思路并不是去发明一个新的并行维度而是把 DDP 里重复保存的那 N 份状态变成每张卡只保存 1/N用通信换显存。以 7B 模型、64 卡为例我粗略算一下单卡显存需求DDP 基线参数 14GB 梯度 14GB 优化器 84GB约 112GB。ZeRO-1只分片优化器状态每卡参数 14GB、梯度 14GB、优化器 84/64≈1.3GB约 29.3GB。ZeRO-2优化器和梯度都分片每卡参数 14GB、梯度 14/64≈0.2GB、优化器 1.3GB约 15.5GB。ZeRO-3参数、梯度、优化器全部按层分片每卡约 1.7GB这时完整模型状态相当于被切到 64 张卡上。注意这里没有算激活值。激活显存另算而且经常比模型状态还猛通常要靠梯度检查点、小 batch、序列切分来压。3.2 通信开销的变化不是免费午餐ZeRO 用显存换通信代价要分开看ZeRO-1 和 ZeRO-2 的通信量基本和 DDP 持平。仍然要做梯度汇总只是把 AllReduce 换成 ReduceScatter传输总量没有明显增加。ZeRO-3 的通信变多因为每个 step 在反向之前需要从其他卡 Gather 当前层权重前向要用完整参数反向算梯度后再 ReduceScatter 梯度。参数、梯度都要动态流动所以通信量大概是 DDP 的 1.5 到 3 倍卡数越多、节点间带宽越差这个开销越明显。所以 ZeRO-3 虽然显存省得最狠但对网络很敏感。跨机训练时最好做到节点内带宽够大或者开启通信重叠和预取参数之类的能力否则可能会看到 GPU 利用率长期在 30% 以下。3.3 DeepSpeed 配置里常见的坑DeepSpeed 的 JSON 配置里最常用的 ZeRO-2/3 配置大概是这样的{ train_batch_size: 256, gradient_accumulation_steps: 2, optimizer: { type: AdamW, params: { lr: 3e-4 } }, zero_optimization: { stage: 3, reduce_scatter: true, contiguous_gradients: true, overlap_comm: true, stage3_max_live_parameters: 1e9, stage3_prefetch_bucket_size: 5e7 } }这里最容易被忽略的是stage3_gather_16bit_weights_on_model_save。ZeRO-3 下权重是分片存在的存 checkpoint 时如果不先把权重 Gather 回来保存出来的模型会是残缺的看起来文件大小不对加载后损失直接崩。很多人训练到一半保存模型恢复再训练发现精度不对排查了半天最后发现是 checkpoint 里没聚合完整权重。另一个经验是ZeRO-3 下不要随意直接访问模型的完整参数。因为参数可能并不在本卡上。如果只是想看一下某个张量的数值最好用 DeepSpeed 提供的上下文管理接口临时 Gather看完再释放。否则轻则触发大量 AllGather重则直接死等。4. 张量并行把一层拆到多张卡通信嵌进每一次前向4.1 Megatron 式拆分的核心列并行和行并行当单层矩阵大到一张卡算不动时就要用张量并行Tensor Parallelism。最经典的是 Megatron 的拆法把 Transformer 里的线性层拆成两类列并行权重矩阵按输出维度切每张卡算其中一块。算完后如果要恢复完整输出就需要 AllGather。行并行权重矩阵按输入维度切每张卡算一个局部部分最后需要 AllReduce 把各部分结果累加。Transformer 里的 MLP 通常先是 h 到 4h 的列并行再是 4h 到 h 的行并行。Attention 里的 QKV 投影常用列并行输出投影用行并行。于是每一层 Transformer 大概要经历 4 次集合通信。这意味着张量并行是把通信写在了每一层、每一次前向和反向里通信频率非常高。4.2 每次前向到底传多少数据用一个具体例子感受一下数量级。假设 hidden size 是 4096一个 micro batch 的 token 数是 4096那么单个张量在 fp16 下的体积是 409640962 字节约 32MB。一层 Transformer 如果触发 4 次通信单层就是 128MB 左右。一个 40 层的模型一次前向光张量并行通信就是几个 GB 量级。这还是在 batch 不算大的情况下。所以张量并行必须跑在高速卡间互联上比如单机 8 卡的 NVLink。跨节点做张量并行不是绝对不行但通信时延会明显拖累每一层计算扩展效率掉得很快。这里有个实操原则优先保证 Tensor Parallel 的通信域不超过一个节点最好就在 8 卡的 NVLink 域内。例如 8 卡节点就做 TP8如果跨节点宁可把流水线并行放在节点之间也不要让张量并行跨交换机。4.3 序列并行降低张量并行带来的激活冗余张量并行一个很麻烦的副作用是为了做 Attention每张卡经常需要拿到完整的 hidden 状态这会放大激活显存。Megatron 后来提出了序列并行Sequence Parallel把 LayerNorm、Dropout 这类按 token 独立计算的操作按序列维度切分只在需要完整 hidden 的地方做 AllGather在 Attention 之后做 ReduceScatter。这既能把一部分激活显存压下来又能减少一部分重复通信。现在很多大模型的训练框架都把这个当作张量并行的默认实现不只是单纯做 ColumnParallel 和 RowParallel。理解这一点对排查显存问题很有用明明已经切了张量并行激活还是涨得离谱十有八九是序列并行没有开启或者序列长度被重复复制到了每张卡上。5. 流水线并行把网络按层切开吞吐由 micro-batch 和气泡决定5.1 朴素切分的问题排队等待张量并行解决了“单层算不动”但模型很深时单张卡要完整加载和计算所有层仍然吃力。流水线并行Pipeline Parallelism的思路是把网络按层顺序切成多段每张卡只负责其中一段。比如 40 层网络切成 4 段第 1 到 10 层在卡 0第 11 到 20 层在卡 1以此类推。朴素切分最直接的问题是流水线气泡。前一个 stage 没算完后面的 stage 只能空等。如果整体只算一个 batch4 个 stage 的利用率大概只有 1/4。为了让后面尽快动起来要引入 micro-batch把一个大 batch 拆成很多个小批量前一个 stage 算完一个 micro-batch 就立刻传给下一个 stage形成流水线。5.2 气泡率怎么算micro-batch 选多大气泡是流水线并行最重要的代价。假设有 p 个 stagem 个 micro-batch理想情况下整个流水线要经历 mp-1 个阶段时间其中只有 m 段时间是让每个 stage 都在有效计算其余是启动和排空阶段的空闲。近似气泡率可以写成bubble ≈ (p-1) / (mp-1)举个例子p8m16 时气泡率大概是 7/23≈30%。m64 时降到 7/71≈9.9%。所以流水线并行要尽可能多做 micro-batch把气泡率压下去。但 micro-batch 太多也有代价梯度累积步数变多显存和调度开销都会上来需要根据显存容量找一个平衡点。通常我看到有人一开始就把流水线并行分开到几十张卡上却没有把 micro-batch 调大结果吞吐极低还以为是卡间网络问题。其实先算一下气泡率把 micro-batch 数量调到 stage 数的 4 到 8 倍情况往往立刻改善。5.3 GPipe、PipeDream、1F1B 调度怎么选GPipe 是经典的 fill-drain 模式先让整个流水线跑完所有 micro-batch 的前向再统一做反向。实现简单但内存峰值高气泡也偏大。PipeDream 尝试把前向和反向交错执行但可能导致不同 micro-batch 的梯度更新顺序不一致影响收敛。实际训练里更常见的是 1F1B 调度也就是一个 micro-batch 做前向紧接着一个 micro-batch 做反向两者交错推进。1F1B 能显著降低显存峰值也能在前向和反向间更均匀地铺满计算是当前很多训练框架里流水线并行的默认调度。选调度时不要只看理论吞吐还要考虑你实际使用的是否支持 1F1B 以及梯度和参数状态的同步策略。如果用的是比较老的流水线实现GPipe 反而更省心但需要准备更大显存。6. 上下文并行长序列场景下把注意力拆到序列维度6.1 为什么长上下文需要一个新的并行维度模型参数量没变但序列变得很长比如训练 128K、1M 上下文长度的模型。这时问题主要在注意力计算。标准 Attention 的分数矩阵大小是 序列长度 × 序列长度序列长度一拉长显存和计算量都会平方级上涨。纯靠数据并行和张量并行解决不了序列维度上的超线性膨胀。上下文并行Context Parallel的思路是把一条长序列沿长度方向切成多段分给多张卡。每张卡只算自己负责的那一段 query以及对应的局部 K/V 块。不同卡之间需要交换 K/V 块才能让每个 query 看到完整的上下文。6.2 环形注意力通信和计算可以重叠最常见的实现是环形注意力Ring Attention。把序列切成 N 段每张卡维护一段 query 和一段 K/V。计算时先算本地的注意力分块然后把 K/V 按环的方向传给下一张卡环形转一圈每张卡累计完整查看一遍所有 K/V 块。每转一轮当前卡把局部注意力结果和 softmax 的局部统计量合并。因为传输的是 K/V 分块不是重复整个序列通信量可控同时 K/V 在环上传输的同时当前卡还可以继续计算已经到手的块实现通信计算重叠。这里有个数值细节softmax 的全局归一化不能简单相加。每拿到一个新的 K/V 块都要记录当前的局部最大值再对已算出的注意力分数做 rescale。环形注意力工程上难的地方就在这里分块之间如何维护数值稳定比想象中容易出错。6.3 上下文并行和序列并行有什么关系上下文并行和我们前面提的序列并行不是一回事。序列并通常指配合张量并行把 LayerNorm、Dropout 这类按 token 独立的激活切分上下文并行则是真正把 Attention 的序列计算拆到设备间多个设备共同完成同一个序列的注意力计算。实际训练长上下文模型时两者经常同时用。注意力部分用上下文并行前向其他部分用张量并行或数据并行。上下文并行尤其适合那种“单序列特别长、batch 又不大的情况”。如果序列很短比如只有 2K 或者 4K强行用上下文并行反而会因为通信开销而没有收益但当序列到了 64K 以上单卡放不下 K/V 或者注意力矩阵算不过来时上下文并行几乎是必须的选择。7. 组合与选型经验分布式策略不是越多越好而是越顺越好7.1 常规组合关系表不同并行方式解决的问题不同放在一张表里看就很清楚并行方式切分对象主要通信频率适用场景DDP数据批次每个 step 一次梯度 AllReduce单卡能放下完整模型ZeRO模型状态梯度 ReduceScatterZeRO-3 还需参数 AllGather单卡放不下完整状态张量并行单层权重每一层多次集合通信单层计算/显存超限流水线并行网络层每个 micro-batch 边界传递激活模型层数深、跨节点上下文并行序列长度注意力层交换 K/V 块超长上下文训练7.2 一个 72B 模型训练任务的实际组合假设一个 72B 模型单卡 80GB总共 8 节点 64 卡。72B 参数光是 bf16 权重就是 144GB单卡绝对放不下因此必须让模型状态在卡间分散。常见的做法是节点内做 TP8利用 NVLink。节点之间用 PP2 或 PP4把网络按层切成 2 到 4 段避免把张量并行跨交换机。其余维度走 DP配合 ZeRO-1 或 ZeRO-2 削减冗余。比如 64 卡可以设 TP8、PP4那么 DP 维度就是 64/(8*4)2再用 ZeRO 把优化器状态分片。每个 micro-batch 的激活值通过梯度检查点控制global batch 用梯度累积去凑。这种“TPPPDPZeRO”的组合是大模型训练里非常常见的基础框架。7.3 启动那一刻NCCL 的报错和拓扑教训组合越多分布式启动越容易出问题。我最常被问到的就是 “NCCL timeout”。原因往往不是网络坏了而是某个进程没有在预期时间完成初始化或者是节点间网卡名不一致导致通信域选错。排查时先开NCCL_DEBUGINFO再配合nvidia-smi topo -m看看卡之间的互联拓扑。注意容器环境里网络接口选择必要时设置NCCL_SOCKET_IFNAME指定实际的业务网卡避免把虚拟网卡和物理网卡混在一起。还有一个容易忽略的点卡数多时torchrun的MASTER_ADDR和MASTER_PORT必须对否则子进程连不上主节点整个训练会卡在初始化。另外不要盲目堆并行度。并行方式越多通信模式越复杂调度和内存碎片叠加后真实吞吐未必高于更简单的组合。我的经验是先用量化评估工具看单卡算力利用率和通信时间占比再决定要不要加一种并行方式。一个并行方案如果只带来很小的显存节省却让通信量翻倍那它就不值得上。最后再分享一个很实在的排查习惯每次改并行配置后先跑一个极小的模型、几个 step验证 loss 能正常下降再上正式规模。张量并行切分不对、流水线阶段切错、ZeRO checkpoint 没 Gather 完整这些坑在几分钟内就能暴露出来。等真跑了几天再发现浪费的时间和算力足够让人崩溃。分布式训练的路上谨慎永远不亏。
阅读完成 · 觉得有帮助？
