名称: 训练大语言模型-Megatron 描述: 使用 NVIDIA Megatron-Core 训练大语言模型（2B-462B 参数），采用先进的并行策略。当训练模型参数超过10亿、需要最大 GPU 效率（在 H100 上达到 47% 模型浮点运算利用率）或需要张量/管道/序列/上下文/专家并行时使用。生产就绪的框架，用于 Nemotron、LLaMA、DeepSeek。版本: 1.0.0 作者: Orchestra Research 许可证: MIT 标签: [Megatron-Core, 大规模训练, NVIDIA, 张量并行, 管道并行, 模型并行, H100, 分布式训练, 生产] 依赖项: [megatron-core, torch, apex, transformer-engine]

Megatron-Core - 大规模 LLM 训练

快速开始

Megatron-Core 通过先进的并行策略，在 H100 GPU 上训练 2B 到 462B 参数的大语言模型，模型浮点运算利用率高达 47%。

安装:

# Docker（推荐）
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:25.04-py3

# 或使用 pip
pip install megatron-core

简单分布式训练:

# 使用数据并行在 2 个 GPU 上训练
torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py

# 或训练 LLaMA-3 8B
./examples/llama/train_llama3_8b_fp8.sh

常见工作流

工作流 1：使用 3D 并行训练 LLaMA 风格模型

复制此清单：

LLaMA 训练设置：
- [ ] 步骤 1：选择并行配置
- [ ] 步骤 2：配置训练超参数
- [ ] 步骤 3：启动分布式训练
- [ ] 步骤 4：监控性能指标

步骤 1：选择并行配置

模型大小决定并行策略：

模型大小	GPU 数量	张量并行	管道并行	数据并行	上下文并行
7B	8	1	1	8	1
13B	8	2	1	4	1
70B	64	4	4	4	1
405B	128	8	8	2	2

步骤 2：配置训练超参数

#!/bin/bash
# train_llama_70b.sh

每节点 GPU 数量=8
节点数量=8  # 总共 64 个 GPU
张量并行=4      # 张量并行
管道并行=4      # 管道并行
上下文并行=1      # 上下文并行

# LLaMA 70B 配置
模型大小=70  # 十亿参数
隐藏大小=8192
层数=80
注意力头数=64
序列长度=4096

# 训练超参数
微批大小=1
全局批大小=1024
学习率=3e-4

torchrun \
  --每节点进程数=$每节点 GPU 数量 \
  --节点数量=$节点数量 \
  pretrain_gpt.py \
  --张量模型并行大小 $张量并行 \
  --管道模型并行大小 $管道并行 \
  --上下文并行大小 $上下文并行 \
  --序列并行 \
  --层数 $层数 \
  --隐藏大小 $隐藏大小 \
  --注意力头数 $注意力头数 \
  --序列长度 $序列长度 \
  --最大位置嵌入 $序列长度 \
  --微批大小 $微批大小 \
  --全局批大小 $全局批大小 \
  --学习率 $学习率 \
  --训练迭代次数 100000 \
  --学习率衰减风格 余弦 \
  --学习率预热迭代次数 2000 \
  --权重衰减 0.1 \
  --梯度裁剪 1.0 \
  --bf16 \
  --使用核心模型 \
  --transformer 实现 transformer_engine \
  --数据路径 /路径/到/数据 \
  --词汇文件 /路径/到/vocab.json \
  --合并文件 /路径/到/merges.txt

步骤 3：启动分布式训练

# 单节点（8 个 GPU）
bash train_llama_70b.sh

# 多节点使用 SLURM
sbatch --节点数=8 --每节点 GPU 数=8 train_llama_70b.sh

步骤 4：监控性能指标

关键指标跟踪：

模型浮点运算利用率：目标 >40% 在 H100 上
吞吐量：令牌/秒/GPU
内存使用：<80GB 每个 GPU 用于 70B 模型
损失：应稳步下降

工作流 2：配置混合专家（MoE）训练

用于稀疏 MoE 模型如 Mixtral。

MoE 训练：
- [ ] 步骤 1：配置专家并行
- [ ] 步骤 2：设置 MoE 超参数
- [ ] 步骤 3：使用专家并行启动训练

步骤 1：配置专家并行

# Mixtral 8x7B 示例
张量并行=2
管道并行=1
专家并行=4  # 在 4 个 GPU 上分配 8 个专家
数据并行=4

总 GPU 数=$((张量并行 * 管道并行 * 专家并行 * 数据并行))
# = 2 * 1 * 4 * 4 = 32 个 GPU

步骤 2：设置 MoE 超参数

torchrun \
  --每节点进程数=8 \
  pretrain_gpt.py \
  --张量模型并行大小 2 \
  --管道模型并行大小 1 \
  --专家模型并行大小 4 \
  --专家数量 8 \
  --moe 路由器 topk 2 \
  --moe 路由器负载均衡类型 辅助损失 \
  --moe 辅助损失系数 0.01 \
  --隐藏大小 4096 \
  --层数 32 \
  --注意力头数 32 \
  --序列长度 4096 \
  --最大位置嵌入 4096 \
  --bf16 \
  --使用核心模型 \
  --transformer 实现 transformer_engine \
  --数据路径 /路径/到/数据 \
  --词汇文件 /路径/到/vocab.json \
  --合并文件 /路径/到/merges.txt

