L3.1 分布式训练范式与 3D 并行

三维坐标　layer: L3（训练层）｜level: Senior｜pillar: 训推框架

本文是 L3 训练层的开篇。当模型参数从 7B 膨胀到万亿，单卡装不下、单机训不动成为铁律。我们将沿着「一个 batch 的前向-反向-更新」这条数据流，拆开 DP / TP / PP 三个正交维度如何把一个巨型模型切碎、铺到上千张卡上，并理解每一种切法的通信代价才是选型的胜负手。

学习目标

前置知识：读过 L0.3（显存账本：知道参数/梯度/优化器状态/激活各占多少显存，会用「参数量 × N 字节」估显存）；理解 L0.2 的通信墙（带宽与延迟为何是分布式的硬约束）；读过 L1 互联拓扑（L1.2 NVLink / InfiniBand / RoCE，知道机内 NVLink 带宽远高于机间网络）。会写基础 PyTorch、知道前向-反向-更新三步即可，无需 NCCL/分布式经验。
学完产出：① 能在白板上画出 TP×PP×DP 的三层正交设备网格，并说清每条通信路径（TP All-Reduce / PP P2P / DP All-Reduce）的频率、通信量与带宽要求；② 能讲清 Megatron 的「列切 → 行切」配对为何能把一个 MLP 的前向通信压到只剩一次 All-Reduce；③ 能用气泡率公式 $\frac{p-1}{m+p-1}$ 解释「为什么 PP 必须切很多 micro-batch」，并说清 1F1B 相比 GPipe 省显存的根因；④ 能在给定卡数与互联条件下做出合理的 TP/PP/DP 选型，并指出把 TP 配成跨机为什么是致命错误；⑤ 亲手用 torchrun 跑通 DDP 与手写 TP，肉眼对照「梯度 All-Reduce」与「激活 All-Reduce」两种通信模式的差异。
阅读姿势：盯住一条主线——「分布式训练的所有切法，本质都是在『显存够不够』与『通信贵不贵』之间做权衡」。每多切一个维度，就是用一笔新的通信开销，换回一块被省下的显存或一段被缩短的时间；选型的胜负手，永远是把最贵的通信放进最快的链路里。

背景与现状

分布式训练（Distributed Training） 要解决的核心矛盾只有一句话：「模型 + 激活 + 优化器状态」的显存需求，远大于单张 GPU 的 HBM 容量；同时单卡算力跑完一个 epoch 的时间无法接受。解法是把计算与显存负担沿不同维度切开，分摊到多张卡，再用集合通信（Collective Communication）把切开的部分缝合回正确的数学结果。

业界把这些切分维度概括为 3D 并行（3D Parallelism），三个维度彼此正交、可叠加：

数据并行（Data Parallelism, DP）：每张卡持有完整模型副本，喂不同的数据分片，反向后用 All-Reduce 同步梯度。解决「训得快」。
张量并行（Tensor Parallelism, TP）：把单个算子的权重矩阵切到多卡（列切 / 行切），层内并行计算。解决「单层装不下」。
流水线并行（Pipeline Parallelism, PP）：把模型**按层（深度方向）**切成多段，每段放一组卡，像工厂流水线一样接力。解决「整个模型装不下」。

从产业演进看：

2019 前：DDP 一招鲜。模型小，单卡能装下，复制 + All-Reduce 即可。
2019–2021：Megatron-LM（NVIDIA）提出 TP，DeepSpeed（微软）提出 ZeRO + PP，10B–100B 模型成为可能。
2021 至今：万亿 MoE 与 128K 长上下文逼出 序列并行（SP）/ 上下文并行（CP）/ 专家并行（EP） 等新维度，3D 并行扩展成「5D 并行」。

业界信号：NVIDIA/Megatron-LM 与 microsoft/DeepSpeed 几乎是所有大模型团队的训练底座，Megatron-Core 更被 NeMo、LLaMA、GLM 等训练栈直接复用——会不会配 3D 并行，是大模型训练工程师的入门资格证。

原理与架构

理解 3D 并行最有效的方式，是先看清三个维度如何把同一个 GPU 集群「切成三层网格」，再逐一拆解每一维的通信代价。

2.1 3D 并行的设备网格切分

假设我们有 16 张 GPU，配置为 TP=2, PP=2, DP=4（注意 2×2×4=16）。它们被组织成三层正交的通信组：

三条通信路径，频率与代价天差地别：

