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

<-->2026-08-01 11:58

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 执行都变成可控的分布式分片问题。
#mermaid-svg-qStgrsGPetwsZFg2{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-qStgrsGPetwsZFg2 .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-qStgrsGPetwsZFg2 .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-qStgrsGPetwsZFg2 .error-icon{fill:#552222;}#mermaid-svg-qStgrsGPetwsZFg2 .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-qStgrsGPetwsZFg2 .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-qStgrsGPetwsZFg2 .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-qStgrsGPetwsZFg2 .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-qStgrsGPetwsZFg2 .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-qStgrsGPetwsZFg2 .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-qStgrsGPetwsZFg2 .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-qStgrsGPetwsZFg2 .marker{fill:#333333;stroke:#333333;}#mermaid-svg-qStgrsGPetwsZFg2 .marker.cross{stroke:#333333;}#mermaid-svg-qStgrsGPetwsZFg2 svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-qStgrsGPetwsZFg2 p{margin:0;}#mermaid-svg-qStgrsGPetwsZFg2 .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-qStgrsGPetwsZFg2 .cluster-label text{fill:#333;}#mermaid-svg-qStgrsGPetwsZFg2 .cluster-label span{color:#333;}#mermaid-svg-qStgrsGPetwsZFg2 .cluster-label span p{background-color:transparent;}#mermaid-svg-qStgrsGPetwsZFg2 .label text,#mermaid-svg-qStgrsGPetwsZFg2 span{fill:#333;color:#333;}#mermaid-svg-qStgrsGPetwsZFg2 .node rect,#mermaid-svg-qStgrsGPetwsZFg2 .node circle,#mermaid-svg-qStgrsGPetwsZFg2 .node ellipse,#mermaid-svg-qStgrsGPetwsZFg2 .node polygon,#mermaid-svg-qStgrsGPetwsZFg2 .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-qStgrsGPetwsZFg2 .rough-node .label text,#mermaid-svg-qStgrsGPetwsZFg2 .node .label text,#mermaid-svg-qStgrsGPetwsZFg2 .image-shape .label,#mermaid-svg-qStgrsGPetwsZFg2 .icon-shape .label{text-anchor:middle;}#mermaid-svg-qStgrsGPetwsZFg2 .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-qStgrsGPetwsZFg2 .rough-node .label,#mermaid-svg-qStgrsGPetwsZFg2 .node .label,#mermaid-svg-qStgrsGPetwsZFg2 .image-shape .label,#mermaid-svg-qStgrsGPetwsZFg2 .icon-shape .label{text-align:center;}#mermaid-svg-qStgrsGPetwsZFg2 .node.clickable{cursor:pointer;}#mermaid-svg-qStgrsGPetwsZFg2 .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-qStgrsGPetwsZFg2 .arrowheadPath{fill:#333333;}#mermaid-svg-qStgrsGPetwsZFg2 .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-qStgrsGPetwsZFg2 .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-qStgrsGPetwsZFg2 .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-qStgrsGPetwsZFg2 .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-qStgrsGPetwsZFg2 .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-qStgrsGPetwsZFg2 .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-qStgrsGPetwsZFg2 .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-qStgrsGPetwsZFg2 .cluster text{fill:#333;}#mermaid-svg-qStgrsGPetwsZFg2 .cluster span{color:#333;}#mermaid-svg-qStgrsGPetwsZFg2 div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-qStgrsGPetwsZFg2 .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-qStgrsGPetwsZFg2 rect.text{fill:none;stroke-width:0;}#mermaid-svg-qStgrsGPetwsZFg2 .icon-shape,#mermaid-svg-qStgrsGPetwsZFg2 .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-qStgrsGPetwsZFg2 .icon-shape p,#mermaid-svg-qStgrsGPetwsZFg2 .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-qStgrsGPetwsZFg2 .icon-shape .label rect,#mermaid-svg-qStgrsGPetwsZFg2 .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-qStgrsGPetwsZFg2 .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-qStgrsGPetwsZFg2 .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-qStgrsGPetwsZFg2 :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Dataset / Global Batch
Micro Batches
Pipeline Parallel Stages
Tensor Parallel Layers
Sequence Parallel Activations
Expert Parallel MoE
Data Parallel Replicas
Distributed Optimizer
NCCL Collectives

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 梯度同步：
DP Rank 3 DP Rank 2 DP Rank 1 DP Rank 0 DP Rank 3 DP Rank 2 DP Rank 1 DP Rank 0 #mermaid-svg-1LLonEAqboDoVzP2{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-1LLonEAqboDoVzP2 .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-1LLonEAqboDoVzP2 .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-1LLonEAqboDoVzP2 .error-icon{fill:#552222;}#mermaid-svg-1LLonEAqboDoVzP2 .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-1LLonEAqboDoVzP2 .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-1LLonEAqboDoVzP2 .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-1LLonEAqboDoVzP2 .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-1LLonEAqboDoVzP2 .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-1LLonEAqboDoVzP2 .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-1LLonEAqboDoVzP2 .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-1LLonEAqboDoVzP2 .marker{fill:#333333;stroke:#333333;}#mermaid-svg-1LLonEAqboDoVzP2 .marker.cross{stroke:#333333;}#mermaid-svg-1LLonEAqboDoVzP2 svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-1LLonEAqboDoVzP2 p{margin:0;}#mermaid-svg-1LLonEAqboDoVzP2 .actor{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-1LLonEAqboDoVzP2 text.actor>tspan{fill:black;stroke:none;}#mermaid-svg-1LLonEAqboDoVzP2 .actor-line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-1LLonEAqboDoVzP2 .innerArc{stroke-width:1.5;stroke-dasharray:none;}#mermaid-svg-1LLonEAqboDoVzP2 .messageLine0{stroke-width:1.5;stroke-dasharray:none;stroke:#333;}#mermaid-svg-1LLonEAqboDoVzP2 .messageLine1{stroke-width:1.5;stroke-dasharray:2,2;stroke:#333;}#mermaid-svg-1LLonEAqboDoVzP2 #arrowhead path{fill:#333;stroke:#333;}#mermaid-svg-1LLonEAqboDoVzP2 .sequenceNumber{fill:white;}#mermaid-svg-1LLonEAqboDoVzP2 #sequencenumber{fill:#333;}#mermaid-svg-1LLonEAqboDoVzP2 #crosshead path{fill:#333;stroke:#333;}#mermaid-svg-1LLonEAqboDoVzP2 .messageText{fill:#333;stroke:none;}#mermaid-svg-1LLonEAqboDoVzP2 .labelBox{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-1LLonEAqboDoVzP2 .labelText,#mermaid-svg-1LLonEAqboDoVzP2 .labelText>tspan{fill:black;stroke:none;}#mermaid-svg-1LLonEAqboDoVzP2 .loopText,#mermaid-svg-1LLonEAqboDoVzP2 .loopText>tspan{fill:black;stroke:none;}#mermaid-svg-1LLonEAqboDoVzP2 .loopLine{stroke-width:2px;stroke-dasharray:2,2;stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-1LLonEAqboDoVzP2 .note{stroke:#aaaa33;fill:#fff5ad;}#mermaid-svg-1LLonEAqboDoVzP2 .noteText,#mermaid-svg-1LLonEAqboDoVzP2 .noteText>tspan{fill:black;stroke:none;}#mermaid-svg-1LLonEAqboDoVzP2 .activation0{fill:#f4f4f4;stroke:#666;}#mermaid-svg-1LLonEAqboDoVzP2 .activation1{fill:#f4f4f4;stroke:#666;}#mermaid-svg-1LLonEAqboDoVzP2 .activation2{fill:#f4f4f4;stroke:#666;}#mermaid-svg-1LLonEAqboDoVzP2 .actorPopupMenu{position:absolute;}#mermaid-svg-1LLonEAqboDoVzP2 .actorPopupMenuPanel{position:absolute;fill:#ECECFF;box-shadow:0px 8px 16px 0px rgba(0,0,0,0.2);filter:drop-shadow(3px 5px 2px rgb(0 0 0 / 0.4));}#mermaid-svg-1LLonEAqboDoVzP2 .actor-man line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-1LLonEAqboDoVzP2 .actor-man circle,#mermaid-svg-1LLonEAqboDoVzP2 line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;stroke-width:2px;}#mermaid-svg-1LLonEAqboDoVzP2 :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} NCCL AllReduce / ReduceScatter grad bucket grad bucket grad bucket grad bucket

