Megatron-LM 源码精读——大模型并行训练全景

https://blog.csdn.net/Antai_ZHU/article/details/144409187

本文基于 NVIDIA Megatron-LM 源码（megatron/core/，即 Megatron-Core）。目标：让你在面试时能从"5D 并行 + rank 排布"讲到"ColumnParallelLinear 的 autograd 通信插入"、从"1F1B vs Interleaved"讲到"DDP 的 param/grad buffer 与 distributed optimizer"、从源码结构讲到工程取舍。读完这篇，Megatron 这关基本通杀。

一、Megatron 是什么、为什么面试要问

Megatron-LM 是 NVIDIA 开源的大模型分布式训练框架，论文从 2019（张量并行）一路到 2022（PP+SP）、2023（CP、MoE）。它的核心价值是把"怎么把一个 Transformer 切到几千张 GPU 上训"这件事做成了可组合的并行库——Megatron-Core（megatron/core）。

面试为什么爱问：

它是工业界 3D/5D 并行的事实标准，DeepSpeed、Megatron-DeepSpeed、Nemo、各厂自研框架都借鉴/对接它的并行接口。
它把"并行"做成了和模型结构解耦的算子级改造——nn.Linear 换成 ColumnParallelLinear，通信藏在 autograd 的 forward/backward 里，对上层透明。
它是理解大模型训练系统全貌最好的教材：通信/计算 overlap、显存/带宽权衡、调度微批次、ZeRO 与 TP 的取舍，全在里面。

源码鸟瞰（megatron/core/）：

parallel_state.py          # 5D 并行的进程组(rank 排布),整篇地基
tensor_parallel/           # 张量并行: ColumnParallelLinear/RowParallelLinear/ mappings/ random
pipeline_parallel/         # 流水并行: schedules(1F1B/Interleaved) + p2p_communication
distributed/               # 数据并行: DDP + param_and_grad_buffer + 分布式优化器(ZeRO)
transformer/               # Transformer 模型: attention/mlp/transformer_layer/moe
models/                    # GPT/Llama/Mamba 等
datasets/                  # 数据加载
dist_checkpointing/        # 分布式 checkpoint(分片保存/加载)
optimizer/                 # 优化器(含 distributed/ZeRO)
num_microbatches_calculator.py  # 流水调度微批次数计算

主线：parallel_state 决定每张卡属于哪些进程组 → tensor_parallel 在组内切张量 → pipeline_parallel 在层间流水 → distributed 在 DP 间同步梯度/分片参数 → transformer 把这些并行算子拼成模型。

二、5D 并行与 rank 排布（最该背的地基）

Megatron 的并行是五个维度的笛卡尔积：

维度	缩写	切什么	谁负责同步	权重是否复制
数据并行	DP	切 batch	梯度 all-reduce	是（复制）
张量并行	TP	切单层权重矩阵	前向/反向 collective	否（每个 TP rank 只有一片）
流水并行	PP	切层（纵向）	P2P send/recv 激活	否（每段不同层）
上下文并行	CP	切序列长度	注意力 all-gather/all-reduce KV	是（权重复制）
专家并行	EP	切 MoE 的 expert	all-to-all dispatch	否（每 EP rank 不同 expert）

rank 排布（面试必背）

源码 parallel_state.py 的 initialize_model_parallel(order="tp-cp-ep-dp-pp") 给出默认 rank 排布顺序。含义：把全局 rank $r \in [0, W)$ 按 tp-cp-ep-dp-pp 从快到慢（内层到外层）的优先级嵌套——

$r = \text{tp} + \text{cp}\cdot T + \text{ep}\cdot T\cdot C + \text{dp}\cdot T\cdot C\cdot E + \text{pp}\cdot T\cdot C\cdot E\cdot D$

其中 $T=\text{tp\_size}, C=\text{cp\_size}, E=\text{ep\_size}, D=\text{dp\_size}$ ， $W=T\cdot C\cdot E\cdot D\cdot P$ 。