TP All-Reduce（最频繁、最贵）：每个 Transformer 层的注意力和 MLP 各触发一次，要求 TP 组必须在同一台机器、走 NVLink，否则 PCIe / 网络带宽会成为瓶颈。
PP P2P（点对点传激活）：只在 stage 边界传一次激活张量，通信量小，可跨机。
DP All-Reduce（每个 step 一次）：反向结束后同步梯度，通信量 = 模型参数量，可被计算重叠（overlap）隐藏。

2.2 TP：列切与行切如何配对（Megatron 的核心技巧）

Megatron-LM 的精髓在于：把 MLP 的两个连续 Linear 设计成 「列切（Column Parallel）→ 行切（Row Parallel）」 的配对，使得中间只需一次 All-Reduce。

列切（Column Parallel）：权重 A 按列拆分 A = [A1, A2]，输出 [XA1, XA2] 各卡持有一片，前向无通信。
行切（Row Parallel）：权重 B 按行拆分，每卡算出部分和，前向用 All-Reduce 求和；反向的 All-Reduce 则由列切算子的 backward 触发（行切算子的 backward 是 identity）。
注意力同理：QKV 投影列切（每卡管几个 head），输出投影行切。一个 Transformer 层前向共 2 次 All-Reduce（Attention + MLP 各一），反向再 2 次。

2.3 PP 与 1F1B 调度：气泡（Bubble）的本质

流水线并行的天敌是气泡（Bubble）——流水线填充（warm-up）和排空（cool-down）阶段，部分 stage 在空转。朴素的 GPipe 调度气泡率高；Megatron/DeepSpeed 采用 1F1B（One-Forward-One-Backward） 调度压缩气泡。

气泡率公式：对于 $p$ 个 stage、 $m$ 个 micro-batch，气泡占比约为 $\frac{p-1}{m+p-1}$ 。

核心结论：micro-batch 数 $m$ 越大，气泡率越低。这就是为什么 PP 必须把一个 global batch 切成很多 micro-batch。
1F1B vs GPipe：GPipe 先做完所有前向再做反向，峰值激活显存 = m 份；1F1B 让前向反向交错进行，激活显存只需约 $p$ 份，在不增加气泡的前提下大幅省显存——这是 1F1B 被普遍采用的真正原因。
Interleaved 1F1B：Megatron 进一步把每个 stage 再切成多个「虚拟 stage」（virtual pipeline），气泡率再降到约 $\frac{1}{v}\cdot\frac{p-1}{m}$ ，代价是更多 P2P 通信。

2.4 三种并行的通信量对比

维度	通信原语	通信频率	单次通信量级	带宽要求	解决的问题
DP / DDP	All-Reduce（梯度）	每 step 1 次	∝ 模型参数量	中（可 overlap）	训得快
TP	All-Reduce（激活）	每层 2 次（×2 反向）	∝ batch×seq×hidden	极高，必须 NVLink	单层装不下
PP	P2P（激活）	每 micro-batch 边界	∝ batch×seq×hidden（仅边界）	低，可跨机	整模型装不下

需要强调的是：选型的黄金法则是——TP 限制在单机内（≤8，吃满 NVLink），PP 跨机（带宽要求低），DP 在最外层扩规模。把高频高带宽的 TP 放到跨机网络上，是新手最常见、也是最致命的配置错误。

动手实践：极简代码实操

实验目标：用 2 个进程（GPU 走 nccl，无 GPU 用 gloo 在 CPU 模拟）亲手跑通 DDP，再手写一个简化 TP，对照观测两者的通信原语与通信量差异，建立「不同并行 = 不同通信模式」的第一手直觉。

3.1 环境准备

# 推荐 Python 3.11；用 uv 或 venv 隔离
python3 -m venv .venv && source .venv/bin/activate
# CPU 路径（默认，gloo backend 即可模拟多进程）
pip install torch --index-url https://download.pytorch.org/whl/cpu
# 若有 NVIDIA GPU，改用官方 CUDA wheel（自动带 nccl）：
# pip install torch

3.2 代码 A：torchrun 跑通 DDP（梯度 All-Reduce）

# ddp_demo.py
import os
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP

def main():
    # torchrun 会注入 RANK / WORLD_SIZE / LOCAL_RANK 环境变量
    rank = int(os.environ["RANK"])
    world = int(os.environ["WORLD_SIZE"])

    # 有 GPU 用 nccl，否则用 gloo 在 CPU 上跑多进程
    ngpu = torch.cuda.device_count()
    use_nccl = torch.cuda.is_available() and ngpu >= world
    backend = "nccl" if use_nccl else "gloo"
    dist.init_process_group(backend=backend, rank=rank, world_size=world)

    dev = f"cuda:{rank}" if use_nccl else "cpu"
    if use_nccl:
        torch.cuda.set_device(rank)

    model = nn.Linear(8, 8).to(dev)
    ddp = DDP(model, device_ids=[rank] if backend == "nccl" else None)

    # 每个 rank 喂不同的数据分片
    x = torch.randn(4, 8, device=dev)
    loss = ddp(x).sum()
    loss.backward()  # ← 反向结束时 DDP 自动触发梯度 All-Reduce

    # 验证：All-Reduce 后所有 rank 的梯度完全一致
    g = next(ddp.parameters()).grad
    print(f"[rank {rank}] backend={backend} grad_mean={g.mean().item():.6f}")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

# 启动 2 个进程（NVIDIA GPU 与 CPU 路径命令完全相同，backend 由代码自动选）
torchrun --nproc_per_node=2 ddp_demo.py

观察点：两个 rank 打印的 grad_mean 完全相等——这正是 All-Reduce 把各卡梯度求和取平均的结果。DDP 在反向传播过程中分桶（bucket）触发 All-Reduce 并与反向计算重叠，这也是它比手动 reduce 快的原因。

3.3 代码 B：手写简化 TP（激活 All-Reduce）

下面手动实现 2.2 节的「列切 → 行切」MLP，看清 TP 的 All-Reduce 发生在激活上，而不是梯度上。

# tp_demo.py
import os, torch
import torch.distributed as dist

def main():
    rank = int(os.environ["RANK"]); world = int(os.environ["WORLD_SIZE"])
    ngpu = torch.cuda.device_count()
    use_nccl = torch.cuda.is_available() and ngpu >= world
    backend = "nccl" if use_nccl else "gloo"
    dist.init_process_group(backend=backend, rank=rank, world_size=world)
    dev = f"cuda:{rank}" if use_nccl else "cpu"
    if use_nccl: torch.cuda.set_device(rank)

    H, Hf = 8, 16                      # hidden / ffn 维度
    torch.manual_seed(0)
    x = torch.randn(4, H, device=dev)  # 输入在所有 TP-rank 上是“复制”的

    # 列切：第一层权重按列(Hf)切成 world 份，每个 rank 拿一片
    shard = Hf // world
    A_full = torch.randn(H, Hf, device=dev)          # 同 seed → 各 rank 一致
    A_local = A_full[:, rank*shard:(rank+1)*shard]   # 本 rank 的列分片
    # 行切：第二层权重按行(Hf)切，行分片与 A 的列分片对齐
    B_full = torch.randn(Hf, H, device=dev)
    B_local = B_full[rank*shard:(rank+1)*shard, :]

    h = torch.relu(x @ A_local)        # 列切输出天然分片，前向无通信
    y_partial = h @ B_local            # 行切：本 rank 只算出“部分和”
    dist.all_reduce(y_partial)         # ← TP 的核心：对激活做 All-Reduce 求和

    if rank == 0:
        # 对照：单机完整计算的参考结果
        y_ref = torch.relu(x @ A_full) @ B_full
        err = (y_partial - y_ref).abs().max().item()
        print(f"[TP] backend={backend} max_err_vs_single={err:.2e}")  # 应≈0

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

torchrun --nproc_per_node=2 tp_demo.py

对照结论：

DDP：All-Reduce 作用在梯度上，每 step 1 次，通信量 ∝ 参数量。
TP：All-Reduce 作用在激活上，每个 MLP/Attention 各 1 次（一层 2 次），通信量 ∝ batch×seq×hidden——层数越多、序列越长，TP 通信越密集，这就是 TP 必须绑定 NVLink 的根因。

踩坑预警 (Gotchas)

