Megatron-LM 大模型训练实践
Megatron-LM 是 NVIDIA 开发的大模型训练框架，支持 3D 并行（TP+PP+DP），是训练千亿参数模型的核心工具。
3D 并行架构
Megatron-LM 将三种并行方式组合：
TP (张量并行) - 节点内，利用 NVLink 高带宽
PP (流水线并行) - 跨节点，点对点通信
DP (数据并行) - 跨节点，梯度 AllReduce总 GPU 数 = TP × PP × DP
启动训练
bash
# 训练 GPT-3 175B 模型
# TP=8, PP=16, DP=64, 总共 1024 GPU
python pretrain_gpt.py \
--num-layers 96 \
--hidden-size 12288 \
--num-attention-heads 96 \
--seq-length 2048 \
--micro-batch-size 1 \
--global-batch-size 1536 \
--tensor-model-parallel-size 8 \
--pipeline-model-parallel-size 16 \
--data-parallel-size 64 \
--optimizer adam \
--adam-beta1 0.9 \
--adam-beta2 0.95 \
--lr 6e-5 \
--min-lr 6e-6 \
--weight-decay 0.1 \
--clip-grad 1.0 \
--bf16 \
--use-flash-attn关键配置
并行度选择
python
# 并行度选择指南
def choose_parallelism(num_gpus, model_params, nvlink_bandwidth):
"""选择最优并行度"""
# TP: 最大化 NVLink 利用
tp = min(8, num_gpus) # 单节点最多 8 GPU
# PP: 根据模型深度
num_layers = estimate_layers(model_params)
pp = min(num_gpus // tp, num_layers // 4)
# DP: 剩余 GPU
dp = num_gpus // (tp * pp)
return tp, pp, dp| 模型规模 | GPU 数 | TP | PP | DP | 微批次 |
|---|---|---|---|---|---|
| 7B | 8 | 4 | 1 | 2 | 4 |
| 13B | 16 | 4 | 2 | 2 | 4 |
| 70B | 64 | 8 | 4 | 2 | 8 |
| 175B | 1024 | 8 | 16 | 8 | 16 |
混合精度训练
bash
# BF16 混合精度（推荐）
--bf16
# FP16 混合精度（需 Loss Scaling）
--fp16 --initial-loss-scale 4294967296Flash Attention
Megatron-LM 原生支持 Flash Attention，可减少 Attention 计算的显存占用和延迟。始终启用 --use-flash-attn。
训练监控
bash
# TensorBoard 监控
tensorboard --logdir /path/to/checkpoints
# 关键指标
# - lm loss: 应平稳下降
# - learning rate: 预热后余弦衰减
# - grad norm: 应稳定，spike 后需检查
# - throughput: tokens/s 或 samples/s检查点保存
3D 并行的检查点保存需要所有 Rank 协调。Megatron-LM 使用分布式检查点格式，每个 Rank 保存本地分片。恢复时需确保并行度不变。