megatron-lm

Megatron-LM 大规模训练

补充说明

Megatron-LM 是 NVIDIA 的大模型训练框架，支持张量并行、流水线并行、数据并行和序列并行。适合千亿参数级别、多机多卡训练。

核心概念：TP（Tensor Parallel）、PP（Pipeline Parallel）、DP（Data Parallel）、micro-batch、global batch 和 checkpoint 转换。

常用示例

准备数据：

python tools/preprocess_data.py \
  --input corpus.jsonl \
  --output-prefix data/corpus \
  --tokenizer-type HuggingFaceTokenizer \
  --tokenizer-model Qwen/Qwen2.5-7B-Instruct

启动训练：

torchrun --nproc_per_node=8 --nnodes=2 \
  --node_rank=0 --master_addr=10.0.0.10 --master_port=29500 \
  pretrain_gpt.py \
  --tensor-model-parallel-size 4 \
  --pipeline-model-parallel-size 2 \
  --num-layers 32 \
  --hidden-size 4096 \
  --num-attention-heads 32 \
  --seq-length 4096 \
  --micro-batch-size 1 \
  --global-batch-size 64

转换 Hugging Face checkpoint：

python tools/checkpoint/convert.py \
  --model-type GPT \
  --loader huggingface \
  --saver megatron \
  --load-dir ./hf-model \
  --save-dir ./megatron-checkpoint

易错点与排错

• TP/PP 的乘积必须能整除总 GPU 数，否则启动失败。
• global batch = micro-batch × DP × gradient accumulation，计算错误会导致训练不稳定。
• 序列长度和位置编码配置不匹配，长文本训练异常。
• 多机通信没有走 IB/NVLink，扩展效率低。
• checkpoint 格式与 Megatron 版本不兼容，转换失败。
• 数据预处理的分词器与模型不一致，token 分布错误。
• 流水线并行下 micro-batch 太少，气泡严重。
• 没有配置激活重计算，显存不足。