多 GPU 训练的四种并行方式 —— 拆分什么,通信什么

• 引言 —— 八张 80GB 显卡为何仍放不下一个 7B 模型
• 内存账本 —— 每参数 16 字节,以及比它更大的激活开销
  • 状态内存:每参数 16 字节
  • 真正把显卡撑爆的是激活
• 数据并行与 ZeRO —— 切分账本里的哪一项
• 张量并行 —— 拆开矩阵,每层都要通信
• 流水线并行 —— 通信最便宜,气泡最昂贵
• 序列并行与上下文并行 —— 第四个维度
• 选择组合 —— 按规模和 GPU 数量读表
• 常见失败模式
• 结语 —— 选择并行方式就是在花带宽预算

引言 —— 八张 80GB 显卡为何仍放不下一个 7B 模型

"7B 模型用 bf16 只要 14GB,一张 80GB 显卡绰绰有余"——这种算法只在推理场景下成立。真要训练同一个模型,一张卡连起步都做不到,哪怕接上八张,依配置不同依然可能爆显存。

原因在于,训练除了参数之外还必须额外背负三块内存:梯度、优化器状态,以及激活。这四块内容里拆分什么、怎么拆,就是并行策略的全部;而拆分之后 GPU 之间要交换什么、交换多少,就是性能的全部。

这篇文章按拆分什么和通信多少这两条轴,梳理四种并行方式。所有数字都以可以自行验算的形式给出。推理侧的内存计算已经在《LLM 推理显存计算法》中讲过,这里只看训练。

内存账本 —— 每参数 16 字节,以及比它更大的激活开销

状态内存:每参数 16 字节

在混合精度训练中使用 Adam 系列优化器时,单个参数占用的字节数可以拆分如下。

这 16 字节和ZeRO 论文给出的账本完全一致。7B 模型对应 112GB,70B 对应 1,120GB。一张 80GB 显卡连 7B 都训练不了的结论,正是由此得出。

换一个优化器,这本账就会跟着变。用 SGD with momentum 的话,fp32 状态就只剩主权重和动量两项,每参数降到 12 字节;用 8 位 Adam 或 Adafactor 还能进一步降低。不过在预训练规模上更换优化器,考验的是收敛而不是内存,所以低精度优化器主要用在微调阶段。

真正把显卡撑爆的是激活

上面这本账里没有算激活。可预训练里制造 OOM 的元凶大多正是激活。Korthikanti 等人的Reducing Activation Recomputation in Large Transformer Models指出,在没有并行化的情况下,单个 Transformer 层的激活内存,以序列长度 s、批大小 b、隐藏维度 h、注意力头数 a 表示,为 s·b·h·(34 + 5as/h) 字节。

def act_gb(s: int, b: int, h: int, a: int, layers: int, mode: str) -> float:
    sbh = s * b * h
    per_layer = {
        "none": sbh * (34 + 5 * a * s / h),   # 全部保存
        "selective": sbh * 34,                # 只重新计算注意力矩阵
        "full": sbh * 2,                      # 只保存层输入,全部重新计算
    }[mode]
    return per_layer * layers / 1e9


s, b, h, a, L = 4096, 1, 4096, 32, 32   # 7B 级配置,单个样本
for mode in ("none", "selective", "full"):
    print(f"{mode:10} {act_gb(s, b, h, a, L, mode):7.2f} GB / 样本")
# none        104.15 GB / 样本
# selective    18.25 GB / 样本
# full          1.07 GB / 样本

单个样本就要 104GB。哪怕微批次只设成 4,也会变成 416GB。全部重新计算的话会降到 1.07GB,代价是多跑一次前向传播,计算量大约增加三分之一。

请注意这里和 34 一起出现的 5as/h 项。这一项与序列长度的平方成正比。把 s 从 4K 增加到 32K,34 那一项会变成 8 倍,而 5as/h 那一项会变成 64 倍。长上下文训练突然变难,正是从这里开始;需要上下文并行,原因也在这里。

数据并行与 ZeRO —— 切分账本里的哪一项

最简单的并行方式 DDP,是让每张 GPU 把上面这 16 字节原样复制一份。只切分批次,把梯度做 all-reduce 之后,各自执行相同的更新。实现最简单,但完全不省内存。

ZeRO 会分阶段把这份复制去掉。

def per_gpu_state_gb(params_b: float, world: int, stage: int) -> float:
    """按 Adam 混合精度计算,单张 GPU 承担的状态内存(GB,以 10^9 为基准)。
    stage 0 = DDP,1 = 优化器切分,2 = +梯度,3 = +参数
    不包含激活。"""
    p = params_b * 1e9
    param, grad, opt = 2 * p, 2 * p, 12 * p
    if stage >= 1:
        opt /= world
    if stage >= 2:
        grad /= world
    if stage >= 3:
        param /= world
    return (param + grad + opt) / 1e9


for stage in (0, 1, 2, 3):
    row = [round(per_gpu_state_gb(n, 8, stage), 1) for n in (7, 13, 70)]
    print(f"stage {stage}  world=8  7B/13B/70B -> {row} GB")
