Megatron-LM 深度学习与源码分析文档

第一部分：整体认知

Megatron-LM 解决的是“单 GPU 无法容纳、单机无法高效训练、普通 PyTorch DDP 通信/显存/调度效率不足”的大模型训练问题。对于 70B、100B、400B 级模型，瓶颈不是单一的 FLOPS，而是参数、梯度、优化器状态、activation、通信拓扑共同组成的系统问题。

以 BF16 Adam 训练为例，每个参数通常需要：

• parameter: 2 bytes
• gradient: 2 bytes 或 FP32 main grad
• optimizer master param / momentum / variance: 常见 8 到 12 bytes+
• activation: 随 batch、sequence、hidden、layers 增长，长上下文时常常比参数更麻烦

70B 模型仅参数 BF16 就约 140GB，Adam 状态会把单副本训练状态推到数百 GB。400B 模型更不可能只靠 DDP 复制完整模型。因此 Megatron-LM 的核心思想是：把模型、数据、序列、专家、优化器状态和 pipeline 执行都变成可控的分布式分片问题。

Megatron-LM 与其他框架关系：

• PyTorch Distributed：底层通信与 process group 基础设施，Megatron 在此之上构造 TP/PP/DP/CP/EP 组。
• DeepSpeed ZeRO：重点是 optimizer state、gradient、parameter 的 DP 维度切分；Megatron 更强调 Transformer 结构内的 TP/PP/SP，以及训练 iteration 调度。
• FSDP：以 parameter shard/all-gather/reduce-scatter 为中心，通用性强；Megatron 更“结构感知”，尤其是 Column/Row Parallel Linear、Attention/MLP 切分、pipeline schedule。
• Megatron + Distributed Optimizer：接近 ZeRO-1/ZeRO-2 思路，在 Megatron 自己的 DDP buffer 和 optimizer layout 上实现，源码入口见 distributed_data_parallel.py 和 distrib_optimizer.py。

实际训练中常见组合：

• 70B：TP=4/8，PP=4/8，DP 若干，SP 开启，activation checkpoint 开启。
• 100B：TP+PP 基本必需，distributed optimizer 减少 DP optimizer state。
• 400B：TP、PP、SP、CP、EP、distributed optimizer、重计算、overlap 基本都要参与。
• MoE 模型：参数量巨大但每 token 激活部分 expert，EP 和 all-to-all 成为核心。

第二部分：核心并行技术

1. Data Parallel，DP

DP 的基本逻辑是每个 DP rank 拥有一份模型副本，global batch 被切成多个 DP shard。每个 rank 做 forward/backward，之后同步梯度，再执行 optimizer step。

在 Megatron 中，DP 不只是 PyTorch DDP。Megatron 的 DDP 使用连续 gradient buffer、bucket、异步 reduce，并区分 expert/non-expert 参数。核心类是 distributed_data_parallel.py。它会把参数按 dtype、grad dtype、expert parallel 属性分组，然后放入 _ParamAndGradBuffer。

DP 梯度同步：

通信开销：

• AllReduce 通信量近似 $2\times \frac{N-1}{N}\times G$，$G$ 为梯度大小，$N$ 为 DP world size。
• 若使用 distributed optimizer，常用 reduce-scatter 让每个 DP rank 只持有一片 grad/optimizer state，然后参数更新后再 all-gather。
• 大 DP size 下 bucket 太小会 latency-bound，Megatron 会设置较大的默认 bucket，源码中可见 bucket size 逻辑 distributed_data_parallel.py。

2. Tensor Parallel，TP

Transformer 很适合 TP，因为主计算是大矩阵乘法：
$$Y=XW$$

设 $X\in {\mathbb{R}}^{SB\times H}$，$W\in {\mathbb{R}}^{H\times 4H}$ 或 $W\in {\mathbb{R}}^{H\times H}$。矩阵可沿输出维或输入维切分。

Column Parallel Linear

Column Parallel 把权重按输出列切：
$$W=[W_1,W_2,...,W_p]$$

每个 TP rank 计算：
$$Y_i=X{W}_i$$

最后可选择 all-gather：
$$Y=[Y_1,Y_2,...,Y_p]$$

源码对应 layers.py。它的 forward 里先计算本地 matmul，再根据 gather_output 调用 gather_from_tensor_model_parallel_region。通信封装在 mappings.py。

Row Parallel Linear

Row Parallel 把输入维切：
$$W=\left[\begin{array}{c}{W}_1\\ W_2\\ ...\\ W_p\end{array}\right],X=[X_1,X_2,...,X_p]$$