为什么这样排：相邻的 rank 在最内层维度（TP）变化，于是同一台机器/同一个 NVLink 域内的 GPU 落在同一个 TP 组——TP 通信（每层都发、带宽最大）走机内 NVLink。越往外层（DP/PP）变化越慢，组间跨机走 IB/RoCE（带宽小、频率低）。"快维度放机内、慢维度放跨机"是排布的核心原则。

面试金句：“Megatron 的 rank 排布默认 tp-cp-ep-dp-pp，把通信最重最频繁的 TP 放最内层落在 NVLink 域、PP 放最外层跨机。每个 rank 同时属于多组 process group（TP/CP/EP/DP/PP），算子通信按需选组。”

源码里这些 group 就是开头那一大串 _TENSOR_MODEL_PARALLEL_GROUP、_PIPELINE_MODEL_PARALLEL_GROUP、_DATA_PARALLEL_GROUP_WITH_CP、_EXPERT_* 等全局变量——initialize_model_parallel 用 torch.distributed.new_group 一个个建出来。读懂这一堆 group 的命名，就读懂了 Megatron 的并行结构。

三、张量并行（TP）：把一个矩阵切到多卡

TP 是 Megatron 的起家之作。核心：把 nn.Linear 的权重矩阵按列或行切到 TP 组各卡，把跨卡的 collective 通信藏在 autograd 的 forward/backward 里，使上层代码看不到通信。

3.1 ColumnParallelLinear（按列切）

权重 $W \in \mathbb{R}^{in\times out}$ ，按输出维度（列）切成 $T$ 片，每卡持有 $W_i \in \mathbb{R}^{in\times out/T}$ 。前向： $Y_i = X W_i$ （ $X$ 在所有卡复制），各卡得到输出的一段。前向不需要通信（输入复制、各算各的），反向需要一次 all-reduce 把梯度 $\frac{\partial L}{\partial X} = \sum_i \frac{\partial L}{\partial Y_i} W_i^T$ 汇总。

源码（tensor_parallel/layers.py）ColumnParallelLinear.__init__：

1 2	self.output_size_per_partition = divide(output_size, world_size) # 按列切 self.weight = Parameter(torch.empty(output_size_per_partition, input_size, ...))

注意 PyTorch 的 linear 是 $X W^T$ ，所以权重 shape 是 [out_per_part, in]。is_expert 时换用 get_expert_tensor_parallel_world_size()，expert 的 TP 与普通层可不同。

3.2 RowParallelLinear（按行切）

按输入维度（行）切，每卡 $W_i \in \mathbb{R}^{out\times in/T}$ ，输入 $X$ 也切了： $Y_i = X_i W_i$ 。前向需要 all-reduce 汇总 $\sum_i Y_i$ ，反向不需要通信（梯度天然分散）。

3.3 MLP 的经典切法（行+列组合，省一次通信）

Megatron MLP = 两个 Linear： $d \to 4d \to d$ 。切法：

1
2
3

X(d) ─► [ColumnParallel: d→4d/T 各卡]  ─► GeLU ─► [RowParallel: 4d/T→d 各卡] ─► all-reduce ─► Y(d)
        前向不通信(输入复制)                      前向 all-reduce 汇总
        反向 all-reduce(对X的梯度)                反向不通信

关键洞察：Column 的前向"不通信"刚好接 Row 的反向"不通信"——中间是分片的，两个通信点合并成整个 MLP 一前一后各一次 all-reduce，而不是每层都通信。Attention 同理（ $QKV$ 列切、 $O$ 投影行切）。这就是 Megatron TP 的精髓：靠算子配对把通信次数压到每层 all-reduce 两次。

3.4 通信藏在 autograd 里（源码核心）

mappings.py 里定义了若干 torch.autograd.Function，把通信塞进前向/反向：

_CopyToModelParallelRegion：forward 不做事、backward 做 all-reduce（对应 Column 前向不通信反向 reduce）。
_ReduceFromModelParallelRegion：forward all-reduce、backward 不做事（对应 Row 前向 reduce）。
_ScatterToModelParallelRegion / _GatherFromModelParallelRegion：序列并行的切/聚。
_ReduceScatterToSequenceParallelRegion / _AllGatherFromSequenceParallelRegion：SP 的 reduce-scatter/all-gather。

