Megatron 架构与 6D 并行完整指南¶

更新日期：2026-04-14

一、Megatron-LM 是什么¶

Megatron-LM 是 NVIDIA 开发的大规模 Transformer 训练框架，是当前工业界训练百亿到万亿参数模型的事实标准。核心论文：Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019)。

Megatron 的核心价值不是模型架构，而是分布式并行策略。它支持 6 种并行维度，可以任意组合。

二、6D 并行全景¶

总 GPU 数 = TP × PP × DP × EP × CP（SP 通常跟随 TP，不独立计数）

三、张量并行（TP）深入¶

3.1 原理¶

TP 将单个矩阵乘法切分到多个 GPU 上并行计算。参考论文 Megatron-LM (2019)。

# MLP 的列并行 + 行并行
# 原始: Y = GeLU(X @ A) @ B
# A: [d, 4d], B: [4d, d]

# TP=2 时:
# GPU 0: A0 = A[:, :2d],  B0 = B[:2d, :]
# GPU 1: A1 = A[:, 2d:],  B1 = B[2d:, :]

# 前向:
# GPU 0: Y0 = GeLU(X @ A0) @ B0   →  AllReduce(Y0, Y1) → Y
# GPU 1: Y1 = GeLU(X @ A1) @ B1   →

# 关键: A 列切分, B 行切分
# GeLU 在列切分的输出上可以独立计算 (因为 GeLU 是逐元素的)
# 最终需要一次 AllReduce 将两个 GPU 的部分和合并

3.2 Attention 的 TP¶

# Attention 的 TP 更自然: 按头切分
# 如果 32 个头, TP=8: 每个 GPU 负责 4 个头
# Q, K, V 投影按列切分, 输出投影按行切分
# 每个 GPU 独立计算自己负责的 4 个头的完整注意力
# 最终 AllReduce 输出投影的结果

3.3 TP 的通信量¶

TP 的 AllReduce 必须在每个 micro-step 的前向和反向中各执行。频率极高，所以必须放在 NVLink 互联的节点内（带宽 900 GB/s）。跨节点做 TP 会成为严重瓶颈。

四、流水线并行（PP）深入¶

4.1 原理¶

PP 将模型的层分成几组，每组放在不同的 GPU（或不同节点）上。参考论文 GPipe (Huang et al., 2019) 和 PipeDream (Narayanan et al., 2019)。

# 60 层模型, PP=4:
# Stage 0 (GPU group 0): Layer 0-14
# Stage 1 (GPU group 1): Layer 15-29
# Stage 2 (GPU group 2): Layer 30-44
# Stage 3 (GPU group 3): Layer 45-59

# 数据流: Stage 0 → Stage 1 → Stage 2 → Stage 3 (前向)
#          Stage 3 → Stage 2 → Stage 1 → Stage 0 (反向)
# 通信: 只需在 stage 边界传递激活值 (P2P Send/Recv)

4.2 PP 调度策略¶

气泡（Bubble）是 PP 最大的效率损失。在气泡时间内，部分 GPU 空闲等待上游/下游。

# 1F1B 调度 (最常用)
# 假设 PP=4, 8 个 micro-batch
#
# 时间线 (F=前向, B=反向, 数字=micro-batch编号):
# Stage 0: F0 F1 F2 F3 B0 F4 B1 F5 B2 F6 B3 F7 B4 B5 B6 B7
# Stage 1:    F0 F1 F2 F3 B0 F4 B1 F5 B2 F6 B3 F7 B4 B5 B6 B7
# Stage 2:       F0 F1 F2 F3 B0 F4 B1 F5 B2 F6 B3 F7 B4 B5 B6 B7
# Stage 3:          F0 F1 F2 B0 F3 B1 F4 B2 F5 B3 F6 B4 F7 B5 B6 B7
#
# Warmup 阶段: 前 PP-1 个 micro-batch 只做前向 (填充流水线)
# Steady 阶段: 交替 1 次前向 + 1 次反向
# Cooldown 阶段: 最后 PP-1 个 micro-batch 只做反向 (排空流水线)

4.3 Zero Bubble PP¶

Zero Bubble Pipeline Parallelism 通过将反向传播拆分为"权重梯度"和"输入梯度"两部分，并重新排列调度顺序，将气泡率降至接近 0。参考 Zero Bubble Pipeline Parallelism (Qi et al., 2024)。

五、数据并行（DP）与 ZeRO¶

5.1 标准 DP¶

每个 GPU 持有完整模型副本，处理不同数据，梯度 AllReduce 后同步更新。

5.2 ZeRO 优化¶

在 Megatron 中通常只用 ZeRO-1（分片优化器状态）。因为 TP+PP 已经大幅减少了内存，不需要 ZeRO-3 的激进分片。ZeRO-3 的通信量太大，会拖慢训练。