通信开销：

  • AllReduce 通信量近似 2 × N − 1 N × G 2 \times \frac{N-1}{N} \times G 2×NN−1×G， G G G 为梯度大小， N N 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 = X W Y = XW Y=XW

设 X ∈ R S B × H X \in \mathbb{R}^{SB \times H} X∈RSB×H， W ∈ R H × 4 H W \in \mathbb{R}^{H \times 4H} W∈RH×4H 或 W ∈ R H × H W \in \mathbb{R}^{H \times H} W∈RH×H。矩阵可沿输出维或输入维切分。

Column Parallel Linear

Column Parallel 把权重按输出列切：

W = W 1 , W 2 , . . . , W p W = W_1, W_2, ..., W_p W=W1,W2,...,Wp

每个 TP rank 计算：

Y i = X W i Y_i = XW_i Yi=XWi

最后可选择 all-gather：

Y = Y 1 , Y 2 , . . . , Y p Y = Y_1, Y_2, ..., Y_p Y=Y1,Y2,...,Yp

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

Row Parallel Linear

Row Parallel 把输入维切：

W = W 1 W 2 . . . W p , X = X 1 , X 2 , . . . , X p W = \begin{bmatrix} W_1 \\ W_2 \\ ... \\ W_p \end{bmatrix}, \quad X = X_1, X_2, ..., X_p W= W1W2...Wp ,X=X1,X2,...,Xp

每个 rank 计算局部输出：