每个 rank 计算局部输出：
$$Y_i=X_i{W}_i$$

最终需要求和：
$$Y=\sum _i{Y}_i$$

源码对应 layers.py。非 SP 模式下 forward 后会 reduce_from_tensor_model_parallel_region，SP 模式下会 reduce_scatter_to_sequence_parallel_region。

Attention 中 QKV 如何切分

Self-Attention 的 QKV projection 通常使用 Column Parallel：
$$[Q,K,V]=X{W}_{qkv}$$

每个 TP rank 拿到部分 attention heads：
$$num\_heads\_per\_rank=\frac{num\_heads}{TP}$$

当前 Attention 基类里计算了 num_attention_heads_per_partition、hidden_size_per_attention_head，SelfAttention 中构建 linear_qkv，源码见 attention.py。

Attention 流程：

GPU kernel 主要发生在：

• QKV GEMM：cuBLAS / Transformer Engine / backend linear kernel
• RoPE：fused 或 elementwise kernel
• attention：FlashAttention / TE attention / PyTorch SDPA
• output projection GEMM
• TP collective：NCCL all-gather/all-reduce/reduce-scatter

MLP 如何切分

MLP 通常为：
$$Z=\varphi (X{W}_1),Y=Z{W}_2$$

Megatron 中 linear_fc1 是 Column Parallel，输出为 $4H/TP$ 或 SwiGLU 下更宽；linear_fc2 是 Row Parallel，把局部 hidden 聚合回 $H$。源码见 mlp.py，forward 见 mlp.py。

3. Pipeline Parallel，PP

PP 把 Transformer layers 按层切到不同 pipeline stage。比如 32 层、PP=4，每个 stage 8 层。当前 TransformerBlock 会按 pipeline rank 计算 layer offset，并只构建本 stage 的 layers，源码见 transformer_block.py。

核心问题是 pipeline bubble。若只有一个 batch：

time ->

GPU0: F0  F1  F2  F3  B3  B2  B1  B0
GPU1:     F0  F1  F2  F3  B3  B2  B1  B0
GPU2:         F0  F1  F2  F3  B3  B2  B1  B0
GPU3:             F0  F1  F2  F3  B3  B2  B1  B0
       warmup        steady 1F1B       cooldown

Megatron 的非 interleaved 1F1B 在 schedules.py。核心流程：

• warmup：前面 stage 先连续 forward 若干 microbatch。
• steady：每个 stage 做 1 forward + 1 backward。
• cooldown：剩余 backward 排空。
• P2P 通信用 recv_forward、send_forward_recv_backward、send_backward_recv_forward。

Interleaved PP 引入 virtual pipeline stage，例如 PP=4，VPP=2，每张 GPU 拿两个非连续 layer chunk：

GPU0: layers [1,2]   + [9,10]
GPU1: layers [3,4]   + [11,12]
GPU2: layers [5,6]   + [13,14]
GPU3: layers [7,8]   + [15,16]

这样可以减少 bubble，但增加调度复杂度和激活驻留压力。初始化参数说明见 parallel_state.py。

4. Sequence Parallel，SP

TP 会导致某些 activation 在每个 TP rank 上复制，尤其 LayerNorm、Dropout、Residual 这类按 hidden 操作的模块。长 sequence 下 activation memory 随 $S\times B\times H\times L$ 增长，SP 的核心是把 sequence 维也切开，让每个 TP rank 只保存 $S/TP$ 的 activation。

SP 通信模式：

• 需要完整 hidden 做某些线性前，all-gather sequence。
• RowParallel 输出后，用 reduce-scatter 回 sequence-sharded activation。
• 源码封装在 mappings.py 和 mappings.py。

SP 与 activation checkpoint 区别：

• SP 是数据布局改变，减少常驻 activation。
• checkpoint 是不存或少存 activation，backward 时重算。
• 二者可以叠加：SP 降低保存量，checkpoint 降低保存集合。

5. Expert Parallel，EP

MoE 模型每层有多个 expert，每个 token 只路由到 top-k expert。若 expert 数很多，不能每张 GPU 保存全部 expert，所以需要 EP：不同 rank 持有不同 experts。

MoE forward：

源码路径：

• MoE layer: moe_layer.py
• Router / TopKRouter: router.py
• AllToAll wrapper: mappings.py
• MoE dispatch/combine forward: moe_layer.py

DeepSeek/Qwen MoE 这类模型通常关心：