gloo 不支持的原语：gloo 后端不支持 reduce_scatter/部分 GPU-only 原语，CPU 模拟时只用 all_reduce/broadcast/all_gather，否则报 RuntimeError: gloo does not support。
端口被占用 / 多次启动残留：torchrun 默认 MASTER_PORT=29500，上次进程没退干净会报 Address already in use，换端口 torchrun --master_port=29501 ... 或 pkill -f torchrun。
TP 切分维度必须整除：Hf // world 不整除会导致各 rank 分片尺寸不一致、All-Reduce 形状对不上。生产代码必须 assert Hf % world == 0。
nccl 必须 set_device：GPU 路径若不 torch.cuda.set_device(rank)，多个进程会挤在 cuda:0 上，要么 OOM 要么 NCCL 卡死。
WORLD_SIZE 必须用 torchrun 启动：直接 python ddp_demo.py 跑会因为没有 RANK/WORLD_SIZE 环境变量而 KeyError，必须用 torchrun。

深入思考

下面三题每题先给题干，再用 <details> 折叠一份图文并茂的参考答案。建议先合上答案自己想 3 分钟，再展开对照。

思考题 1：64 卡训 70B 的 3D 配置选型

现有 64 张 A100（8 卡/机，机内 NVLink、机间 InfiniBand）要训一个 70B 模型。请给出你的 TP / PP / DP 三维配置，并结合 2.1 与 2.4 节的通信量表解释为什么 TP 不能跨机、PP 适合跨机、DP 放最外层；再说明把 TP 粗暴配成 16（跨机）会发生什么灾难。

展开参考答案（含 3D 配置到物理链路的映射图 + 算一遍）

结论：合理配置是 TP=8（绑死单机 8 卡 NVLink）、PP=2（跨 2 台机走 InfiniBand）、DP=4（在最外层扩规模），恰好 8×2×4=64。原则就一句话——把最频繁、最大带宽的 TP All-Reduce 关进机内 NVLink，把低频小通信的 PP P2P 放到机间网络，DP 梯度同步因可与计算重叠而对带宽最不敏感，丢到最外层。

用数量级算一遍「TP 配成 16 会怎样」（呼应 2.4 节通信量表）：

TP All-Reduce 的单次通信量 ∝ batch × seq × hidden，且每个 Transformer 层前向 2 次、反向 2 次。70B 模型有几十层，一个 step 内 TP 通信会触发上百次 All-Reduce。
机内 NVLink（A100 第三代）双向带宽约 600 GB/s；机间 InfiniBand（HDR 200Gb）单向约 25 GB/s——相差约 20~24 倍。
把 TP=16 意味着 TP 组跨越 2 台机，上百次 All-Reduce 里必有一半要走 InfiniBand。原本几十微秒的 NVLink 通信被拉长到毫秒级，乘以「每层 ×4、几十层、每 step」，通信彻底压倒计算，GPU 利用率（MFU）会从 ~50% 跌到个位数。
正确做法：TP 上限锚在「单机 NVLink 卡数」（A100/H100 通常 ≤8），需要再切就往 PP（跨机但低频）和 DP（最外层）上叠——这与 2.4 节「TP 限制在单机内、PP 跨机、DP 在最外层扩规模」的黄金法则完全一致。

思考题 2：气泡与显存的权衡

你的 PP 配置 p=8，global batch 切成 m=8 个 micro-batch 时气泡率约 47%。运维要求把气泡压到 10% 以下。你有两条路——增大 $m$ 还是 启用 Interleaved 1F1B？分别会付出什么代价（显存 / P2P 通信 / 实现复杂度）？

展开参考答案（含两条降气泡路径的决策图 + 算一遍）

结论：两条路都能压气泡，但代价方向不同——增大 $m$ 几乎零额外通信、实现简单，但受限于「global batch 不能无限大」和 micro-batch 太小算不满算力；Interleaved 1F1B 不增大 batch 就能降气泡，但把每个 stage 切成 $v$ 个虚拟 stage，P2P 通信次数翻 $v$ 倍、实现复杂度显著上升。实务上先尽量增大 $m$ ，触顶后再上 Interleaved。

用气泡率公式算一遍（呼应 2.3 节）：