参考文献¶

[1] Shoeybi et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. 论文
[2] Narayanan et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. 2021. 论文
[3] Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. 2022. 论文

参考文献¶

[1] Shoeybi et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. 论文
[2] Narayanan et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. 2021. 论文
[3] Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. 2022. 论文

Megatron 架构与 6D 并行完整指南¶

更新日期：2026-04-14

一、Megatron-LM 是什么¶

Megatron-LM 是 NVIDIA 开发的大规模 Transformer 训练框架，是当前工业界训练百亿到万亿参数模型的事实标准。核心论文：Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019)。

Megatron 的核心价值不是模型架构，而是分布式并行策略。它支持 6 种并行维度，可以任意组合。

二、6D 并行全景¶

总 GPU 数 = TP × PP × DP × EP × CP（SP 通常跟随 TP，不独立计数）

三、张量并行（TP）深入¶

3.1 原理¶

TP 将单个矩阵乘法切分到多个 GPU 上并行计算。参考论文 Megatron-LM (2019)。

MLP 由两个矩阵乘法组成：\(Y = \text{GeLU}(X \cdot A) \cdot B\)，其中 \(A \in \mathbb{R}^{d \times 4d}\)，\(B \in \mathbb{R}^{4d \times d}\)。

TP=2 时切分方式：前向计算：每个 GPU 独立计算 \(Y_i = \text{GeLU}(X \cdot A_i) \cdot B_i\)，最终 AllReduce 合并 \(Y = Y_0 + Y_1\)。

为什么 A 列切、B 行切？因为 GeLU 是逐元素函数——列切分后每个 GPU 的 GeLU 可以独立计算。B 行切分保证最终只需一次 AllReduce（不是 AllGather）。

3.2 Attention 的 TP¶

Attention 的 TP 更自然——按头切分。32 头 TP=8 → 每 GPU 负责 4 个头。Q、K、V 投影按列切分，每个 GPU 独立计算自己负责的 4 个头的完整注意力（不需要跨 GPU 通信）。最终输出投影按行切分 + AllReduce 合并结果。

3.3 TP 的通信量¶

TP 的 AllReduce 必须在每个 micro-step 的前向和反向中各执行。频率极高，所以必须放在 NVLink 互联的节点内（带宽 900 GB/s）。跨节点做 TP 会成为严重瓶颈。

四、流水线并行（PP）深入¶

4.1 原理¶

PP 将模型的层分成几组，每组放在不同的 GPU（或不同节点）上。参考论文 GPipe (Huang et al., 2019) 和 PipeDream (Narayanan et al., 2019)。

60 层模型 PP=4 时的切分：数据流：Stage 0 → 1 → 2 → 3（前向），Stage 3 → 2 → 1 → 0（反向）。通信只在 stage 边界发生，传递激活值（P2P Send/Recv）。

4.2 PP 调度策略¶

气泡（Bubble）是 PP 最大的效率损失。在气泡时间内，部分 GPU 空闲等待上游/下游。 1F1B 调度（最常用）。假设 PP=4，8 个 micro-batch：

Warmup 阶段：前 PP-1=3 个 micro-batch 只做前向（填充流水线）。此时 Stage 3 空闲 → 气泡。

Steady 阶段：交替 1 次前向 + 1 次反向。所有 Stage 都在工作。

Cooldown 阶段：最后 PP-1=3 个 micro-batch 只做反向（排空流水线）。此时 Stage 0 空闲 → 气泡。

气泡率 = Warmup + Cooldown 时间 / 总时间 = \((PP-1) / \text{total\_microbatches}\)。PP=4, 8 micro-batch → 气泡率 = ⅜ = 37.5%。增加 micro-batch 数 → 气泡率下降。

4.3 Zero Bubble PP¶

Zero Bubble Pipeline Parallelism 通过将反向传播拆分为"权重梯度"和"输入梯度"两部分，并重新排列调度顺序，将气泡率降至接近 0。参考 Zero Bubble Pipeline Parallelism (Qi et al., 2024)。

五、数据并行（DP）与 ZeRO¶

5.1 标准 DP¶

每个 GPU 持有完整模型副本，处理不同数据，梯度 AllReduce 后同步更新。

5.2 ZeRO 优化¶

在 Megatron 中通常只用 ZeRO-1（分片优化器状态）。因为 TP+PP 已经大幅减少了内存，不需要 ZeRO-3 的激进分片。ZeRO-3 的通信量太大，会拖慢训练。

参考文献¶

[1] Shoeybi et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. 论文
[2] Narayanan et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. 2021. 论文
[3] Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. 2022. 论文

参考文献¶

[1] Shoeybi et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. 论文
[2] Narayanan et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. 2021. 论文
[3] Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. 2022. 论文

↑ 上级 · C1. Megatron 架构与 6D 并行完整指南