• top-k routing 负载均衡
• token dropping / capacity factor
• all-to-all-v 不均匀通信
• grouped GEMM 性能
• shared expert overlap
• EP 与 TP 的组合，即 expert tensor parallel

第三部分：Megatron-LM 源码分析

当前重点目录

megatron/
  training/
    initialize.py
    training.py
  core/
    parallel_state.py
    tensor_parallel/
    pipeline_parallel/
    distributed/
    transformer/
    optimizer/

用户请求中的目录 megatron/core/tensor_parallel、pipeline_parallel、distributed、transformer、optimizer 是当前代码中的真实主路径。

启动流程：从 torchrun 开始

典型启动：

torchrun --nproc_per_node=8 pretrain_gpt.py \
  --tensor-model-parallel-size 4 \
  --pipeline-model-parallel-size 2 \
  ...

调用链：

关键源码：

• pretrain 入口：training.py
• 初始化：initialize.py
• torch.distributed 初始化：initialize.py
• 并行组创建：parallel_state.py
• model/optimizer setup：training.py

Training Loop：一次 iteration

Megatron 主循环在 train，单步在 train_step。

一次 iteration：

对应源码行为：

• model_chunk.zero_grad_buffer() 清空连续 grad buffer。
• forward_backward_func(...) 根据 PP/VPP 状态选择 no-pipeline、1F1B 或 interleaved。
• config.finalize_model_grads_func = finalize_model_grads 在训练初始化时设置。
• optimizer.step() 执行混合精度、grad norm、distributed optimizer 更新。

Transformer Layer 调用关系

关键源码：

• TransformerBlock 构建 layers：transformer_block.py
• TransformerLayer 组装 attention/mlp：transformer_layer.py
• Attention 抽象和 SelfAttention：attention.py
• MLP：mlp.py
• Column/Row Linear：layers.py

哪些地方发生 GPU kernel：

• Linear GEMM：QKV、O projection、MLP fc1/fc2。
• Attention kernel：FlashAttention、TE attention、SDPA。
• Elementwise/fused：LayerNorm、RoPE、bias+GELU/SwiGLU、bias+dropout+add。
• NCCL kernel：TP/DP/EP/PP 通信。
• Optimizer kernel：FusedAdam、multi-tensor update、grad scaling、norm clipping。
• CUDA Graph：当前训练代码支持 full-iteration graph 和 optimizer graph 包装，见 training.py。

第四部分：GPU 通信和性能优化

NCCL 通信模式

Collective	典型用途	Megatron 场景
AllReduce	多 rank 求和并广播结果	DP gradient sync，RowParallel 输出聚合
ReduceScatter	求和后每 rank 保留一片	SP、distributed optimizer、DP grad shard
AllGather	每 rank 收集所有 shard	ColumnParallel output gather，SP 前恢复 sequence
AllToAll	每 rank 给每 rank 发送不同 token	MoE EP token dispatch/combine

TP 通信：

• ColumnParallel：forward 可 all-gather；backward 可能 reduce-scatter/split。
• RowParallel：forward all-reduce 或 reduce-scatter；backward 需要处理 input grad。
• 封装见 mappings.py。

DP 通信：

• 普通 DDP：bucket all-reduce。
• Distributed optimizer：grad reduce-scatter，param all-gather。
• Megatron DDP buffer 负责 overlap 和 bucket 管理。

EP 通信：

• router 后 token dispatch 走 all-to-all 或 all-gather。
• combine 阶段反向通信恢复 token 顺序。
• 性能瓶颈通常是不均匀 all-to-all-v、token permutation、grouped GEMM 小矩阵效率。

通信计算 Overlap

Overlap 的本质是用 CUDA stream 和异步 NCCL kernel 把通信藏在 GEMM/attention 计算后面。

Megatron 中常见 overlap：

• overlap grad reduce：backward 产生某个 bucket 的 grad 后立刻 async reduce。
• overlap parameter gather：下一层/下一 microbatch 需要参数前提前 all-gather。
• pipeline communication overlap：P2P send/recv 与本 stage compute 交错。
• TP comm overlap：Transformer Engine user buffer 支持 GEMM 与 TP collective overlap，初始化见 initialize.py。

为什么提升 MFU：
$$T_{iter}\approx T_{compute}+T_{exposed\_comm}$$

Overlap 不是减少通信总量，而是减少 exposed communication：
$$T_{exposed\_comm}=\max (0,T_{comm}-T_{overlapped})$$

性能指标