# stage 0  world=8  7B/13B/70B -> [112.0, 208.0, 1120.0] GB
# stage 1  world=8  7B/13B/70B -> [38.5, 71.5, 385.0] GB
# stage 2  world=8  7B/13B/70B -> [26.2, 48.8, 262.5] GB
# stage 3  world=8  7B/13B/70B -> [14.0, 26.0, 140.0] GB

读法是这样的。第 3 阶段把整整 16 字节都按 GPU 数量切分,所以 GPU 越多,内存占用就能无限缩小。相反,第 1 阶段把参数和梯度这 4 字节留了下来、没有切分,所以不管加多少 GPU,7B 模型都不会低于 28GB。如果加了更多 GPU 内存却没降下来,通常是 ZeRO 阶段选低了。

在 PyTorch 里,FSDP 实际上就对应 ZeRO 第 3 阶段。截至 2026 年 8 月 2 日,PyTorch 官方教程明确写着"FSDP1 is deprecated",并推荐使用基于 torch.distributed.fully_shard 的 FSDP2。教程里没写清楚是从哪个版本开始废弃的,这一点最好对照你所用 PyTorch 版本的发行说明再确认一遍。这里确认过的 PyTorch 稳定版本号是 2.13.0,发布日期是 2026 年 7 月 8 日。

# FSDP2 的最小形式——不包裹模块,而是原地分片
import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

mesh = init_device_mesh("cuda", (torch.distributed.get_world_size(),), mesh_dim_names=("dp",))
policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32)

# 不能只作用于根模块,子模块也要应用,通信单位才会被切得够细
for block in model.layers:
    fully_shard(block, mesh=mesh, mp_policy=policy)
fully_shard(model, mesh=mesh, mp_policy=policy)

如果只写最后一行、漏掉循环,整个模型就会变成一个通信单位,在前向传播开始时一次性把全部参数收集起来。内存节省效果消失,通信和计算的重叠也没有了。"开了 FSDP,为什么内存没降"这类问题,很大一部分都出在这个错误上。

张量并行 —— 拆开矩阵,每层都要通信

张量并行是把一次矩阵乘法拆给多张 GPU 分担计算。Megatron-LM 论文里的经典做法,是把 MLP 的第一个矩阵按列切分、第二个矩阵按行切分,中间不需要通信就能直接拼接,只在模块末尾做一次 all-reduce。注意力模块也用同样的方式按头切分。

结果是,每个 Transformer 层会产生前向 2 次、反向 2 次的 all-reduce。通信的对象不是参数,而是激活张量,大小为 s·b·h。

def tp_bytes_per_step(s, b, h, layers, tp, dtype_bytes=2):
    """张量并行中单张 GPU 每步搬运的字节数(环形 all-reduce 近似)。"""
    tensor = s * b * h * dtype_bytes
    ring = 2 * (tp - 1) / tp          # 环形 all-reduce 中每 GPU 的搬运量系数
    per_layer = 4 * tensor * ring     # 前向 2 + 反向 2
    return per_layer * layers


print(tp_bytes_per_step(4096, 1, 4096, 32, tp=8) / 1e9, "GB/step")
# 7.516192768 GB/step

每一步都要搬运 7.5GB。用 NVLink 4 的 H100 SXM 官方标称每 GPU 双向 900GB/s,这次通信大约 10 毫秒就能完成。用一个 400Gb/s 的 InfiniBand 端口做同样的通信,也就是 50GB/s,就会变成 150 毫秒。张量并行组一旦跨越节点边界,训练速度就会掉一个数量级,原因就是这道算术题。

上面这些都是厂商标称的数值,实测有效带宽会比这更低。更准确的做法是在集群上直接跑一遍 nccl-tests 的 all_reduce_perf 拿到真实数字。实务上的规则很简单:张量并行度不应超过单节点内的 GPU 数量。如果是 8-GPU 节点,TP 就应该在 8 以内。

流水线并行 —— 通信最便宜,气泡最昂贵

流水线并行把层切成若干阶段,分别放到不同的 GPU 上。只有在阶段边界才交换激活,所以通信量少得多。把 32 层切成 4 个阶段,边界只有 3 处,每个边界上只有一个 s·b·h 大小的张量做点对点传输。所以流水线是可以跨节点边界的。

代价是气泡。第一个阶段处理第一个微批次的时候,其余阶段都在空转;等到最后一个微批次的反向传播结束之前,前面的阶段又会再空转一轮。

def bubble_ratio(stages: int, micro_batches: int) -> float:
    return (stages - 1) / micro_batches

for m in (4, 8, 16, 32, 64):
    print(f"micro_batches={m:3}  8 阶段气泡 {bubble_ratio(8, m):.1%}")
# micro_batches=  4  8 阶段气泡 175.0%
# micro_batches=  8  8 阶段气泡 87.5%
# micro_batches= 16  8 阶段气泡 43.8%
# micro_batches= 32  8 阶段气泡 21.9%
# micro_batches= 64  8 阶段气泡 10.9%

