# Megatron-LM 3D 混合并行：张量与流水线并行工程化训练万亿参数 Transformer > 面向数千 GPU 训练万亿参数模型，给出 Megatron-LM 中 TP、PP 与 3D 混合并行的配置参数、拓扑策略与性能优化要点。 ## 元数据 - 路径: /posts/2026/02/26/engineering-tensor-pipeline-parallelism-megatron-lm-3d-hybrid/ - 发布时间: 2026-02-26T15:46:21+08:00 - 分类: [mlops](/categories/mlops/) - 站点: https://blog.hotdry.top ## 正文 在大型语言模型（LLM）训练进入万亿参数时代，单一GPU内存已无法承载模型权重与激活值，分布式模型并行成为必然。NVIDIA的Megatron-LM框架通过张量并行（Tensor Parallelism, TP）、流水线并行（Pipeline Parallelism, PP）与数据并行（Data Parallelism, DP）的3D混合策略，实现了在数千H100 GPU上高效训练462B参数模型，MFU高达47%。本文聚焦工程化实现，剖析核心机制、配置参数与落地清单，帮助ML工程师快速上手大规模Transformer训练。 ### 张量并行（TP）：处理宽层模型的核心 TP将Transformer单层的权重矩阵（如注意力QKV投影、MLP前馈）拆分到多个GPU上，每个GPU仅计算局部shard，避免单卡内存爆炸。Megatron-LM采用column-parallel和row-parallel线性层实现： - **Column-parallel Linear**：输入X（B, H）乘权重W（H, 4H）时，将W按列拆分到TP_size个GPU，每个GPU计算Y_i = X * W_i（输出B, 4H/TP_size），前向仅需all-gather合并。 - **Row-parallel Linear**：后续投影Z = Y * V（V按行拆分），每个GPU计算局部Z_i，后向需all-reduce求和。 注意力机制中，Q/K/V投影column-parallel（头数按TP拆分），输出投影row-parallel。这种设计确保每个Transformer block仅2次集体通信（all-gather + all-reduce），极大降低开销。 **工程参数建议**： - TP_size：优先2/4/8，必须整除注意力头数与MLP hidden_size。宽模型（如hidden_size>8192）用8；窄模型用4。 - 拓扑映射：TP组限于单节点内NVLink/NVSwitch，避免跨节点all-reduce延迟。 - 内存节省：TP减小每卡参数量约1/TP_size，但激活需通信合并，结合FP8/BF16进一步压缩。 证据显示，在6144 H100上，TP=8的462B模型训练稳定MFU 47%，“Our codebase efficiently trains models from 2B to 462B parameters across thousands of GPUs, achieving up to 47% Model FLOP Utilization (MFU) on H100 clusters.”（来源：Megatron-LM GitHub）。 ### 流水线并行（PP）：征服深层模型的利器 PP将模型层栈分为P个连续stage，每个stage（含多层）置于TP组上，microbatch沿stage流水前进/后向重叠（1F1B调度），隐藏bubble。 - **Microbatch调度**：全局batch拆为M个microbatch，M >> P以填充pipeline。Bubble比例≈(P-1)/M。 - **激活传递**：仅stage边界激活send/recv，点对点通信容忍跨节点延迟。 **工程参数建议**： - PP_size：层数/每stage层数，目标每stage FLOPs均衡（注意力层≈FFN层1.3倍FLOPs）。深模型（>100层）用16-32。 - Microbatch_size：每GPU 1-4 seq（视seq_len=4096），global_batch = microbatch * M * DP_size * TP_size * PP_size。 - 内存优化：激活checkpointing（--checkpoint-activations），每stage内存减半，但重算代价≈20%额外FLOPs。 PP跨节点映射优选环形拓扑（ring all-reduce备选），与TP互补。 ### 3D混合并行：万亿参数的黄金配置 3D策略下，总GPU数 = TP_size × PP_size × DP_size。设计原则：TP控宽（intra-node）、PP控深（inter-node）、DP提吞吐（跨rack，ZeRO分片）。 **示例配置（1024 H100，~1T参数）**： | 维度 | 大小 | 作用 | 通信范围 | |------|------|------|----------| | TP | 8 | 层内shard | 单节点 NVLink | | PP | 16 | 层间流水 | 跨节点 InfiniBand | | DP | 8 | 数据复制 | ZeRO-3 全分片 | 启动命令模板（Megatron-LM examples/pretrain_gpt.sh）： ``` mpirun -np 1024 \ python pretrain_gpt.py \ --num-layers 128 --hidden-size 16384 --num-attention-heads 128 \ --tensor-model-parallel-size 8 \ --pipeline-model-parallel-size 16 \ --data-parallel-size 8 \ --micro-batch-size 1 --global-batch-size 1024 \ --seq-len 4096 \ --fp16 --activation-checkpointing \ --zero-stage 3 # ZeRO via DeepSpeed集成 ``` **性能调优清单**： 1. **MFU监控**：目标>45%。低则检查：matmul太小（TP过大）、bubble大（M