• MFU：Model FLOPs Utilization，模型理论 FLOPs / GPU 峰值 FLOPs / 时间。Megatron 会计算训练 FLOPs，相关逻辑在 training.py 中。
• FLOPS：关注 GEMM/Attention 是否接近硬件峰值。
• Bandwidth：HBM、NVLink、PCIe、IB/RDMA、collective bandwidth。
• Latency：小 bucket、小 all-to-all、pipeline P2P 对 latency 敏感。
• Scaling efficiency：
  $$Efficiency(N)=\frac{Throughput(N)}{N\times Throughput(1)}$$

第五部分：GPU 公司视角

如果你加入 GPU 公司做 AI Infra/GPU 软件生态，Megatron-LM 里最值得盯的不是 API，而是这些边界：

1. GPU 适配相关模块：

• torch.cuda 假设：初始化、device、stream、event、CUDA Graph。
• NCCL process group：TP/DP/PP/EP collective。
• Transformer Engine / Apex / FlashAttention：大量高性能 kernel 依赖。
• Fused optimizer、fused activation、fused norm、grouped GEMM。
• FP8/FP4 路径：fp8_utils.py、fp4_utils.py、TE integration。

2. CUDA/NCCL 依赖点：

• initialize.py 强制 CUDA availability。
• TP collective 在 mappings.py。
• PP P2P 在 schedules.py 和 p2p_communication.py。
• MoE all-to-all 在 token dispatcher 和 mappings.py。
• optimizer 依赖 TE/Apex FusedAdam，见 distrib_optimizer.py。

3. Kernel optimization 重点：

• QKV GEMM + bias fusion。
• FlashAttention / MLA / GQA/MQA attention。
• SwiGLU/GELU fused activation。
• LayerNorm/RMSNorm。
• RowParallel reduce-scatter fused add/norm。
• MoE grouped GEMM。
• token permutation/unpermutation。
• FP8/FP4 quant/dequant and scaling。
• optimizer multi-tensor update。

4. 支持国产 GPU 需要解决：

• PyTorch backend 是否完整支持该设备。
• torch.distributed backend 是否能替代 NCCL。
• collective library 是否支持 all-reduce/reduce-scatter/all-gather/all-to-all，并有足够性能。
• CUDA Graph 等价能力。
• TE/Apex/FlashAttention 替代 kernel。
• dtype 支持：BF16、FP16、FP8、FP4。
• memory allocator、stream/event、profiling 工具链。
• 编译器：Triton/CUTLASS/CUTE 等生态替代。

5. GPU 软件栈替代路线：

第六部分：30 天源码阅读路线

Day	阅读文件	目标	实验
1	README、docs/get-started	了解训练参数体系	跑 help
2	pretrain_gpt.py	入口脚本	找 model_provider
3	training.py	pretrain 主流程	画调用链
4	initialize.py	分布式初始化	打印 rank/world
5	parallel_state.py	TP/PP/DP group	验证 rank groups
6	arguments/config	参数如何进 config	改 TP/PP 参数
7	总结 Week1	整体架构	写一页图
8	layers.py	ColumnParallel	手算矩阵 shape
9	layers.py	RowParallel	手算通信点
10	mappings.py	autograd collective	跟 backward
11	attention.py	QKV/head 切分	打印 qkv shape
12	mlp.py	MLP TP	比较 SwiGLU
13	schedules.py	1F1B	画 timeline
14	总结 Week2	TP/PP/DP/SP	跑 2/4 GPU
15	transformer_block.py	block 构建	layer offset
16	transformer_layer.py	layer forward	hook module
17	distributed_data_parallel.py	DDP buffer	bucket 观察
18	param_and_grad_buffer.py	grad layout	看 buffer range
19	distrib_optimizer.py	optimizer shard	看 state shard
20	checkpointing	分布式 checkpoint	保存/恢复
21	总结 Week3	iteration 全链路	写 sequence diagram
22	moe_layer.py	MoE forward	打印 routing
23	router.py	top-k router	观察负载
24	token_dispatcher	EP all-to-all	Nsight/NCCL log
25	fusions/TE	fused kernels	开关 fusion
26	CUDA Graph	graph capture	比较 iteration time
27	overlap 参数	grad/param overlap	比较 MFU
28	profiler	Nsight Systems	标通信空洞
29	GPU 适配点	backend/kernels	列替代清单
30	综合复盘	训练框架设计	写技术报告

第七部分：实践实验

实验 1：单 GPU 运行 GPT 小模型

目标：理解最小训练链路，不引入 TP/PP 干扰。

