training-llms-megatron

Megatron-Core - 大规模 LLM 训练

快速开始

Megatron-Core 通过先进的并行策略，在 H100 GPU 上以最高 47% 的模型 FLOP 利用率训练参数量为 2B 至 462B 的 LLM。

安装：

# Docker (recommended)
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:25.04-py3

# Or pip
pip install megatron-core

简单的分布式训练：

# Train with 2 GPUs using data parallelism
torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py

# Or LLaMA-3 8B training
./examples/llama/train_llama3_8b_fp8.sh

常见工作流

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

复制此检查清单：

LLaMA Training Setup:
- [ ] Step 1: Choose parallelism configuration
- [ ] Step 2: Configure training hyperparameters
- [ ] Step 3: Launch distributed training
- [ ] Step 4: Monitor performance metrics

步骤 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

GPUS_PER_NODE=8
NNODES=8  # 64 GPUs total
TP=4      # Tensor parallel
PP=4      # Pipeline parallel
CP=1      # Context parallel

# LLaMA 70B configuration
MODEL_SIZE=70  # Billion parameters
HIDDEN_SIZE=8192
NUM_LAYERS=80
NUM_HEADS=64
SEQ_LENGTH=4096

# Training hyperparameters
MICRO_BATCH=1
GLOBAL_BATCH=1024
LR=3e-4

torchrun \
  --nproc_per_node=$GPUS_PER_NODE \
  --nnodes=$NNODES \
  pretrain_gpt.py \
  --tensor-model-parallel-size $TP \
  --pipeline-model-parallel-size $PP \
  --context-parallel-size $CP \
  --sequence-parallel \
  --num-layers $NUM_LAYERS \
  --hidden-size $HIDDEN_SIZE \
  --num-attention-heads $NUM_HEADS \
  --seq-length $SEQ_LENGTH \
  --max-position-embeddings $SEQ_LENGTH \
  --micro-batch-size $MICRO_BATCH \
  --global-batch-size $GLOBAL_BATCH \
  --lr $LR \
  --train-iters 100000 \
  --lr-decay-style cosine \
  --lr-warmup-iters 2000 \
  --weight-decay 0.1 \
  --clip-grad 1.0 \
  --bf16 \
  --use-mcore-models \
  --transformer-impl transformer_engine \
  --data-path /path/to/data \
  --vocab-file /path/to/vocab.json \
  --merge-file /path/to/merges.txt

步骤 3：启动分布式训练

# Single node (8 GPUs)
bash train_llama_70b.sh

# Multi-node with SLURM
sbatch --nodes=8 --gpus-per-node=8 train_llama_70b.sh

步骤 4：监控性能指标

需要跟踪的关键指标：

Model FLOP Utilization (MFU): Target >40% on H100
Throughput: Tokens/sec/GPU
Memory usage: <80GB per GPU for 70B model
Loss: Should decrease steadily

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

适用于 Mixtral 等稀疏 MoE 模型。

MoE Training:
- [ ] Step 1: Configure expert parallelism
- [ ] Step 2: Set MoE hyperparameters
- [ ] Step 3: Launch training with EP

步骤 1：配置专家并行

# Mixtral 8x7B example
TENSOR_PARALLEL=2
PIPELINE_PARALLEL=1
EXPERT_PARALLEL=4  # Split 8 experts across 4 GPUs
DATA_PARALLEL=4

TOTAL_GPUS=$((TENSOR_PARALLEL * PIPELINE_PARALLEL * EXPERT_PARALLEL * DATA_PARALLEL))
# = 2 * 1 * 4 * 4 = 32 GPUs

步骤 2：设置 MoE 超参数

torchrun \
  --nproc_per_node=8 \
  pretrain_gpt.py \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 1 \
  --expert-model-parallel-size 4 \
  --num-experts 8 \
  --moe-router-topk 2 \
  --moe-router-load-balancing-type aux_loss \
  --moe-aux-loss-coeff 0.01 \
  --hidden-size 4096 \
  --num-layers 32 \
  --num-attention-heads 32 \
  --seq-length 4096 \
  --max-position-embeddings 4096 \
  --bf16 \
  --use-mcore-models \
  --transformer-impl transformer_engine \
  --data-path /path/to/data \
  --vocab-file /path/to/vocab.json \
  --merge-file /path/to/merges.txt

步骤 3：使用 EP 启动训练

专家并行会将不同专家分布到各个 GPU 上，在保持模型容量的同时减少内存占用。

Memory without EP: 8 experts × 7B = 56GB per GPU
Memory with EP=4: 2 experts × 7B = 14GB per GPU
Savings: 75% memory reduction

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

在 H100 上达到 47% MFU。

Performance Optimization:
- [ ] Step 1: Enable Flash Attention
- [ ] Step 2: Use FP8 precision (H100)
- [ ] Step 3: Optimize micro-batch size
- [ ] Step 4: Tune parallelism degrees

步骤 1：启用优化

--use-mcore-models  # Use Megatron Core models
--transformer-impl transformer_engine  # Use Transformer Engine
--sequence-parallel  # Reduce activation memory (use with TP)

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

--fp8-hybrid  # FP8 mixed precision training
# Transformer Engine handles FP8 automatically

结果：在 H100 上，相较于 BF16 可提速 1.5-2 倍。