Y i = X i W i Y_i = X_i W_i Yi=XiWi

最终需要求和：

Y = ∑ i Y i Y = \sum_i Y_i Y=i∑Yi

源码对应 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 q k v Q, K, V = X W_{qkv} Q,K,V=XWqkv

每个 TP rank 拿到部分 attention heads：

n u m _ h e a d s _ p e r _ r a n k = n u m _ h e a d s T P num\_heads\_per\_rank = \frac{num\_heads}{TP} num_heads_per_rank=TPnum_heads

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

Attention 流程：
#mermaid-svg-YEOH7sL30XAeq3bu{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-YEOH7sL30XAeq3bu .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-YEOH7sL30XAeq3bu .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-YEOH7sL30XAeq3bu .error-icon{fill:#552222;}#mermaid-svg-YEOH7sL30XAeq3bu .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-YEOH7sL30XAeq3bu .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-YEOH7sL30XAeq3bu .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-YEOH7sL30XAeq3bu .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-YEOH7sL30XAeq3bu .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-YEOH7sL30XAeq3bu .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-YEOH7sL30XAeq3bu .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-YEOH7sL30XAeq3bu .marker{fill:#333333;stroke:#333333;}#mermaid-svg-YEOH7sL30XAeq3bu .marker.cross{stroke:#333333;}#mermaid-svg-YEOH7sL30XAeq3bu svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-YEOH7sL30XAeq3bu p{margin:0;}#mermaid-svg-YEOH7sL30XAeq3bu .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-YEOH7sL30XAeq3bu .cluster-label text{fill:#333;}#mermaid-svg-YEOH7sL30XAeq3bu .cluster-label span{color:#333;}#mermaid-svg-YEOH7sL30XAeq3bu .cluster-label span p{background-color:transparent;}#mermaid-svg-YEOH7sL30XAeq3bu .label text,#mermaid-svg-YEOH7sL30XAeq3bu span{fill:#333;color:#333;}#mermaid-svg-YEOH7sL30XAeq3bu .node rect,#mermaid-svg-YEOH7sL30XAeq3bu .node circle,#mermaid-svg-YEOH7sL30XAeq3bu .node ellipse,#mermaid-svg-YEOH7sL30XAeq3bu .node polygon,#mermaid-svg-YEOH7sL30XAeq3bu .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-YEOH7sL30XAeq3bu .rough-node .label text,#mermaid-svg-YEOH7sL30XAeq3bu .node .label text,#mermaid-svg-YEOH7sL30XAeq3bu .image-shape .label,#mermaid-svg-YEOH7sL30XAeq3bu .icon-shape .label{text-anchor:middle;}#mermaid-svg-YEOH7sL30XAeq3bu .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-YEOH7sL30XAeq3bu .rough-node .label,#mermaid-svg-YEOH7sL30XAeq3bu .node .label,#mermaid-svg-YEOH7sL30XAeq3bu .image-shape .label,#mermaid-svg-YEOH7sL30XAeq3bu .icon-shape .label{text-align:center;}#mermaid-svg-YEOH7sL30XAeq3bu .node.clickable{cursor:pointer;}#mermaid-svg-YEOH7sL30XAeq3bu .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-YEOH7sL30XAeq3bu .arrowheadPath{fill:#333333;}#mermaid-svg-YEOH7sL30XAeq3bu .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-YEOH7sL30XAeq3bu .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-YEOH7sL30XAeq3bu .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-YEOH7sL30XAeq3bu .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-YEOH7sL30XAeq3bu .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-YEOH7sL30XAeq3bu .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-YEOH7sL30XAeq3bu .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-YEOH7sL30XAeq3bu .cluster text{fill:#333;}#mermaid-svg-YEOH7sL30XAeq3bu .cluster span{color:#333;}#mermaid-svg-YEOH7sL30XAeq3bu div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-YEOH7sL30XAeq3bu .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-YEOH7sL30XAeq3bu rect.text{fill:none;stroke-width:0;}#mermaid-svg-YEOH7sL30XAeq3bu .icon-shape,#mermaid-svg-YEOH7sL30XAeq3bu .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-YEOH7sL30XAeq3bu .icon-shape p,#mermaid-svg-YEOH7sL30XAeq3bu .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-YEOH7sL30XAeq3bu .icon-shape .label rect,#mermaid-svg-YEOH7sL30XAeq3bu .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-YEOH7sL30XAeq3bu .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-YEOH7sL30XAeq3bu .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-YEOH7sL30XAeq3bu :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} hidden_states
ColumnParallel linear_qkv
split Q/K/V by local heads
RoPE / position embedding
FlashAttention / DotProductAttention kernel
RowParallel linear_proj
TP all-reduce or SP reduce-scatter

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 = ϕ ( X W 1 ) , Y = Z W 2 Z = \phi(XW_1), \quad Y = ZW_2 Z=ϕ(XW1),Y=ZW2

Megatron 中 linear_fc1 是 Column Parallel，输出为 4 H / T P 4H/TP 4H/TP 或 SwiGLU 下更宽；linear_fc2 是 Row Parallel，把局部 hidden 聚合回 H H 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：