ColumnParallelLinear.forward 里就一句 output = linear_with_grad_accumulation_and_async_allreduce(..., sequence_parallel, ...)——它内部根据 sequence_parallel 开关决定走 all-reduce 还是 all-gather。通信被包成一个 autograd Function，自动出现在正确的反向位置，对写模型的人完全透明。

3.5 序列并行（SP）

标准 TP 在非层通信点要 all-reduce 全 $[s,b,d]$ 激活，通信量 $\propto s\cdot b\cdot d$ 。SP 把序列维度 $s$ 也切到 TP 组各卡：层间用 reduce-scatter / all-gather（通信量从 $sbd$ 降到 $sbd/T$ ，是 all-reduce 的一半），代价是 attention 里需要再聚回全序列。SP 对长序列显存和带宽都友好，是当前默认。

3.6 初始化：CPU vs GPU、affine

源码 _initialize_affine_weight_cpu / _gpu：CPU 初始化时按 rank 切好再各自加载（省显存，适合大模型 init）；GPU 初始化时各卡独立 init 同一个分布、靠 offset 保证各卡拿到的是不同切片（快，省 CPU↔GPU 传输）。这就是为什么大模型初始化不爆显存——权重一上来就是分片的。

四、流水并行（PP）：把层切到多卡再流水

TP 切不动太多（受机内 GPU 数限制），扩展到上千卡靠 PP + DP。PP 把模型按层切成 $P$ 段，每段一张/一组卡，前向激活从第 0 段 P2P 发到第 1 段，逐段流水。

4.1 朴素 PP 与气泡

朴素做法：前向跑完整个 batch 再反向——前向时只有一段在算，其它段空等，气泡巨大。改进是把 batch 切成 $M$ 个微批次（microbatch），让多段像流水线一样同时干不同微批次。

4.2 1F1B 调度（Megatron 默认）

pipeline_parallel/schedules.py 的 forward_backward_pipelining_without_interleaving。核心思想：一个前向紧接着一个反向（One Forward One Backward），让前向和反向交错进行，尽早启动反向释放激活显存。

阶段（ $P$ 段、 $M$ 个微批）：

Warmup：第 0 段连做若干前向填满流水线（前向数 $\approx P-\text{rank}-1$ ）。
Steady state（1F1B）：每段每轮做"1 前向 + 1 反向"，前向推进、反向回收激活。
Cooldown：剩下的反向收尾。

气泡占比 $\approx \frac{P-1}{M}$ —— $M$ 越大气泡越小，但要更多微批（更多显存存激活）。这就是为什么大模型训练要把 batch 切很多 microbatch。

4.3 Interleaved 1F1B（Virtual Pipeline， interleaved）

forward_backward_pipelining_with_interleaving（论文 Interleaved Pipeline Parallelism）。思路：把每段的 $L/P$ 层再分成 $V$ 个 chunk，每卡持有 $V$ 个不连续的层段，轮流前向/反向。例（注释里给的例子， $P=4, V=2, L=16$ ）：

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

效果：流水线"颗粒"从"段"变"段/V"，气泡降到 $\approx \frac{P-1}{M\cdot V}$ ，代价是 P2P 通信次数 $\times V$ 、激活显存略增。通信带宽紧张时不划算，带宽富裕时显著降气泡——这就是为什么 PP 通信要 NVLink/IB 充足。

4.4 P2P 通信与异步

p2p_communication.py：send_forward_recv_backward、send_forward_backward_recv_forward_backward 等——前向 send 和反向 recv 重叠（batch_isend_irecv 一次发多个），让 P2P 通信藏在前一卡的算力里。recv_forward/recv_backward 先占好 buffer，避免动态分配。这就是 PP 调优的微观功夫。

4.5 微批次数与调度表