步骤 3：使用专家并行启动训练

专家并行在不同 GPU 上分配不同专家，减少内存同时保持容量。

无专家并行内存：8 专家 × 7B = 56GB 每个 GPU
有专家并行=4：2 专家 × 7B = 14GB 每个 GPU
节省：75% 内存减少

工作流 3：优化最大吞吐量

在 H100 上实现 47% 模型浮点运算利用率。

性能优化：
- [ ] 步骤 1：启用 Flash Attention
- [ ] 步骤 2：使用 FP8 精度（H100）
- [ ] 步骤 3：优化微批大小
- [ ] 步骤 4：调优并行度

步骤 1：启用优化

--使用核心模型  # 使用 Megatron Core 模型
--transformer 实现 transformer_engine  # 使用 Transformer Engine
--序列并行  # 减少激活内存（与张量并行使用）

步骤 2：使用 FP8 精度（仅 H100）

--fp8 混合  # FP8 混合精度训练
# Transformer Engine 自动处理 FP8

结果：在 H100 上比 BF16 提速 1.5-2 倍。

步骤 3：优化微批大小

找到能适应内存的最大微批大小：

# 从 1 开始，增加直到内存不足
for 微批大小 in 1 2 4 8; do
  echo "测试微批大小=$微批大小"
  torchrun ... --微批大小 $微批大小
done

典型值：

7B 模型：4-8
70B 模型：1-2
405B 模型：1

步骤 4：调优并行度

经验法则：

张量并行：使用 ≤8（受节点内 NVLink 限制）
管道并行：用于 >70B 模型
上下文并行：用于序列 >8K 令牌
数据并行：填充剩余 GPU

示例 405B 在 128 个 H100 上：

张量并行=8（1 个节点）
管道并行=8（跨节点）
上下文并行=2（长序列）
数据并行=1
总计 = 8 × 8 × 2 × 1 = 128 个 GPU

何时使用 vs 替代方案

使用 Megatron-Core 当：

训练模型 >10B 参数
需要最大效率（目标 >40% 模型浮点运算利用率）
使用 NVIDIA GPU（A100, H100）
生产规模训练
想要细粒度并行控制

使用替代方案代替：

PyTorch FSDP：模型 <70B，更简单 API，PyTorch 原生
DeepSpeed：更容易设置，适合 <100B 模型
HuggingFace Accelerate：原型设计，更简单工作流
LitGPT：教育用途，单文件实现

常见问题

问题：低 GPU 利用率（<30% 模型浮点运算利用率）

原因：

微批大小太小
并行开销过多
未使用 Flash Attention

修复：

# 增加微批大小
--微批大小 4  # 原为 1

# 启用优化
--使用闪存注意力
--序列并行

# 减少张量并行如果 >8
--张量模型并行大小 4  # 原为 16

问题：内存不足

减少内存：

--张量模型并行大小 2  # 在 GPU 间分割模型
--重新计算粒度 完全  # 梯度检查点
--重新计算方法 块  # 检查点 transformer 块
--重新计算层数 1  # 每层检查点

或使用 CPU/NVMe 卸载：

--CPU 优化器  # 将优化器卸载到 CPU
--CPU 优化器类型 ADAM  # CPU Adam 变体

问题：训练速度低于预期

检查：

网络瓶颈：确保启用 InfiniBand/NVLink
管道气泡：使用交错管道调度
```
--每个虚拟管道阶段层数 2
```
数据加载：使用快速数据加载器
```
--数据加载器类型 循环
```

问题：损失发散

稳定训练：

--学习率预热迭代次数 2000  # 更长预热
--梯度裁剪 1.0  # 梯度裁剪
--初始化方法标准差 0.006  # 更小初始化
--注意力丢弃 0.0  # 注意力中无丢弃
--隐藏丢弃 0.0  # FFN 中无丢弃

高级主题

并行策略：参见 references/parallelism-guide.md 详细比较张量并行/管道并行/数据并行/上下文并行/专家并行，附性能分析和使用时机。

性能基准：参见 references/benchmarks.md 不同模型大小和 GPU 配置的模型浮点运算利用率数据。

生产配置：参见 references/production-examples.md 来自 LLaMA 3 405B、Nemotron-4 340B 和 DeepSeek-V3 671B 的真实世界设置。

训练配方：参见 references/training-recipes.md 完整 GPT/LLaMA/Mixtral 架构的超参数配置。

硬件要求

GPU：NVIDIA Ampere+（A100, H100, B200）
- Turing 可用但较慢
- FP8 需要 Hopper/Ada/Blackwell
网络：InfiniBand 或 400Gb+ 以太网用于多节点
每个 GPU 内存：
- 7B 模型：40GB+
- 70B 模型：80GB（带张量并行=4）
- 405B 模型：80GB（带张量并行=8, 管道并行=8）
存储：快速 NVMe 用于检查点（70B+ 模型需要 1TB+）

资源

文档：https://docs.nvidia.com/megatron-core/
GitHub：https://github.com/NVIDIA/Megatron-LM
论文：
- “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism” (2019)
- “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM” (2021)
NeMo 框架：https://docs.nvidia.com/nemo-framework/（基于 Megatron-Core 构建）