text 复制代码 
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：

text 复制代码 
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 × B × H × L S \times B \times H \times L S×B×H×L 增长，SP 的核心是把 sequence 维也切开，让每个 TP rank 只保存 S / T P S/TP 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：
#mermaid-svg-AzTctr6N9I44DFpl{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-AzTctr6N9I44DFpl .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-AzTctr6N9I44DFpl .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-AzTctr6N9I44DFpl .error-icon{fill:#552222;}#mermaid-svg-AzTctr6N9I44DFpl .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-AzTctr6N9I44DFpl .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-AzTctr6N9I44DFpl .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-AzTctr6N9I44DFpl .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-AzTctr6N9I44DFpl .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-AzTctr6N9I44DFpl .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-AzTctr6N9I44DFpl .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-AzTctr6N9I44DFpl .marker{fill:#333333;stroke:#333333;}#mermaid-svg-AzTctr6N9I44DFpl .marker.cross{stroke:#333333;}#mermaid-svg-AzTctr6N9I44DFpl svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-AzTctr6N9I44DFpl p{margin:0;}#mermaid-svg-AzTctr6N9I44DFpl .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-AzTctr6N9I44DFpl .cluster-label text{fill:#333;}#mermaid-svg-AzTctr6N9I44DFpl .cluster-label span{color:#333;}#mermaid-svg-AzTctr6N9I44DFpl .cluster-label span p{background-color:transparent;}#mermaid-svg-AzTctr6N9I44DFpl .label text,#mermaid-svg-AzTctr6N9I44DFpl span{fill:#333;color:#333;}#mermaid-svg-AzTctr6N9I44DFpl .node rect,#mermaid-svg-AzTctr6N9I44DFpl .node circle,#mermaid-svg-AzTctr6N9I44DFpl .node ellipse,#mermaid-svg-AzTctr6N9I44DFpl .node polygon,#mermaid-svg-AzTctr6N9I44DFpl .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-AzTctr6N9I44DFpl .rough-node .label text,#mermaid-svg-AzTctr6N9I44DFpl .node .label text,#mermaid-svg-AzTctr6N9I44DFpl .image-shape .label,#mermaid-svg-AzTctr6N9I44DFpl .icon-shape .label{text-anchor:middle;}#mermaid-svg-AzTctr6N9I44DFpl .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-AzTctr6N9I44DFpl .rough-node .label,#mermaid-svg-AzTctr6N9I44DFpl .node .label,#mermaid-svg-AzTctr6N9I44DFpl .image-shape .label,#mermaid-svg-AzTctr6N9I44DFpl .icon-shape .label{text-align:center;}#mermaid-svg-AzTctr6N9I44DFpl .node.clickable{cursor:pointer;}#mermaid-svg-AzTctr6N9I44DFpl .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-AzTctr6N9I44DFpl .arrowheadPath{fill:#333333;}#mermaid-svg-AzTctr6N9I44DFpl .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-AzTctr6N9I44DFpl .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-AzTctr6N9I44DFpl .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-AzTctr6N9I44DFpl .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-AzTctr6N9I44DFpl .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-AzTctr6N9I44DFpl .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-AzTctr6N9I44DFpl .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-AzTctr6N9I44DFpl .cluster text{fill:#333;}#mermaid-svg-AzTctr6N9I44DFpl .cluster span{color:#333;}#mermaid-svg-AzTctr6N9I44DFpl div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-AzTctr6N9I44DFpl .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-AzTctr6N9I44DFpl rect.text{fill:none;stroke-width:0;}#mermaid-svg-AzTctr6N9I44DFpl .icon-shape,#mermaid-svg-AzTctr6N9I44DFpl .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-AzTctr6N9I44DFpl .icon-shape p,#mermaid-svg-AzTctr6N9I44DFpl .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-AzTctr6N9I44DFpl .icon-shape .label rect,#mermaid-svg-AzTctr6N9I44DFpl .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-AzTctr6N9I44DFpl .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-AzTctr6N9I44DFpl .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-AzTctr6N9I44DFpl :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Tokens
TopK Router
routing_map + probs
AllToAll / AllGather token dispatch
Local grouped experts GEMM
AllToAll combine
Output tokens restored

源码路径：

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 源码分析

当前重点目录

text 复制代码 
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 开始

典型启动：

bash 复制代码 
torchrun --nproc_per_node=8 pretrain_gpt.py \
  --tensor-model-parallel-size 4 \
  --pipeline-model-parallel-size 2 \
  ...

