ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

大模型3D并行训练怎么配才不浪费算力

大模型3D并行训练怎么配才不浪费算力 摘要模型大到单卡放不下就必须把张量并行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 卡节点先试 TP4 或 TP8官方指南把 TP 的适用场景定为 hidden size 4096 以上。hidden 很小的模型开大 TP通信占比会盖过计算收益反而变慢。第三步确定 PPTP 切完仍放不下再加 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_node8pretrain_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这里 TP4、PP4、CP2DP 就是 64 / (4 x 4 x 2) 2。CP 是上下文并行只在序列很长官方说明针对 8K 以上这个示例的序列长度正好是 8192时才需要序列不长就设为 1把卡留给 DP。第六步看指标再迭代先别猜看三个数单卡 MFU、显存峰值、各 rank 的通信耗时。显存有富余就把 TP 或 PP 降一档把卡让给 DPMFU 低且 PP 气泡明显就开交错调度或增大微批次数通信耗时集中在 TP all-reduce就确认 TP 组没有跨出节点。官方给的总原则很朴素先只用 DP放不下再加 TP模型更大再加 PP序列很长再加 CP能不加的维度就不加。论文中的万亿参数模型在 3072 张 GPU 上跑出 502 petaFLOP/s单卡吞吐达到理论峰值的 52%说明这套组合在超大规模下也能把算力吃满前提是每一维都配得克制。我的看法3D 并行的调优没有万能配方但有稳定的顺序先让模型放得下再让通信不跨越慢链路最后用 DP 吃掉剩余的卡。多数人踩坑是一上来就把 TP 开得很大或者为了省显存无脑加 PP。反过来想每一维并行都是一笔通信账能少付就少付。
返回列表