步骤 3：优化微批次大小

找到适合内存的最大微批次：

# Start with 1, increase until OOM
for MBS in 1 2 4 8; do
  echo "Testing micro-batch-size=$MBS"
  torchrun ... --micro-batch-size $MBS
done

典型值：

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

步骤 4：调整并行度

经验法则：

Tensor Parallel: Use ≤8 (limited by NVLink within node)
Pipeline Parallel: Use for >70B models
Context Parallel: Use for sequences >8K tokens
Data Parallel: Fill remaining GPUs

128 张 H100 上运行 405B 模型的示例：

TP=8 (1 node)
PP=8 (across nodes)
CP=2 (long sequences)
DP=1
Total = 8 × 8 × 2 × 1 = 128 GPUs

何时使用及替代方案

适合使用 Megatron-Core 的情况：

• 训练参数量 >10B 的模型
• 需要最高效率（目标 >40% MFU）
• 使用 NVIDIA GPU（A100、H100）
• 大规模生产训练
• 希望精细控制并行策略

适合使用的替代方案：

• PyTorch FSDP：模型 <70B、API 更简单、PyTorch 原生支持
• DeepSpeed：设置更简单，适用于 <100B 的模型
• HuggingFace Accelerate：适合原型开发，工作流更简单
• LitGPT：适合教学，采用单文件实现

常见问题

问题：GPU 利用率低（<30% MFU）

原因：

1. 微批次过小
2. 并行开销过大
3. 未使用 Flash Attention

解决方法：

# Increase micro-batch
--micro-batch-size 4  # Was 1

# Enable optimizations
--use-flash-attn
--sequence-parallel

# Reduce TP if >8
--tensor-model-parallel-size 4  # Was 16

问题：内存不足

通过以下方式减少内存占用：

--tensor-model-parallel-size 2  # Split model across GPUs
--recompute-granularity full  # Gradient checkpointing
--recompute-method block  # Checkpoint transformer blocks
--recompute-num-layers 1  # Checkpoint every layer

或者使用 CPU/NVMe 卸载：

--cpu-optimizer  # Offload optimizer to CPU
--cpu-optimizer-type ADAM  # CPU Adam variant

问题：训练速度低于预期

请检查：

1. 网络瓶颈：确保已启用 InfiniBand/NVLink
2. 流水线气泡：使用交错流水线调度

   --num-layers-per-virtual-pipeline-stage 2
   

3. 数据加载：使用快速数据加载器

   --dataloader-type cyclic
   

问题：损失发散

稳定训练：

--lr-warmup-iters 2000  # Longer warmup
--clip-grad 1.0  # Gradient clipping
--init-method-std 0.006  # Smaller init
--attention-dropout 0.0  # No dropout in attention
--hidden-dropout 0.0  # No dropout in FFN

进阶主题

并行策略：参阅 references/parallelism-guide.md，了解 TP/PP/DP/CP/EP 的详细对比、性能分析以及各自的适用场景。

性能基准：参阅 references/benchmarks.md，了解不同模型大小和 GPU 配置下的 MFU 数据。

生产配置：参阅 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（TP=4）
  • 405B 模型：80GB（TP=8，PP=8）
• 存储：检查点需要高速 NVMe（70B+ 模型需 1TB+）

资源

• 文档：https://docs.nvidia.com/megatron-core/
• GitHub：https://github.com/NVIDIA/Megatron-LM
• 论文：
  • “Megatron-LM：使用模型并行训练数十亿参数语言模型”（2019）
  • “使用 Megatron-LM 在 GPU 集群上高效训练大规模语言模型”（2021）
• NeMo Framework：https://docs.nvidia.com/nemo-framework/（基于 Megatron-Core 构建）

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存、驱动程序/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 确定输入。 根据数据或用户确认格式、标识符和参数。任何会改变结果的值都应询问，不要猜测。
3. 先运行小规模版本。 执行冒烟测试：使用小型模型或子集，运行少量步骤，并确认损失下降且检查点已保存。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件。
6. 汇报。 说明运行了什么（版本、命令、参数）、检查了什么，以及仍有哪些不确定之处。

如果发生这种情况	请执行此操作
CUDA 显存不足	减小批次大小、启用梯度累积/检查点或混合精度，或者对模型进行分片。
损失为 NaN 或发散	降低学习率、检查数据中是否有无效值，并启用梯度裁剪。
这些说明中的函数、标志或端点在已安装版本中不存在	查阅已安装版本自己的文档（help()、--help、官方文档），进行调整，并告知用户。切勿编造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明假设。

诚信规则

• 切勿捏造结果、参数、标识符、引用或统计数据。如果无法运行或验证某项内容，请直言说明。
• 没有日志时，切勿声称取得了训练结果；启动大型任务前，先估算计算成本。
• 此处的版本特定细节可能已过时：请对照已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，先征得同意。

相关技能

• tensorrt-llm：使用 NVIDIA TensorRT 优化 LLM 推理，以实现最高吞吐量和最低延迟。
• moe-training：使用 DeepSpeed 或 HuggingFace 训练混合专家（MoE）模型。
• axolotl：提供使用 Axolotl 微调大语言模型的指南，涵盖 YAML 训练配置、LoRA 和 QLoRA、偏好训练等内容……