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

更新日期：2026-04-14

---

六、上下文并行（CP）深入

CP 将长序列切分到多个 GPU 上，每个 GPU 只处理序列的一部分。参考 Ring Attention (Liu et al., 2023)。

# 场景: 训练 128K 上下文, 但单 GPU 的 KV Cache 放不下 128K
# CP=4: 每个 GPU 处理 32K

# 问题: Attention 需要每个 Q 与所有 K/V 交互
# 解决: Ring Attention — 按环形拓扑传递 KV 块

def ring_attention(Q_local, KV_local, cp_group):
    # Q_local: [B, S/CP, H, d]  — 本地的 Q
    # KV_local: [B, S/CP, H, d] — 本地的 KV (初始)

    O = zeros_like(Q_local)
    m = full(..., -inf)  # running max for online softmax
    l = zeros(...)       # running sum

    kv_current = KV_local
    for step in range(cp_size):
        # 计算本地 Q 与当前 KV 块的注意力 (用 Flash Attention kernel)
        O_step, m_step, l_step = flash_attention(Q_local, kv_current)

        # Online softmax 合并 (无需全量 softmax)
        m_new = max(m, m_step)
        alpha = exp(m - m_new)
        beta = exp(m_step - m_new)
        l = alpha  l + beta  l_step
        O = alpha  O + beta  O_step
        m = m_new

        # Ring 通信: 把 KV 传给下一个 GPU, 从上一个 GPU 接收
        kv_current = ring_send_recv(kv_current, cp_group)
        # 关键: 通信和计算可以重叠!

    return O / l

6.1 Dynamic CP (2026.01 新特性)

Megatron Core 2026.01 引入了 Dynamic Context Parallelism，可以根据 batch 内实际序列长度动态调整 CP 大小，避免短序列浪费。最高 1.48x 加速。

---

七、序列并行（SP）

SP 是 TP 的补充。TP 切分了注意力和 FFN 的计算，但 LayerNorm 和 Dropout 是在完整序列上操作的 — 这部分 TP 帮不上忙。SP 将这些操作的序列维度也切分。参考 Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022)。

SP 的核心思路：

• 不用 SP 时： LayerNorm 需要完整 [B, S, D] 激活 → 显存占用大

• 用 SP 后： LayerNorm 在 [B, S/TP, D] 上计算 → 激活内存减少 TP 倍

实现方式： 将 TP 中的 AllReduce 替换为 ReduceScatter + AllGather。

操作	作用	原理
ReduceScatter	对 TP 组内各 GPU 的部分结果先做 Reduce（求和），再 Scatter 到每张卡上，使每个 GPU 只保留 1/TP 的聚合结果，激活内存立即降为原来的 1/TP	在 LayerNorm / Dropout 前，各 GPU 已持有完整但冗余的中间激活；ReduceScatter 同时完成"聚合 + 切分"，一次通信替代 AllReduce，且输出天然是切片形式，后续 LayerNorm 可直接在局部切片上计算
AllGather	在进入下一个需要完整张量的算子（如 Attention QKV 线性层）前，将各 GPU 持有的 1/TP 切片拼回完整张量	TP 的矩阵乘法要求输入是完整的隐藏维度；AllGather 通信量与 ReduceScatter 相同（每卡发送 data_size/TP），因此 SP 总通信量 = 原 AllReduce，不多不少，但中间状态的显存节省了 TP 倍
>

总通信量不变，但激活内存减少 TP 倍。

---

八、6D 并行的实际配置

8.1 配置实例

模型	总 GPU	TP	PP	DP	EP	CP	节点数	MFU	为什么选这组配置
LLaMA-2 70B	64×A100（80 GB SXM）	8：70B 模型 BF16 权重约 140 GB，TP=8 让每卡只持有 ~17.5 GB 权重，NVLink 带宽足够	1：70B 用 TP=8 已可放入单节点，无需流水线切分	8：8 节点全部用于数据并行，最大化吞吐	-：Dense 模型无专家	1：标准 4K 上下文，无需序列切分	8	~48%：Dense 模型通信简单，TP 内 NVLink 高效，DP AllReduce 可与计算重叠
LLaMA-3 405B	16384×H100（80 GB SXM）	8：405B 权重 ~810 GB，TP=8 让每卡持有 ~101 GB，配合激活重计算刚好放入 80 GB 显存	16：TP=8 仍不够切分 405B 的全部层，PP=16 将 126 层分布到 16 个流水线 stage	128：16384/(8×16)=128 路数据并行，全局 batch 可达百万 token	-：Dense 模型	1：标准上下文长度	2048	~43%：PP=16 引入较大气泡开销（~6%），跨节点 PP 通信延迟进一步降低 MFU
DeepSeek-V3 (671B MoE)	2048×H800（80 GB）	1（注意力头用 TP=8 切分）：MoE FFN 本身已被 EP 切分，无需对 FFN 做 TP；仅注意力部分用 TP=8 节省 KV Cache	16：671B 总参数量巨大，PP=16 按层切分，每 stage ~42B 参数	2：EP=64 已占满大量 GPU，DP=2 保持最小数据并行度以维持梯度统计稳定性	64：256 个专家 / EP=64 = 每卡 4 个专家，All-to-All 通信限制在 rack 内高带宽网络	1：标准 4K-8K 上下文	256	~40%：MoE 的 All-to-All dispatch 开销 + 专家负载不均衡是 MFU 下降主因
128K 长上下文训练	256×H100（80 GB SXM）	8：长序列的 QKV 矩阵极大，TP=8 按注意力头切分以降低单卡激活内存	4：长序列 + 大模型需要 PP 进一步切分层，PP=4 将气泡控制在 ~12%	2：256/(8×4×1×4)=2，仅 2 路 DP，长序列场景 batch 本身很小	-：Dense 模型	4：128K / CP=4 = 每卡处理 32K token，KV Cache 从 128K 降到 32K，显存占用降 4 倍；Ring Attention 在 NVLink + IB 下可将通信与 Flash Attention 重叠	32	~35%：CP 引入 Ring 通信开销 + 长序列导致计算 / 通信比下降
视频 DiT 训练	512×H100（80 GB SXM）	4：DiT 隐藏维度较小（~3072），TP=4 已足够切分；TP=8 会因通信开销超过计算收益	4：视频 DiT 通常 48-64 层，PP=4 每 stage 12-16 层	8：512/(4×4×1×4)=8 路数据并行，视频训练单样本很大，DP=8 平衡吞吐	-：非 MoE 架构	4：视频帧展开后序列极长（如 16 帧×256 patch=4096+），CP=4 按时间帧切分，每卡处理 4 帧的 patch 序列	64	~32%：视频数据 I/O 瓶颈 + DiT 的 cross-attention 通信密集，整体利用率偏低