调用链：
#mermaid-svg-9bYUV2jUra4gRknR{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-9bYUV2jUra4gRknR .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-9bYUV2jUra4gRknR .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-9bYUV2jUra4gRknR .error-icon{fill:#552222;}#mermaid-svg-9bYUV2jUra4gRknR .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-9bYUV2jUra4gRknR .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-9bYUV2jUra4gRknR .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-9bYUV2jUra4gRknR .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-9bYUV2jUra4gRknR .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-9bYUV2jUra4gRknR .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-9bYUV2jUra4gRknR .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-9bYUV2jUra4gRknR .marker{fill:#333333;stroke:#333333;}#mermaid-svg-9bYUV2jUra4gRknR .marker.cross{stroke:#333333;}#mermaid-svg-9bYUV2jUra4gRknR svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-9bYUV2jUra4gRknR p{margin:0;}#mermaid-svg-9bYUV2jUra4gRknR .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-9bYUV2jUra4gRknR .cluster-label text{fill:#333;}#mermaid-svg-9bYUV2jUra4gRknR .cluster-label span{color:#333;}#mermaid-svg-9bYUV2jUra4gRknR .cluster-label span p{background-color:transparent;}#mermaid-svg-9bYUV2jUra4gRknR .label text,#mermaid-svg-9bYUV2jUra4gRknR span{fill:#333;color:#333;}#mermaid-svg-9bYUV2jUra4gRknR .node rect,#mermaid-svg-9bYUV2jUra4gRknR .node circle,#mermaid-svg-9bYUV2jUra4gRknR .node ellipse,#mermaid-svg-9bYUV2jUra4gRknR .node polygon,#mermaid-svg-9bYUV2jUra4gRknR .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-9bYUV2jUra4gRknR .rough-node .label text,#mermaid-svg-9bYUV2jUra4gRknR .node .label text,#mermaid-svg-9bYUV2jUra4gRknR .image-shape .label,#mermaid-svg-9bYUV2jUra4gRknR .icon-shape .label{text-anchor:middle;}#mermaid-svg-9bYUV2jUra4gRknR .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-9bYUV2jUra4gRknR .rough-node .label,#mermaid-svg-9bYUV2jUra4gRknR .node .label,#mermaid-svg-9bYUV2jUra4gRknR .image-shape .label,#mermaid-svg-9bYUV2jUra4gRknR .icon-shape .label{text-align:center;}#mermaid-svg-9bYUV2jUra4gRknR .node.clickable{cursor:pointer;}#mermaid-svg-9bYUV2jUra4gRknR .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-9bYUV2jUra4gRknR .arrowheadPath{fill:#333333;}#mermaid-svg-9bYUV2jUra4gRknR .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-9bYUV2jUra4gRknR .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-9bYUV2jUra4gRknR .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-9bYUV2jUra4gRknR .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-9bYUV2jUra4gRknR .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-9bYUV2jUra4gRknR .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-9bYUV2jUra4gRknR .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-9bYUV2jUra4gRknR .cluster text{fill:#333;}#mermaid-svg-9bYUV2jUra4gRknR .cluster span{color:#333;}#mermaid-svg-9bYUV2jUra4gRknR div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-9bYUV2jUra4gRknR .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-9bYUV2jUra4gRknR rect.text{fill:none;stroke-width:0;}#mermaid-svg-9bYUV2jUra4gRknR .icon-shape,#mermaid-svg-9bYUV2jUra4gRknR .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-9bYUV2jUra4gRknR .icon-shape p,#mermaid-svg-9bYUV2jUra4gRknR .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-9bYUV2jUra4gRknR .icon-shape .label rect,#mermaid-svg-9bYUV2jUra4gRknR .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-9bYUV2jUra4gRknR .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-9bYUV2jUra4gRknR .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-9bYUV2jUra4gRknR :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} torchrun creates ranks
pretrain_gpt.py
training.pretrain
initialize_megatron
torch.distributed.init_process_group
parallel_state.initialize_model_parallel
TP/PP/DP/CP/EP groups
get_model / model_provider
DDP/FSDP wrapping
get_megatron_optimizer
train loop

关键源码：

Training Loop：一次 iteration

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

