☰
大模型3D并行训练怎么配才不浪费算力
2026/10/4 7:04:01 网站建设 项目流程

摘要

模型大到单卡放不下,就必须把张量并行(TP)、流水线并行(PP)和数据并行(DP)组合起来用。难点不在概念,而在怎么配:TP 开多大,PP 切几段,剩下的卡给 DP,三个数一乘必须等于总卡数。这篇文章按"先放得下,再跑得快"的顺序,讲清每一维并行的代价和一套可以照着走的调优流程。

背景与问题

前面几篇讲过 DDP、FSDP、DeepSpeed ZeRO 和 Megatron 的基本用法。它们各管一段:数据并行解决"吞吐",ZeRO 和 FSDP 解决"状态冗余",而当单层权重或整个模型的激活都放不下时,就得切计算本身。

实际项目里没有谁单独够用。官方指南给出的关系式是:

总 GPU 数 = TP x PP x CP x EP x DP

不用长上下文和 MoE 时,CP 和 EP 都是 1,公式退化成 DP = 总卡数 / (TP x PP)。所以所谓"调优",本质是在满足显存约束的前提下,选一组 TP 和 PP,让通信开销最小、DP 尽量大。

核心思路与优势

三个维度各自的代价

维度切什么通信方式主要代价
TP单层内部的矩阵每层前后向都有 all-reduce,频率最高对带宽极敏感
PP按层分段只在相邻段之间传激活,量小流水线气泡,空转的卡
DP数据批次每步一次梯度同步,可与计算重叠每份副本都要装下模型分片

由此得到第一条铁律:TP 放在节点内,PP 和 DP 跨节点。TP 的通信太密,必须走 NVLink;跨节点的 InfiniBand 或 RoCE 带宽低得多,只适合 PP 的点对点传输和 DP 的梯度同步。常见的 8 卡节点上,TP 一般不超过 8。

流水线气泡与交错调度

PP 切成 p 段、一个批次拆成 m 个微批次时,非交错的 1F1B 调度中气泡占比约为 (p-1)/m。要压低气泡,有两个办法:

  • 增大微批次数 m,也就是增大全局批大小或减小微批大小
  • 开启交错调度(虚拟流水线),每张卡负责多段不连续的层,气泡约降到原来的 1/v,代价是通信量相应增加

Megatron 的论文报告,交错调度在显存占用相当的情况下能把吞吐提升 10% 以上。对应的参数是--num-layers-per-virtual-pipeline-stage。

配套的三个开关

  • 序列并行:--sequence-parallel,把 LayerNorm 和 Dropout 的激活沿序列维度切开。用了 TP 就应该打开,TP 与专家并行同时使用时则是必选项
  • 分布式优化器:--use-distributed-optimizer,把优化器状态在 DP 组内切分,相当于 ZeRO-1,省下的显存可以换更小的 TP 或 PP
  • 通信重叠:--overlap-grad-reduce、--overlap-param-gather、--tp-comm-overlap,让通信藏在计算后面

面向人群

  • 已经会用 DDP 或 FSDP,开始碰到单卡放不下的模型的工程师
  • 在 K8s 或 Slurm 集群上负责训练平台,需要给不同模型定并行配置的人
  • 想读懂 Megatron 论文和官方文档里那些并行参数的人

实践步骤

第一步:先算清楚放不放得下

用 bf16 加 Adam 估算,每个参数约需要 16 字节(权重、梯度、优化器状态合计,混合精度下常见的估法;若梯度按 fp32 累加,则约 18 字节)。70B 模型光是这些状态就超过 1TB,远超单卡 80GB。所以必须靠 TP x PP 把模型切成能装进单卡的分片,分布式优化器再把优化器状态在 DP 组内摊薄。

第二步:确定 TP

从节点内卡数出发:8 卡节点先试 TP=4 或 TP=8,官方指南把 TP 的适用场景定为 hidden size 4096 以上。hidden 很小的模型开大 TP,通信占比会盖过计算收益,反而变慢。

第三步:确定 PP

TP 切完仍放不下,再加 PP。PP 只在模型层数足够多(官方建议 50 层以上)时才划算,并且要保证每段层数大致均衡。PP 能不加就不加,每多一段,气泡就多一份。

第四步:剩下的给 DP,再调批大小

DP = 总卡数 / (TP x PP)。然后检查全局批大小能否被 DP x 微批大小整除,并保证微批次数 m 远大于 PP 段数 p,否则气泡会很大。

第五步:一个官方示例

官方指南里的 LLaMA-3 70B、64 卡示例(这里额外加上了指南推荐的--sequence-parallel;这只是并行相关参数的节选,真正跑起来还要补上数据、分词器等参数,并在 8 个节点上分别用 torchrun 配好--nnodes与 rendezvous 参数启动):

torchrun--nproc_per_node=8pretrain_gpt.py\--tensor-model-parallel-size4\--pipeline-model-parallel-size4\--context-parallel-size2\--num-layers80\--hidden-size8192\--num-attention-heads64\--seq-length8192\--sequence-parallel\--micro-batch-size1\--global-batch-size512\--bf16

这里 TP=4、PP=4、CP=2,DP 就是 64 / (4 x 4 x 2) = 2。CP 是上下文并行,只在序列很长(官方说明针对 8K 以上,这个示例的序列长度正好是 8192)时才需要,序列不长就设为 1,把卡留给 DP。

第六步:看指标再迭代

先别猜,看三个数:单卡 MFU、显存峰值、各 rank 的通信耗时。显存有富余就把 TP 或 PP 降一档,把卡让给 DP;MFU 低且 PP 气泡明显,就开交错调度或增大微批次数;通信耗时集中在 TP all-reduce,就确认 TP 组没有跨出节点。

官方给的总原则很朴素:先只用 DP,放不下再加 TP,模型更大再加 PP,序列很长再加 CP,能不加的维度就不加。论文中的万亿参数模型在 3072 张 GPU 上跑出 502 petaFLOP/s,单卡吞吐达到理论峰值的 52%,说明这套组合在超大规模下也能把算力吃满,前提是每一维都配得克制。

我的看法

3D 并行的调优没有万能配方,但有稳定的顺序:先让模型放得下,再让通信不跨越慢链路,最后用 DP 吃掉剩余的卡。多数人踩坑,是一上来就把 TP 开得很大,或者为了省显存无脑加 PP。反过来想,每一维并行都是一笔通信账,能少付就少付。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询