8.2 配置设计原则

黄金规则（按优先级排列）： 1. TP 在节点内：TP 通信频率最高，必须用 NVLink（900 GB/s）。TP ≤ GPUs_per_node。

1. PP 在相邻节点：PP 是 P2P 通信，延迟敏感但带宽要求低。

2. EP 在同一 rack：All-to-All 通信对带宽敏感。

3. DP 在最外层：梯度 AllReduce 可以与计算重叠，带宽可聚合。

4. CP 按需：只有训练超长序列（>32K）才需要。

8.3 配置计算示例

def design_parallel_config(
    model_params_B,      # 模型参数 (十亿)
    is_moe,              # 是否 MoE
    n_experts,           # MoE 专家数
    seq_len,             # 训练序列长度
    total_gpus,          # 总 GPU 数
    gpus_per_node,       # 每节点 GPU 数 (通常 8)
    gpu_mem_gb,          # 单卡显存 (如 80)
):
    # Step 1: TP — 确保权重能放进节点内
    # 模型权重 (BF16) ≈ 2 × model_params_B GB
    # 加上优化器状态 ≈ 16 × model_params_B GB (AdamW)
    # 加上梯度 ≈ 2 × model_params_B GB
    # 加上激活 ≈ 变化大
    tp = 1
    mem_per_gpu = (16 + 2 + 2) * model_params_B  # 粗估
    while mem_per_gpu / tp > gpu_mem_gb * 0.8:
        tp *= 2
    tp = min(tp, gpus_per_node)

    # Step 2: PP — 如果 TP 还不够
    pp = 1
    while mem_per_gpu / (tp  pp) > gpu_mem_gb  0.8:
        pp *= 2

    # Step 3: EP (MoE only)
    ep = 1
    if is_moe:
        # 每个 GPU 上的专家数 = n_experts / EP
        # 希望每 GPU 上 4-16 个专家
        ep = max(1, n_experts // 16)

    # Step 4: CP (长序列)
    cp = 1
    if seq_len > 32768:
        cp = min(seq_len // 16384, 8)

    # Step 5: DP
    dp = total_gpus // (tp  pp  ep * cp)

    return {'TP': tp, 'PP': pp, 'DP': dp, 'EP': ep, 'CP': cp}

---

九、通信重叠

通信重叠（Communication-Computation Overlap）是提高 MFU 的关键技术。目标：在 GPU 计算时，同时在后台执行通信。

# Megatron 中开启通信重叠
args = {
    '--overlap-grad-reduce': True,          # DP 梯度通信与计算重叠
    '--overlap-param-gather': True,         # ZeRO 参数收集与计算重叠
    '--use-distributed-optimizer': True,    # 启用分布式优化器 (ZeRO-1)
}

---

十、Megatron 数据格式

10.1 .bin / .idx 格式

数据预处理命令：

python tools/preprocess_data.py \
    --input data.jsonl \
    --output-prefix my_data \
    --tokenizer-model tokenizer.model \
    --workers 64 \
    --append-eod

训练时指定多数据集混合与权重：

--data-path "0.7 web_data 0.15 code_data 0.15 math_data"

10.2 图片不存在 .bin 中

多模态数据：.bin 只存文本 token。图片通过路径引用或预编码的特征文件加载。详见 [B3 多模态数据工程]。

---

参考文献

• [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. 论文

• [4] Huang et al. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. 2019. 论文

• [5] Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. 2020. 论文

• [6] Liu et al. Ring Attention with Blockwise Transformers. 2023. 论文

• [7] Qi et al. Zero Bubble Pipeline Parallelism. 2024. 论文

• [8] Megatron-LM GitHub

• [9] Megatron Parallelisms Guide

---

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