代码/运行方式：

torchrun --nproc_per_node=1 pretrain_gpt.py \
  --num-layers 2 \
  --hidden-size 128 \
  --num-attention-heads 4 \
  --seq-length 128 \
  --max-position-embeddings 128 \
  --micro-batch-size 2 \
  --global-batch-size 8 \
  --train-iters 20 \
  --lr 1e-4 \
  --bf16 \
  --tensor-model-parallel-size 1 \
  --pipeline-model-parallel-size 1

观察指标：

• iteration time
• loss 是否下降
• GPU memory
• 单 GPU kernel 组成

预期结果：无 collective 或 collective 很少，主要是 GEMM/attention/optimizer kernel。

实验 2：2 GPU Tensor Parallel

目标：观察 Column/Row Parallel 通信。

运行方式：

NCCL_DEBUG=INFO torchrun --nproc_per_node=2 pretrain_gpt.py \
  --num-layers 2 \
  --hidden-size 512 \
  --num-attention-heads 8 \
  --seq-length 512 \
  --micro-batch-size 1 \
  --global-batch-size 2 \
  --train-iters 20 \
  --bf16 \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 1

观察指标：

• NCCL all-reduce/all-gather/reduce-scatter
• linear_qkv 输出 shape
• RowParallel 输出后的 reduce
• MFU 与单 GPU 对比

预期结果：显存下降，通信出现，隐藏层维度被 TP 切分。

实验 3：4 GPU Pipeline Parallel

目标：观察 microbatch、bubble、1F1B。

运行方式：

torchrun --nproc_per_node=4 pretrain_gpt.py \
  --num-layers 8 \
  --hidden-size 512 \
  --num-attention-heads 8 \
  --seq-length 512 \
  --micro-batch-size 1 \
  --global-batch-size 8 \
  --train-iters 20 \
  --bf16 \
  --tensor-model-parallel-size 1 \
  --pipeline-model-parallel-size 4

观察指标：

• 每个 rank 层数
• PP P2P send/recv
• warmup/steady/cooldown 时间
• microbatch 数变化对 bubble 的影响

预期结果：microbatch 越多，bubble 占比越低，但 activation 驻留越多。

实验 4：分析 NCCL 通信

目标：建立 TP/DP/EP collective 与性能的直觉。

运行方式：

NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=COLL torchrun --nproc_per_node=4 pretrain_gpt.py \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 1 \
  --num-layers 4 \
  --hidden-size 1024 \
  --num-attention-heads 16 \
  --seq-length 1024 \
  --micro-batch-size 1 \
  --global-batch-size 4 \
  --train-iters 10 \
  --bf16

观察指标：

• collective 类型
• message size
• bucket size
• 是否有通信暴露在 compute 间隙
• Nsight Systems timeline

预期结果：TP collective 在 layer 内高频出现，DP collective 在 backward bucket 维度出现。

实验 5：修改一个模块并测试性能

目标：从 GPU Infra 角度完成一次小型优化闭环。

建议修改点：

• 在 mlp.py 周围加入 NVTX range 或对比 fused/non-fused activation。
• 或切换 bias_activation_fusion，观察 SwiGLU/GELU kernel 数量。
• 或调整 DDP bucket size，观察 overlap 与 latency。

运行方式：

torchrun --nproc_per_node=2 pretrain_gpt.py \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 1 \
  --num-layers 4 \
  --hidden-size 1024 \
  --num-attention-heads 16 \
  --seq-length 1024 \
  --micro-batch-size 1 \
  --global-batch-size 4 \
  --train-iters 50 \
  --bf16

观察指标：

• tokens/sec
• iteration time p50/p95
• MFU
• kernel count
• exposed NCCL time
• peak memory

预期结果：融合 kernel 能减少 launch 和 HBM 往返；bucket/overlap 参数会改变通信暴露时间，但不一定总是更快。

最后的阅读心法

Megatron-LM 不是“一个训练 API”，而是一个把 Transformer 数学结构、GPU kernel、NCCL collective、显存布局、pipeline 调度和 optimizer state sharding 绑在一起的训练系统。你阅读源码时应始终追问四件事：

• 这个 tensor 当前被按哪个维度切了？
• 这个 rank 拥有哪些参数、activation、grad、optimizer state？
• 下一次 collective 是为了恢复完整值、求和，还是重新分片？
• 这段通信是否在 critical path 上，能不能 overlap 或 fuse？

我没有修改仓库文件；本次工作是基于当前源码做结构化分析与学习路线设计。