num_microbatches_calculator.py 算每个 iter 用几个 microbatch；get_pp_rank_microbatches / get_schedule_table 生成 interleaved 的调度表（哪个虚拟微批次什么时候做前向/反向），避免死锁、平衡各卡负载。

面试金句：“PP 的本质是减小气泡。朴素 1F1B 把气泡压到 $(P-1)/M$ ；Interleaved 用 $V$ 个虚拟段把颗粒变细气泡再降到 $(P-1)/(MV)$ ，代价是 P2P 通信 $\times V$ 。所以 PP 的优化永远是拿通信带宽换气泡、拿显存换气泡的权衡。”

五、数据并行（DP）：DDP、buffer、distributed optimizer

DP 是最老牌的并行——每卡持完整模型副本，各算不同 batch，反向后 all-reduce 梯度。Megatron 的 distributed/distributed_data_parallel.py（DDP）在 PyTorch DDP 之上做了为 TP/PP 量身定制的工程强化。

5.1 Param/Grad Buffer：把碎片参数攒成大桶

问题：Transformer 有几十上百个 Parameter，每个都 all-reduce 一次效率极低（NCCL 每次有固定开销、小消息带宽利用率差）。源码 param_and_grad_buffer.py 的 _ParamAndGradBuffer / _ParamAndGradBucketGroup：

把所有参数和梯度连续地铺进一两个大 buffer（_ParamAndGradBucket），每个参数是 buffer 里的一个 view。
梯度反向时直接写进 grad buffer（同一块连续显存），攒满一桶再 start_grad_sync 一次 all-reduce。
好处：一次大 all-reduce 替代几十次小 all-reduce，带宽利用率高；连续显存利于异步、利于 overlap。

5.2 Gradient Accumulation Fusion（梯度累加融合）

反向算权重梯度 $W^T \cdot g$ 时，Megatron 把"多个 microbatch 的权重梯度直接累加进 grad buffer"这一步融合进 wgrad kernel——CPU/Python 不介入、不分配中间张量、直接 atomic add 到 buffer。这就是 layers.py 里 gradient_accumulation_fusion 开关的作用。一个 iter 内 $M$ 个 microbatch 的权重梯度只算一次（累加），最后才 all-reduce。

5.3 通信/计算 overlap：bucket ready + async

register_grad_ready(param)：某个参数的反向算完就标记该 bucket"ready"，bucket 内所有参数 ready 立刻发起异步 all-reduce，和后续反向计算 overlap——反向还在算别的参数，前面的梯度已经在网线上飞了。这是 DDP 性能的关键。no_sync() 上下文则关掉同步、攒梯度（用于 ZeRO-2/3 跨 step 累积）。

5.4 Distributed Optimizer（ZeRO-style 分片）

distributed/ 下的 param_and_grad_buffer 支持分布式优化器：把 optimizer state（Adam 的 $m,v$ ）甚至权重分片到 DP 组各卡（对应 ZeRO-1/2/3）。配合 optimizer/ 下的实现，每卡只存/算自己那片参数的 $m,v$ ，更新后 all-gather 回完整权重。这就是 Megatron 自己的 ZeRO——和 DeepSpeed 思路一致，但和 TP/PP 深度耦合，rank 排布上 DP 是最外层之一。

面试金句：“Megatron DDP 的三板斧：param/grad 连续 buffer 替代碎片 all-reduce、gradient accumulation fusion 把多 microbatch 权重梯度融进 wgrad kernel 一次累加、bucket ready 触发 async all-reduce 和反向 overlap。再叠加 distributed optimizer 分片优化器状态/权重，等价 ZeRO-1/2/3 但与 TP/PP 共生。”

六、上下文并行（CP）：长序列不爆显存

context_parallel_size：把输入序列长度 $s$ 切到 CP 组各卡。动机：长上下文（32k/128k）时 attention 的 $O(s^2)$ 激活爆显存，单卡放不下。

机制（源码注释给的例子， $T=4, C=2$ ，8 GPU）：序列分两段，GPU0-3 算前半、GPU4-7 算后半，靠 [GPU0,GPU4] 等 4 个 CP 组交换 KV。

实现要点：

