Hotdry.
归档
mlops

Megatron-LM 3D 混合并行：张量与流水线并行工程化训练万亿参数 Transformer

面向数千 GPU 训练万亿参数模型，给出 Megatron-LM 中 TP、PP 与 3D 混合并行的配置参数、拓扑策略与性能优化要点。

在大型语言模型（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<P*2）、net 饱和（nvidia-smi -q -d UTILIZATION）。
  2. 内存阈值：<85% GPU mem。OOM 时：减 microbatch、增 checkpoint、用 FP8（--fp8）。
  3. 通信重叠：--overlap-grad-reduce --tp-comm-overlap，默认 PP 重叠。
  4. 容错：--checkpoint-interval 1000 steps，分布式 checkpoint reshaping 支持 TP/PP 变更。
  5. 扩展性：弱 scaling（增模型 / GPU）MFU 升至 48%；强 scaling（定 batch 增 GPU）MFU 降至 42%，DP 主导。

常见陷阱与回滚：

  • TP 不整除维度→padding 浪费，fallback TP=4。
  • PP 不均衡→慢 stage 瓶颈，用 profiler（nsys）定位，重分 stage。
  • 大 DP 下 ZeRO 通信峰值高，监控 p2p throughput，回滚小 DP + 大 global-batch accumulation。

MoE 扩展：加 EP_size，专家路由负载均衡 < 1.2 倍方差。

通过以上参数与策略，Megatron-LM 将万亿训练从概念落地生产。实际部署前，基准小规模（TP=1,PP=1）验证 config，再线性扩展。

资料来源：

（正文约 1250 字）

查看归档