一次 iteration：
#mermaid-svg-dCmUezggkrZLsvJ2{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-dCmUezggkrZLsvJ2 .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-dCmUezggkrZLsvJ2 .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-dCmUezggkrZLsvJ2 .error-icon{fill:#552222;}#mermaid-svg-dCmUezggkrZLsvJ2 .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-dCmUezggkrZLsvJ2 .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-dCmUezggkrZLsvJ2 .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-dCmUezggkrZLsvJ2 .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-dCmUezggkrZLsvJ2 .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-dCmUezggkrZLsvJ2 .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-dCmUezggkrZLsvJ2 .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-dCmUezggkrZLsvJ2 .marker{fill:#333333;stroke:#333333;}#mermaid-svg-dCmUezggkrZLsvJ2 .marker.cross{stroke:#333333;}#mermaid-svg-dCmUezggkrZLsvJ2 svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-dCmUezggkrZLsvJ2 p{margin:0;}#mermaid-svg-dCmUezggkrZLsvJ2 .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-dCmUezggkrZLsvJ2 .cluster-label text{fill:#333;}#mermaid-svg-dCmUezggkrZLsvJ2 .cluster-label span{color:#333;}#mermaid-svg-dCmUezggkrZLsvJ2 .cluster-label span p{background-color:transparent;}#mermaid-svg-dCmUezggkrZLsvJ2 .label text,#mermaid-svg-dCmUezggkrZLsvJ2 span{fill:#333;color:#333;}#mermaid-svg-dCmUezggkrZLsvJ2 .node rect,#mermaid-svg-dCmUezggkrZLsvJ2 .node circle,#mermaid-svg-dCmUezggkrZLsvJ2 .node ellipse,#mermaid-svg-dCmUezggkrZLsvJ2 .node polygon,#mermaid-svg-dCmUezggkrZLsvJ2 .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-dCmUezggkrZLsvJ2 .rough-node .label text,#mermaid-svg-dCmUezggkrZLsvJ2 .node .label text,#mermaid-svg-dCmUezggkrZLsvJ2 .image-shape .label,#mermaid-svg-dCmUezggkrZLsvJ2 .icon-shape .label{text-anchor:middle;}#mermaid-svg-dCmUezggkrZLsvJ2 .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-dCmUezggkrZLsvJ2 .rough-node .label,#mermaid-svg-dCmUezggkrZLsvJ2 .node .label,#mermaid-svg-dCmUezggkrZLsvJ2 .image-shape .label,#mermaid-svg-dCmUezggkrZLsvJ2 .icon-shape .label{text-align:center;}#mermaid-svg-dCmUezggkrZLsvJ2 .node.clickable{cursor:pointer;}#mermaid-svg-dCmUezggkrZLsvJ2 .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-dCmUezggkrZLsvJ2 .arrowheadPath{fill:#333333;}#mermaid-svg-dCmUezggkrZLsvJ2 .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-dCmUezggkrZLsvJ2 .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-dCmUezggkrZLsvJ2 .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-dCmUezggkrZLsvJ2 .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-dCmUezggkrZLsvJ2 .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-dCmUezggkrZLsvJ2 .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-dCmUezggkrZLsvJ2 .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-dCmUezggkrZLsvJ2 .cluster text{fill:#333;}#mermaid-svg-dCmUezggkrZLsvJ2 .cluster span{color:#333;}#mermaid-svg-dCmUezggkrZLsvJ2 div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-dCmUezggkrZLsvJ2 .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-dCmUezggkrZLsvJ2 rect.text{fill:none;stroke-width:0;}#mermaid-svg-dCmUezggkrZLsvJ2 .icon-shape,#mermaid-svg-dCmUezggkrZLsvJ2 .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-dCmUezggkrZLsvJ2 .icon-shape p,#mermaid-svg-dCmUezggkrZLsvJ2 .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-dCmUezggkrZLsvJ2 .icon-shape .label rect,#mermaid-svg-dCmUezggkrZLsvJ2 .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-dCmUezggkrZLsvJ2 .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-dCmUezggkrZLsvJ2 .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-dCmUezggkrZLsvJ2 :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} zero grad buffer + optimizer.zero_grad
forward_backward_func
forward microbatches
loss calculation
backward microbatches
finalize_model_grads / DDP grad sync
optimizer.step
LR scheduler + logging

对应源码行为：

  • 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 调用关系