权重在 CP 组复制，所以反向时权重梯度要 all-reduce（CP 复用 DP 组做这个，所以有 _DATA_PARALLEL_GROUP_WITH_CP）。
Attention 里每张卡只算自己那段 query 对全序列 K/V 的注意力，需要 all-gather KV 或 ring 交换——靠 FlashAttention 的 ring/all-gather 变体实现，CP 通信量 $\propto s/T$ 级别、藏在前向算力里。
注意 CP 和 TP 都"切序列维度"——TP 切的是激活的 $s$ （SP），CP 切的是输入序列，CP 跨的是"不同 GPU 处理不同 token 段"，和 SP 互补可叠加。

CP 是长上下文训练的标配（Llama-3 128k、Qwen 长上下文都靠它）。面试可点：“CP 切序列长度降 attention 激活显存、权重复制故复用 DP 组 reduce 梯度，靠 FlashAttention ring attention 在组内传 KV。”

七、Transformer 模块与 MoE

transformer/ 下把并行算子拼成可训练的 Transformer。

transformer_layer.py：一个 Transformer 层（attn + mlp + residual + norm + fp8 等）。层与层之间就是 PP 切的单位。
attention.py / dot_product_attention.py：注意力。TP 切 $Q/K/V/O$ （QKV 列切、O 行切），CP 在这里交换 KV，集成 FlashAttention（把 $O(s^2)$ 激活降到 $O(s)$ ）。新版还有 multi_latent_attention.py（MLA，DeepSeek-V2 那套，KV 压缩）。
mlp.py：MLP，前面讲过的 Column→Row 配对切法。Llama 系用 SwiGLU（三个矩阵）。
moe/：MoE 层。dispatcher 把 token 按路由分发到 expert（all-to-all），expert 算完再 all-to-all 回收。EP 切 expert、TP 切 expert 内部矩阵——_EXPERT_* 那一堆 group 就是给 MoE 的。token-per-expert 负载均衡靠辅助 loss。
transformer_block.py：整层堆叠，PP 切到各 stage。
cuda_graphs.py：对固定 shape 的小层用 CUDA Graph 抓图，减少 kernel launch 开销（decode/小算子多的场景提速）。
fusions/：各种融合 kernel（FlashFusedMLP、rotary embedding、norm+bias 等），把多个小算子合成一个大 kernel，减少访存和 launch。

面试可点：“Megatron 的注意力天生支持 TP+CP+FlashAttention，MLP 靠 Column/Row 配对把通信压到每层两次 all-reduce，MoE 用 all-to-all dispatch + EP/TP 嵌套切 expert。”

八、显存、精度与通信 overlap 的工程取舍

把这些串成面试能讲的"系统权衡"：

8.1 显存账（和显存计算法则篇呼应）

权重：TP/PP/EP 把权重切到各卡，每卡 $\frac{P_{\text{本层}}}{T\cdot P}$ （PP 是按层切，各卡只存自己段）。
激活：SP 把序列维度切了降激活；activation checkpointing 重计算省激活；CP 切长序列降 attention 激活；FlashAttention 把 $s^2$ 降 $s$ 。
优化器状态：distributed optimizer 分片到 DP，等价 ZeRO-1/2/3。
TP 省权重不省激活，DP/ZeRO 省优化器状态不省权重，PP 省权重按层切——三者正交可叠加。

8.2 混合精度 / FP8

主流 bf16/fp16 训练 + fp32 master weight + Adam（每参数 16 bytes，见显存篇）。
fp8_utils.py：H100 的 FP8（E4M3 前向、E5M2 反向），权重/激活/梯度量化到 1 byte，权重显存减半、算力翻倍。靠 per-tensor scaling + amax 历史，配合 _TENSOR_AND_DATA_PARALLEL_GROUP 跨 TP/DP 做 amax all-reduce。

8.3 通信/计算 overlap

TP：all-reduce 藏在 autograd，和反向计算 overlap。
PP：P2P send/recv 用 batch_isend_irecv 重叠。
DDP：bucket ready 触发 async all-reduce 与反向 overlap。
SP：reduce-scatter 替代 all-reduce 省一半带宽。
通用工具 timers.py 测各段开销定位瓶颈。

