总目录 › Megatron-LM › 第 3 章
Chapter 03

快速上手：跑通最小可见的训练循环

📌 本章命令均基于 NVIDIA/Megatron-LM · commit 432d76b2b8bf 。如果你的本地 repo 已经更新，部分函数签名可能已经动过。

3.1 两条路径的选择

Megatron 入门有两条路：

路径文件学到什么跑得动吗
A. 最小可见循环 examples/run_simple_mcore_train_loop.py 用 Megatron-Core 的最小积木从零搭一个 2 层 GPT，看清初始化 → 数据 → forward → backward → checkpoint 的全链路 单卡 1 GB 显存都能跑
B. 工业级预训练入口 pretrain_gpt.py + gpt_builders.py 看 NVIDIA 自己用来训 175B 的入口长什么样，调用栈深 至少 8 × A100 才有意义

本章走 A 路径，因为它全部代码都在一个文件里，不到 300 行，看完就懂"Megatron 在搞什么"。看完 A 再去翻 B 会顺很多。

3.2 把这个最小循环读一遍

examples/run_simple_mcore_train_loop.py 在你本地 repos/megatron-lm/examples/run_simple_mcore_train_loop.py。它的核心 5 个函数：

initialize_distributed()        # 设置 NCCL + Megatron parallel groups
model_provider()                # 造一个 2 层、hidden=12 的 GPTModel
get_train_data_iterator()       # 用 MockGPTDataset 造随机数据
forward_step_func()             # 跑一次 forward 并返回 loss
save_distributed_checkpoint()   # （可选）保存

把它和你熟悉的"普通 PyTorch 训练"对照一下：

普通 PyTorchMegatron 这里差别在哪
torch.distributed.init_process_group() initialize_distributed() 内部还会调 parallel_state.initialize_model_parallel() 多了"模型并行 group"的概念
model = MyTransformer() model = GPTModel(config=TransformerConfig(...), ...) 用 TransformerConfig 一份配置驱动模型构建
自己写 for x in loader: loss = model(x).backward() forward_backward_func(...) 一次包办 forward/backward 被封装因为要支持 pipeline 调度

这就是 Megatron 的"长相"：用一组高度抽象的工具函数代替你自己写训练循环。

3.3 真的跑一次

# 进入克隆好的源码目录
cd /Users/zjw/Documents/LLM/train/repos/megatron-lm

# 单卡跑（torchrun 即使单卡也用，方便和多卡共享代码）
torchrun --nproc_per_node 1 examples/run_simple_mcore_train_loop.py

你会看到类似这样的输出（数字可能略有不同）：

[Megatron] step 1 ... lm loss = 4.6051
[Megatron] step 2 ... lm loss = 4.5912
...

跑得通说明你已经完成了一次完整的 Megatron-Core 训练循环。出错走 2.5 节去查。整体调度流程在脑子里是这样的：

flowchart TB A["torchrun --nproc_per_node N"] --> B["initialize_distributed()
init_process_group + initialize_model_parallel"] B --> C["model_provider() →
TransformerConfig + GPTModel"] B --> D["get_train_data_iterator() →
BlendedMegatronDatasetBuilder + MockGPTDataset"] C --> E["forward_backward_func(
forward_step_func, data_iterator,
model, num_microbatches, ...)"] D --> E E --> F["optimizer.step()"] F --> G{"达到 max_steps?"} G -- 否 --> E G -- 是 --> H["save_distributed_checkpoint()"]

3.4 看一眼"工业版"入口

现在打开另一份脚本，对照感受落差：

repos/megatron-lm/examples/gpt3/train_gpt3_175b_distributed.sh

关键几行（行号取自基线 commit）：

GPUS_PER_NODE=8                                            # L7
WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))                  # L13

torchrun \                                                  # L21~
    --nproc_per_node $GPUS_PER_NODE \
    ... \
    pretrain_gpt.py \
    --num-layers 96 \                                       # 175B 的层数
    --hidden-size 12288 \
    --seq-length 2048 \
    --micro-batch-size 1 \
    --global-batch-size 1536 \
    --tensor-model-parallel-size 8 \                        # TP=8
    --pipeline-model-parallel-size 16                       # PP=16

你看到的就是 GPT-3 175B 的标准训练命令模板：8 卡 TP × 16 路 PP × 12 路 DP = 1536 卡训练。用第 1 章 1.2 节的公式验证：

$$W = \mathrm{TP} \times \mathrm{PP} \times \mathrm{DP} = 8 \times 16 \times 12 = 1536$$

pretrain_gpt.py 文件本身在仓库根目录。它和 examples/run_simple_mcore_train_loop.py 解决的是同一件事，但：

3.5 跑 345M 小预训练（中等难度，单机 8 卡）

如果你手上有一台 8 卡 A100/H100，可以试试官方的 examples/gpt3 配置改小版（README 链接给的"Your First Training Run"页面里有完整流程）。简化的最小步骤：

  1. 准备 tokenizer 和数据：使用 GPT-2 词表 + wikipedia 或 openwebtext 的一个小子集，经过 tools/preprocess_data.py 处理成 Megatron 的 indexed dataset 二进制；
  2. 把 examples/gpt3/train_gpt3_175b_distributed.sh 复制出来，改这些 flag：
    • --num-layers 12 --hidden-size 1024 --num-attention-heads 16（345M 配置）
    • --micro-batch-size 4 --global-batch-size 32
    • --tensor-model-parallel-size 1 --pipeline-model-parallel-size 1（小模型不用并行）
  3. bash train_gpt345m.sh 启动。

完整数据处理流程留到第 6 章。本章目标只是让你看到 loss 真的在降，对 Megatron 的"长相"和"启动方式"有个直观印象。

3.6 这章给你留下的应该是什么