路 A（增大 $m$ ）：气泡率 $\frac{p-1}{m+p-1}=\frac{7}{m+7}$ 。要 ≤10% 需 $m+7\ge 70$ ，即 $m\ge 63$ 。代价：global batch = m × micro-batch × DP， $m$ 涨 8 倍意味着要么 global batch 涨 8 倍（可能伤收敛、需重调学习率），要么把单个 micro-batch 缩到很小（算不满 Tensor Core，单步更慢）。但几乎不增加 P2P 通信，1F1B 的激活显存仍只需约 $p$ 份。
路 B（Interleaved 1F1B）：气泡率 ≈ $\frac{1}{v}\cdot\frac{p-1}{m}=\frac{1}{v}\cdot\frac{7}{8}$ 。取 $v=4$ 时约 ≈ 22%（仍未破 10%，需 $v\approx8$ 或同时增大 $m$ ），无需增大 batch。代价：每个物理 stage 持有 $v$ 段不连续的层，stage 边界的 P2P 通信次数 × $v$ ，对机间带宽更敏感；调度逻辑（Megatron 的 schedules.py）复杂得多，显存上每个虚拟 stage 都要留激活，峰值略升。
决策：若 global batch 还有余量、micro-batch 不至于太小 → 优先增大 $m$ （零成本）；若 batch 已触顶（再大伤收敛）→ 上 Interleaved，用 P2P 通信换气泡。

思考题 3：新维度（SP / CP / EP）的触发条件

当序列长度从 4K 拉到 128K，单纯靠 TP+PP+DP 为什么会撞墙（提示：激活显存与 attention 的 $O(n^2)$ ）？序列并行（SP）/ 上下文并行（CP） 各自切的是哪个张量维度，它们与 TP 是叠加还是替代关系？对于 MoE 模型，专家并行（EP） 又是把什么切到不同卡上的？

展开参考答案（含新维度切分对象示意图 + 对比表）

结论：长序列把「激活显存」这一项撑爆——它随 seq 线性涨、attention 中间量随 $O(n^2)$ 涨，而 DP/TP/PP 切的是参数和层、压根没切 seq 维度，于是撞墙。SP 沿 seq 维切 LayerNorm/Dropout 等逐元素激活（与 TP 叠加、几乎不增通信），CP 沿 seq 维切整个 attention（用 Ring Attention 跨卡传 K/V），EP 则把 MoE 的不同专家（Expert FFN）分到不同卡上（靠 All-to-All 路由 token）——三者都是新增的正交维度，是叠加而非替代。

用数量级算一遍「为什么会撞墙」：

单层 attention 的激活量级：QKV 与中间激活 ∝ batch × seq × hidden，随 seq 线性涨；而 attention 分数矩阵 QKᵀ 是 seq × seq，随 $O(\text{seq}^2)$ 暴涨。seq 从 4K → 128K（×32），线性项 ×32，平方项 ×1024。
DP 复制整模型（不省单卡激活）、TP 切 hidden 维（激活的 seq 维仍整份留在每张卡）、PP 切层（每个 stage 内单层激活仍是整条 seq）——三者都没碰 seq 维，所以激活显存随长序列爆炸时，原 3D 无解，必须引入沿 seq 维切的新维度。

新维度对比表：

维度	切的张量维度	通信原语	与 TP 的关系	解决的痛点
SP（序列并行）	seq 维上的逐元素激活（LayerNorm/Dropout，TP 未切的部分）	All-Gather / Reduce-Scatter（替换 TP 的 All-Reduce，总量不变）	叠加（Megatron SP 与 TP 配套使用）	进一步省 TP 没覆盖的激活显存
CP（上下文并行）	seq 维上的整个 attention 计算	Ring 形 P2P 传递 K/V 分块（Ring Attention）	叠加（正交新维度）	直接拆解 attention 的 $O(n^2)$ 长序列墙
EP（专家并行）	expert 维：不同 Expert FFN 放不同卡	All-to-All（按 router 把 token 发到对应专家）	叠加（MoE 专属，与 TP/DP 共存）	万亿 MoE 的专家权重单卡放不下

三者与原 TP/PP/DP 共同构成业界所说的「5D 甚至更高维并行」——每多一个维度，就是为某一类「原维度切不动」的资源（长序列激活、海量专家权重）新开一条正交的切分轴。

学习目标​

背景与现状​

原理与架构​

2.1 3D 并行的设备网格切分​

2.2 TP：列切与行切如何配对（Megatron 的核心技巧）​

2.3 PP 与 1F1B 调度：气泡（Bubble）的本质​

2.4 三种并行的通信量对比​

动手实践：极简代码实操​

3.1 环境准备​

3.2 代码 A：torchrun 跑通 DDP（梯度 All-Reduce）​

3.3 代码 B：手写简化 TP（激活 All-Reduce）​

踩坑预警 (Gotchas)​

深入思考​

思考题 1：64 卡训 70B 的 3D 配置选型​

思考题 2：气泡与显存的权衡​

思考题 3：新维度（SP / CP / EP）的触发条件​

延伸阅读​