要把气泡压到 10% 左右,微批次数量至少要达到阶段数的 8 倍。但微批次一多,激活内存就会增加,全局批次大小也会变大,进而影响学习率调度。这三个值彼此牵连,正是流水线调优困难的原因。

目前已经有多种用来压缩气泡的调度方式,比如交错调度(interleaved schedule),或者 DeepSeek-V3 的 DualPipe。DualPipe 通过把前向和反向的计算-通信区间相互重叠来压缩流水线气泡,具体说明见DeepSeek-V3 技术报告。在 PyTorch 一侧,torch.distributed.pipelining 作为标准 API 提供。

序列并行与上下文并行 —— 第四个维度

前面三条轴都是按参数和批次来切分的。还有一条轴,是按序列长度来切分的。这块术语常被混用,这里做一下区分。

• 序列并行:Megatron 系列里与张量并行搭配使用的技术。它把张量并行没能切分、只能重复保留的 LayerNorm 和 dropout 部分的激活,按序列轴切分,从而消除这部分的重复存储。它在与张量并行相同的通信组内运作。
• 上下文并行:把注意力计算本身按序列轴切分。每张 GPU 只持有序列的一部分,把键和值以环形传递,逐步累加出部分注意力结果。环形注意力(ring attention)系列都属于这一类。DeepSpeed 的 Ulysses 用 all-to-all 通信实现了类似的目标,是另一种实现方式。

该不该用上下文并行,判断标准很明确:如果已经增加了序列长度、把微批次降到 1、激活也全部改成重新计算,却依然 OOM,那就是打开上下文并行的时候了。Llama 3 405B 训练之所以用 4D 并行(TP、CP、PP、DP),也是因为 128K 上下文扩展阶段的缘故。

如果训练的是 MoE 模型,还会再加上一条专家并行。做法是把专家分散放到不同 GPU 上,再用 all-to-all 把 token 路由过去。因为通信模式是 all-to-all 而不是 all-reduce,对网络拓扑要敏感得多;一旦负载均衡崩掉,就会变成只有部分 GPU 在干活,其余的都在等。

选择组合 —— 按规模和 GPU 数量读表

实务中要决定的不是"哪种并行方式更好",而是"在这个规模下该按什么顺序叠加"。通行的做法是把通信代价高的轴放在最内层:TP 放在节点内,CP 紧随其后,PP 放在节点之间,DP 放在最外层。

这张表只是起点,不是标准答案。实际情况下,目标全局批次大小会强烈制约组合方式。全局批次等于微批次 × 梯度累积 × 数据并行度,而 DP 的度数固定为 GPU 总数除以 TP、PP、CP 的度数。把 TP 从 8 改成 4,DP 就会翻倍,全局批次也跟着翻倍,学习率就得重新调。改变并行化配置不是单纯的基础设施工作,而是一次超参数变更。

常见失败模式

这里只列现场反复出现的那些情况。

1. TP 组跨了节点。把 16-GPU 的 TP 横跨在两台 8-GPU 节点上,速度会直接崩掉。用 nvidia-smi topo -m 确认实际链路,并检查启动器把各个 rank 按什么顺序分配到节点上。
2. FSDP 只套在根模块上。这是前面提到过的错误。必须以子模块为单位包裹,通信才能和计算重叠。
3. 开了激活检查点还是 OOM。不参与重新计算的嵌入层、输出层,以及 logits 张量都会留下来。在词表很大的模型里,logits 的大小是 s·b·V,光是它自己就能达到几十 GB。需要拆成分块计算损失来处理。
4. 靠减少微批次来避开 OOM,却放着气泡不管。在 PP 环境下减少微批次确实能解决内存问题,但气泡会变大,吞吐量会腰斩。这两者是同一根杠杆的两端。
5. 梯度累积和归一化对不上。用累积步数去除损失的位置一旦搞错,实际生效的学习率就会变。每次并行配置发生变化,都应该跑一次短程训练,确认损失曲线和之前的配置是否重合。
6. 靠调大 NCCL 超时时间来掩盖问题。出现超时,通常是某个 rank 挂了,或者是有一个慢节点拖后腿。应该先用 NCCL_DEBUG=INFO 看清楚是哪个 rank 卡住了。

结语 —— 选择并行方式就是在花带宽预算

这四种并行方式彼此并不是替代关系,而是各自消耗不同资源的工具。ZeRO 和 FSDP 是拿内存换带宽,张量并行是拿内存换节点内带宽,流水线是拿内存换空闲时间,上下文并行是拿序列轴上的内存换环形通信。

所以在选定配置之前,真正该做的不是去读框架文档,而是画出你这个集群的带宽地图:节点内 GPU 间的实测带宽、节点间的实测带宽,以及这两者的比值。掌握这三个数字之后,上面那张表基本就能自动决定了。下一篇文章会按谱系整理《训练框架技术栈》,讲这些决策如何真正落实到代码里。