#mermaid-svg-wktWQoeZfAA0pi7U{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-wktWQoeZfAA0pi7U .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-wktWQoeZfAA0pi7U .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-wktWQoeZfAA0pi7U .error-icon{fill:#552222;}#mermaid-svg-wktWQoeZfAA0pi7U .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-wktWQoeZfAA0pi7U .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-wktWQoeZfAA0pi7U .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-wktWQoeZfAA0pi7U .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-wktWQoeZfAA0pi7U .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-wktWQoeZfAA0pi7U .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-wktWQoeZfAA0pi7U .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-wktWQoeZfAA0pi7U .marker{fill:#333333;stroke:#333333;}#mermaid-svg-wktWQoeZfAA0pi7U .marker.cross{stroke:#333333;}#mermaid-svg-wktWQoeZfAA0pi7U svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-wktWQoeZfAA0pi7U p{margin:0;}#mermaid-svg-wktWQoeZfAA0pi7U .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-wktWQoeZfAA0pi7U .cluster-label text{fill:#333;}#mermaid-svg-wktWQoeZfAA0pi7U .cluster-label span{color:#333;}#mermaid-svg-wktWQoeZfAA0pi7U .cluster-label span p{background-color:transparent;}#mermaid-svg-wktWQoeZfAA0pi7U .label text,#mermaid-svg-wktWQoeZfAA0pi7U span{fill:#333;color:#333;}#mermaid-svg-wktWQoeZfAA0pi7U .node rect,#mermaid-svg-wktWQoeZfAA0pi7U .node circle,#mermaid-svg-wktWQoeZfAA0pi7U .node ellipse,#mermaid-svg-wktWQoeZfAA0pi7U .node polygon,#mermaid-svg-wktWQoeZfAA0pi7U .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-wktWQoeZfAA0pi7U .rough-node .label text,#mermaid-svg-wktWQoeZfAA0pi7U .node .label text,#mermaid-svg-wktWQoeZfAA0pi7U .image-shape .label,#mermaid-svg-wktWQoeZfAA0pi7U .icon-shape .label{text-anchor:middle;}#mermaid-svg-wktWQoeZfAA0pi7U .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-wktWQoeZfAA0pi7U .rough-node .label,#mermaid-svg-wktWQoeZfAA0pi7U .node .label,#mermaid-svg-wktWQoeZfAA0pi7U .image-shape .label,#mermaid-svg-wktWQoeZfAA0pi7U .icon-shape .label{text-align:center;}#mermaid-svg-wktWQoeZfAA0pi7U .node.clickable{cursor:pointer;}#mermaid-svg-wktWQoeZfAA0pi7U .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-wktWQoeZfAA0pi7U .arrowheadPath{fill:#333333;}#mermaid-svg-wktWQoeZfAA0pi7U .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-wktWQoeZfAA0pi7U .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-wktWQoeZfAA0pi7U .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-wktWQoeZfAA0pi7U .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-wktWQoeZfAA0pi7U .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-wktWQoeZfAA0pi7U .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-wktWQoeZfAA0pi7U .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-wktWQoeZfAA0pi7U .cluster text{fill:#333;}#mermaid-svg-wktWQoeZfAA0pi7U .cluster span{color:#333;}#mermaid-svg-wktWQoeZfAA0pi7U div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-wktWQoeZfAA0pi7U .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-wktWQoeZfAA0pi7U rect.text{fill:none;stroke-width:0;}#mermaid-svg-wktWQoeZfAA0pi7U .icon-shape,#mermaid-svg-wktWQoeZfAA0pi7U .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-wktWQoeZfAA0pi7U .icon-shape p,#mermaid-svg-wktWQoeZfAA0pi7U .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-wktWQoeZfAA0pi7U .icon-shape .label rect,#mermaid-svg-wktWQoeZfAA0pi7U .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-wktWQoeZfAA0pi7U .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-wktWQoeZfAA0pi7U .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-wktWQoeZfAA0pi7U :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} TransformerBlock
TransformerLayer
input_layernorm
SelfAttention
linear_qkv ColumnParallel
core_attention / FlashAttention
linear_proj RowParallel
bias dropout add
pre_mlp_layernorm
MLP or MoELayer
linear_fc1 ColumnParallel
GELU/SwiGLU fused
linear_fc2 RowParallel
bias dropout add

关键源码：

哪些地方发生 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 i t e r ≈ T c o m p u t e + T e x p o s e d _ c o m m T_{iter} \approx T_{compute} + T_{exposed\_comm} Titer≈Tcompute+Texposed_comm

Overlap 不是减少通信总量，而是减少 exposed communication：

T e x p o s e d _ c o m m = max ⁡ ( 0 , T c o m m − T o v e r l a p p e d ) T_{exposed\comm} = \max(0, T{comm} - T_{overlapped}) Texposed_comm=max(0,Tcomm−Toverlapped)

性能指标

  • 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：
    E f f i c i e n c y ( N ) = T h r o u g h p u t ( N ) N × T h r o u g h p u t ( 1 ) Efficiency(N) = \frac{Throughput(N)}{N \times Throughput(1)} Efficiency(N)=N×Throughput(1)Throughput(N)

第五部分：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。
  1. 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。
  1. 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。
  1. 支持国产 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 等生态替代。
  1. GPU 软件栈替代路线：

