Megatron-LM 大模型训练 — 英伟达出品的千亿级模型训练框架¶

# Megatron-LM 需要从源码安装（暂无 pip 包）
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM

# 安装依赖
pip install -r requirements.txt

# 安装 Apex（英伟达的混合精度加速库，强烈推荐）
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check \
    --no-cache-dir --no-build-isolation \
    --config-settings "--build-option=--cpp_ext" \
    --config-settings "--build-option=--cuda_ext" ./

# 安装 Flash Attention（可选，但推荐）
pip install flash-attn --no-build-isolation

# 验证（需要在 GPU 环境下）
python -c "import megatron; print('Megatron-LM ready')"

# 启动 GPT 预训练（8 GPU，张量并行=2，流水线并行=2，数据并行=2）
torchrun --nproc_per_node=8 pretrain_gpt.py \
    --num-layers 32 \                              # Transformer 层数
    --hidden-size 4096 \                           # 隐藏层维度
    --num-attention-heads 32 \                     # 注意力头数
    --seq-length 2048 \                            # 序列长度
    --max-position-embeddings 2048 \
    --micro-batch-size 4 \                         # 每个微批次大小
    --global-batch-size 512 \                      # 全局 batch 大小
    --train-iters 500000 \                         # 训练迭代次数
    --data-path /data/gpt_text_document \          # 预处理后的数据路径
    --vocab-file gpt2-vocab.json \
    --merge-file gpt2-merges.txt \
    --tensor-model-parallel-size 2 \              # 张量并行数（切分同一层）
    --pipeline-model-parallel-size 2 \            # 流水线并行数（切分不同层）
    --lr 0.00015 \
    --bf16

并行策略说明（3D 并行）：

张量并行（Tensor Parallelism，--tensor-model-parallel-size N）：
  - 把同一个矩阵乘法切分到 N 张 GPU 上同时计算
  - 适合单层参数很大的情况（如 FFN 层）
  - N 张 GPU 之间需要频繁通信（需要高速 NVLink）

流水线并行（Pipeline Parallelism，--pipeline-model-parallel-size N）：
  - 把不同的 Transformer 层分配到不同 GPU 上
  - 类似工厂流水线：GPU1 处理完第 1-8 层传给 GPU2 处理 9-16 层
  - 通信少，但 GPU 会有"气泡"（等待时间）

数据并行（Data Parallelism）：
  - 每个 GPU 处理不同的数据，梯度汇总取平均
  - 总 GPU 数 = 张量并行 × 流水线并行 × 数据并行

# 数据预处理（训练前必须先处理数据）
python tools/preprocess_data.py \
    --input /data/raw_text.json \                  # 原始文本（JSONL 格式）
    --output-prefix /data/gpt_text \               # 输出文件前缀
    --vocab-file gpt2-vocab.json \
    --dataset-impl mmap \                          # 内存映射格式（大文件友好）
    --tokenizer-type GPT2BPETokenizer \
    --workers 4 \                                  # 并行处理线程数
    --append-eod                                   # 在每个文档末尾加 EOD token

# 在 4 张 A100 上训练 13B 参数 LLaMA 风格模型

# 参数规划：
# LLaMA-13B：num-layers=40, hidden-size=5120, num-attention-heads=40
# 并行策略：张量并行=2, 流水线并行=2, 数据并行=1
# 总 GPU = 2 × 2 × 1 = 4 张

torchrun --nproc_per_node=4 pretrain_gpt.py \
    --num-layers 40 \
    --hidden-size 5120 \
    --num-attention-heads 40 \
    --seq-length 4096 \
    --max-position-embeddings 4096 \
    --micro-batch-size 1 \
    --global-batch-size 32 \
    --train-iters 100000 \
    --tensor-model-parallel-size 2 \
    --pipeline-model-parallel-size 2 \
    --lr 1e-4 \
    --min-lr 1e-5 \
    --lr-decay-style cosine \
    --weight-decay 0.1 \
    --clip-grad 1.0 \
    --bf16 \
    --use-flash-attn \                             # 使用 Flash Attention 加速
    --use-rotary-position-embeddings \             # 使用 RoPE 位置编码（LLaMA 同款）
    --no-position-embedding \
    --data-path /data/train_text_document \
    --save /checkpoints/llama-13b \
    --save-interval 1000 \                         # 每 1000 步保存一次
    --log-interval 100