8.4 容错与 checkpoint

dist_checkpointing/：分片保存/加载 checkpoint——每卡只存自己那片参数/优化器状态（按 TP/PP/EP/ZeRO 分片），加载时按当前并行配置重排（reshard），支持不同并行度间迁移。
训练中断恢复靠保存 optimizer state + RNG state + 数据指针，rerun_state_machine.py 处理重试。

九、怎么读源码（给你的路线）

parallel_state.py 先读 initialize_model_parallel，画一张"五个维度 → 一堆 process group"的对应表。这是地基，不懂它后面全悬空。
tensor_parallel/layers.py 读 ColumnParallelLinear/RowParallelLinear 的 __init__ 和 forward，配合 mappings.py 的几个 autograd.Function，理解"通信藏在哪一步"。
pipeline_parallel/schedules.py 读 get_forward_backward_func 的分派 + forward_backward_pipelining_without_interleaving 的 warmup/steady/cooldown 三段，再扫 interleaved。
distributed/param_and_grad_buffer.py + distributed_data_parallel.py 读 buffer 分桶、register_grad_ready、no_sync、distributed optimizer 分片。
transformer/transformer_layer.py + attention.py 把上面的并行算子拼起来，看一个 layer 全流程。
跑 examples/ 下一个最小 GPT 训练脚本，用 --tensor-model-parallel-size 2 --pipeline-model-parallel-size 2 等参数实际跑，对照日志理解 rank 排布。

十、面试速答清单

Q1：Megatron 的并行有几个维度？rank 怎么排布、为什么？

5 维：DP/TP/PP/CP/EP。默认 rank 排布 tp-cp-ep-dp-pp——把通信最重最频繁的 TP 放最内层落在同一 NVLink 域（机内高带宽），PP 放最外层跨机走 IB/RoCE。每个 rank 同时属于多组 process group（TP/CP/EP/DP/PP），算子按需选组。原则：快维度机内、慢维度跨机。

Q2：ColumnParallelLinear 和 RowParallelLinear 各自怎么切、前后向在哪通信？

Column 按输出维度切、输入复制，前向不通信、反向 all-reduce 对输入的梯度；Row 按输入维度切、输入也切，前向 all-reduce 汇总输出、反向不通信。MLP 里 Column→激活→Row 配对，使整个 MLP 一前一后各一次 all-reduce。通信藏在 _CopyToModelParallelRegion 等 autograd.Function 里，对模型代码透明。

Q3：序列并行 SP 省了什么？

TP 在非层通信点 all-reduce 全 $[s,b,d]$ 激活，通信量 $\propto sbd$ 。SP 把序列 $s$ 也切到 TP 组各卡，层间用 reduce-scatter/all-gather，通信量降到 $sbd/T$ （all-reduce 的一半），且显存省 $T$ 倍激活。代价是 attention 内要再聚回全序列。长序列默认开。

Q4：1F1B 和 Interleaved 的气泡分别是多少？

朴素 PP 气泡 $\approx (P-1)/M$ 。1F1B（warmup+steady 1F1B+cooldown）显存上更优，气泡同量级。Interleaved 把每段再切 $V$ 个虚拟 chunk 轮流做，气泡降到 $(P-1)/(M\cdot V)$ ，代价是 P2P 通信次数 $\times V$ 、激活显存略增——带宽富裕时显著降气泡，带宽紧张时不划算。

Q5：Megatron 的 DDP 比 PyTorch 原生 DDP 强在哪？

三板斧：①param/grad 连续 buffer 把几十个碎片参数攒成大桶，一次大 all-reduce 替代多次小 all-reduce；②gradient accumulation fusion 把多 microbatch 的权重梯度融进 wgrad kernel 直接累加，一个 iter 只算一次；③bucket ready 触发 async all-reduce 与反向 overlap。再加 distributed optimizer 分片优化器状态/权重等价 ZeRO-1/2/3 但与 TP/PP 共生。