#mermaid-svg-M0pk5m1jVVqLcf7n{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-M0pk5m1jVVqLcf7n .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-M0pk5m1jVVqLcf7n .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-M0pk5m1jVVqLcf7n .error-icon{fill:#552222;}#mermaid-svg-M0pk5m1jVVqLcf7n .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-M0pk5m1jVVqLcf7n .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-M0pk5m1jVVqLcf7n .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-M0pk5m1jVVqLcf7n .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-M0pk5m1jVVqLcf7n .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-M0pk5m1jVVqLcf7n .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-M0pk5m1jVVqLcf7n .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-M0pk5m1jVVqLcf7n .marker{fill:#333333;stroke:#333333;}#mermaid-svg-M0pk5m1jVVqLcf7n .marker.cross{stroke:#333333;}#mermaid-svg-M0pk5m1jVVqLcf7n svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-M0pk5m1jVVqLcf7n p{margin:0;}#mermaid-svg-M0pk5m1jVVqLcf7n .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-M0pk5m1jVVqLcf7n .cluster-label text{fill:#333;}#mermaid-svg-M0pk5m1jVVqLcf7n .cluster-label span{color:#333;}#mermaid-svg-M0pk5m1jVVqLcf7n .cluster-label span p{background-color:transparent;}#mermaid-svg-M0pk5m1jVVqLcf7n .label text,#mermaid-svg-M0pk5m1jVVqLcf7n span{fill:#333;color:#333;}#mermaid-svg-M0pk5m1jVVqLcf7n .node rect,#mermaid-svg-M0pk5m1jVVqLcf7n .node circle,#mermaid-svg-M0pk5m1jVVqLcf7n .node ellipse,#mermaid-svg-M0pk5m1jVVqLcf7n .node polygon,#mermaid-svg-M0pk5m1jVVqLcf7n .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-M0pk5m1jVVqLcf7n .rough-node .label text,#mermaid-svg-M0pk5m1jVVqLcf7n .node .label text,#mermaid-svg-M0pk5m1jVVqLcf7n .image-shape .label,#mermaid-svg-M0pk5m1jVVqLcf7n .icon-shape .label{text-anchor:middle;}#mermaid-svg-M0pk5m1jVVqLcf7n .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-M0pk5m1jVVqLcf7n .rough-node .label,#mermaid-svg-M0pk5m1jVVqLcf7n .node .label,#mermaid-svg-M0pk5m1jVVqLcf7n .image-shape .label,#mermaid-svg-M0pk5m1jVVqLcf7n .icon-shape .label{text-align:center;}#mermaid-svg-M0pk5m1jVVqLcf7n .node.clickable{cursor:pointer;}#mermaid-svg-M0pk5m1jVVqLcf7n .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-M0pk5m1jVVqLcf7n .arrowheadPath{fill:#333333;}#mermaid-svg-M0pk5m1jVVqLcf7n .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-M0pk5m1jVVqLcf7n .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-M0pk5m1jVVqLcf7n .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-M0pk5m1jVVqLcf7n .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-M0pk5m1jVVqLcf7n .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-M0pk5m1jVVqLcf7n .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-M0pk5m1jVVqLcf7n .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-M0pk5m1jVVqLcf7n .cluster text{fill:#333;}#mermaid-svg-M0pk5m1jVVqLcf7n .cluster span{color:#333;}#mermaid-svg-M0pk5m1jVVqLcf7n div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-M0pk5m1jVVqLcf7n .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-M0pk5m1jVVqLcf7n rect.text{fill:none;stroke-width:0;}#mermaid-svg-M0pk5m1jVVqLcf7n .icon-shape,#mermaid-svg-M0pk5m1jVVqLcf7n .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-M0pk5m1jVVqLcf7n .icon-shape p,#mermaid-svg-M0pk5m1jVVqLcf7n .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-M0pk5m1jVVqLcf7n .icon-shape .label rect,#mermaid-svg-M0pk5m1jVVqLcf7n .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-M0pk5m1jVVqLcf7n .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-M0pk5m1jVVqLcf7n .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-M0pk5m1jVVqLcf7n :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Megatron-LM
PyTorch Device Backend
torch.distributed Backend
Fused Kernels
Collective Library: NCCL Alternative
Compiler: Triton/CUTLASS Alternative
Runtime Streams/Events/Graphs
Driver + Runtime API
NVLink/PCIe/RDMA Equivalent

第六部分：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 干扰。

代码/运行方式：

bash 复制代码 
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 通信。

运行方式：

bash 复制代码 
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。

运行方式：

bash 复制代码 
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 与性能的直觉。

运行方式：

bash 复制代码 
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。

运行方式：

bash 复制代码 
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？

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

相关推荐
aqi001 小时前
15天学会AI应用开发（十八）使用LangGraph实现精确记忆功能
人工智能·python·大模型·ai编程·ai应用
Michaelliu_dev1 小时前
RoPE通俗讲解
人工智能·llm·位置编码·多模态大模型·rope·旋转位置编码·mllm
呆呆敲代码的小Y1 小时前
5 分钟上手 OpenMontage：把 AI 编程助手变成视频工作室
人工智能·aigc·音视频·ai视频生成·claude code·openmontage
ShallWeL1 小时前
【机器学习】（30）—— 嵌入空间
人工智能·神经网络·机器学习·embedding
MindUp1 小时前
AI辅助PPT生成工具实测：百度文库等平台在技术学习场景下的内容生成与排版能力对比
人工智能·百度·powerpoint
甲维斯1 小时前
DeepSeekFlash前端依旧拉垮，而且变慢了很多！
前端·人工智能
console.log('npc')1 小时前
OpenClaw 使用教程：开源 AI Agent 编排框架完全指南
人工智能·microsoft·ai编程·openclaw
独行侠影a1 小时前
AI绘画与音乐：生成式艺术是创作还是抄袭？
人工智能
数字融合1 小时前
黎阳之光一数字孪生多能源可视化系统管理技术
人工智能·自动化·virtualenv
热门推荐
01如何新建文件夹？ 电脑新建文件夹的4种方法02GitHub 镜像站点032026年7月AI圈大地震：GPT-5.6被政府限制、Claude入驻Slack、Anthropic自研芯片04国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）05AI 编程 IDE 全景解析 2026：Agent 全面接管开发链路06AI科技热点日报 | 2026年07月01日07微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新082026上半年大模型全景技术解读：推理融合、Agent 爆发与多模态统一092026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？102026 年 AI 大模型 & AI 编程工具实战全总结