Q6：CP 切什么、权重复制吗、梯度怎么同步？

CP 切输入序列长度，降 attention $O(s^2)$ 激活显存，支持长上下文。CP 组内权重复制，所以反向时权重梯度要 all-reduce——CP 复用 DP 组做这个（_DATA_PARALLEL_GROUP_WITH_CP）。Attention 内靠 FlashAttention 的 ring/all-gather 变体在 CP 组传 KV，通信量 $O(s/T)$ 级、藏在前向算力里。

Q7：MoE 在 Megatron 里怎么并行？

expert 用 EP 切到各卡（每卡不同 expert），expert 内部矩阵还能 TP 切（expert_tensor_parallel_size）。dispatcher 用 all-to-all 把 token 按路由分发到对应 expert 所在卡，算完再 all-to-all 回收。_EXPERT_* 那一堆 group 就是给 MoE 通信用的。负载均衡靠辅助 loss。显存按总 expert 参数算、算力按激活 expert 参数算。

Q8：distributed optimizer 和 DeepSpeed ZeRO 是一回事吗？

思想一致——分片优化器状态/权重/梯度降低 DP 冗余显存。区别是 Megatron 的实现和 TP/PP/EP 深度耦合：rank 排布里 DP 在外层，分片粒度和 TP 切片、PP 分段、EP 分 expert 正交组合，param/grad buffer 既服务 all-reduce 也服务分片。dist_checkpointing 还支持分片保存和不同并行度间 reshard 加载。

Q9：通信/计算 overlap 在 Megatron 哪些地方体现？

TP 的 collective 藏在 autograd 与反向 overlap；PP 的 P2P 用 batch_isend_irecv 收发重叠；DDP 的 bucket ready 触发 async all-reduce 与反向 overlap；SP 用 reduce-scatter 省一半带宽。通用 timers 定位瓶颈。

十一、一张图收口

rank 排布:  r = tp + cp·T + ep·T·C + dp·T·C·E + pp·T·C·E·D   (默认 tp-cp-ep-dp-pp)
            快维度(TP)落机内 NVLink, 慢维度(PP)跨机 IB

TP:   Column 列切(前向不通信/反向reduce) + Row 行切(前向reduce/反向不通信)
      MLP=Column→act→Row, 每层通信压到两次 all-reduce; SP 用 reduce-scatter 省一半

PP:   1F1B(warmup→steady→cooldown), 气泡 (P-1)/M
      Interleaved V 虚拟段, 气泡 (P-1)/(MV), P2P 通信×V

DP:   param/grad 连续 buffer + gradient accum fusion + bucket ready async all-reduce
      + distributed optimizer(分片 opt state/weight = ZeRO-1/2/3, 与TP/PP共生)

CP:   切序列长度降 attention 激活, 权重复用 DP 组 reduce, FlashAttention ring 传 KV

EP:   切 expert, all-to-all dispatch, _EXPERT_* group

模型: transformer_layer = attention(TP+CP+FlashAttn/MLA) + mlp(Column/Row) + 残差/范数/fp8
      + MoE(dispatch all-to-all) + cuda_graphs + fusions

主线一句话：Megatron 把"切 Transformer"做成了五维正交并行——TP 切张量靠 Column/Row 配对把通信压到每层两次、PP 切层靠微批流水把气泡压到 $(P-1)/M$ 、DP 切 batch 靠 buffer+fusion+async overlap 把 all-reduce 藏进反向、CP 切序列降长上下文显存、EP 切 expert 靠 all-to-all——所有通信都尽力藏在 autograd 和算力背后，让上层只管写模型。把这五维和它们的通信点讲顺，Megatron 这关就稳了。

参考资料

Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019（TP）
Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, 2021（PP+DP，1F1B/Interleaved）
Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022（SP）
Liu et al., Ring Attention / Megatron CP, 2023
Megatron-Core 源码：https://github.com/NVIDIA/Megatron-LM （megatron/core/）
与本文显存计算法则、RDMA、x86 体系结构、操作系统接口篇交叉对照