万卡训练的张力:2026 年 3D 并行与 ZeRO 组合的工程真相
约 22 分钟6529 字2 次阅读
一句话摘要
当模型规模从 7B 跃迁到 671B(DeepSeek-V3)再到万亿参数,单卡显存已不再是约束,真正的瓶颈是多维并行策略的张力——张量并行的通信粒度、流水线并行的 bubble、ZeRO 的内存节省与通信开销之间的帕累托前沿。本文从工程一线视角拆解 FSDP、DeepSpeed ZeRO、Megatron-LM TP/PP 在 2026 年生产级训练基础设施中的角色定位与组合范式。
一、引言:为什么大模型训练从单卡到万卡的工程跨越如此痛苦
2024 年训练 7B 模型只需要 16 张 H100 + DDP(Distributed Data Parallel)就够了;到 2026 年,主流的前沿模型——DeepSeek-V3(671B MoE)、Llama-3.1-405B、Qwen3-235B——已经把训练基础设施推到了 万卡 GPU 集群 + 3D 并行(TP+PP+DP)+ ZeRO-1 + 序列并行 + 专家并行的复合架构上。
训练的"痛苦"主要来自三个相互冲突的物理约束:
- 显存墙:单个 H100 有 80GB 显存,但训练一个 671B 的稠密模型光是参数 + 优化器状态(Adam 的一阶/二阶动量)就需要约 2.7TB 显存——是单卡容量的 33 倍。
- 通信墙:跨卡同步梯度、跨机通信 all-reduce 的延迟是毫秒级,而 GPU 计算的吞吐是 TFLOPS 级——任何不合理的并行切分都会让 GPU 闲置等通信。
- 气泡墙(Pipeline Bubble):流水线并行把模型按层切到不同卡上,前向/反向必须等待整个 pipeline 填满才能开始计算——空泡时间可以吃掉 30%+ 的算力。
这三条墙同时存在,使得分布式训练不是一个"把单卡代码包一层 DDP"就能解决的事,而是一整套并行策略的组合优化问题。下面我们沿着"内存数学 → DP/ZeRO → TP → PP → 3D 组合 → 生产实战"的脉络,把这套基础设施拆开看。
二、内存墙:训练 LLM 的根本约束
在讨论任何并行策略之前,必须先建立单卡显存占用的精确数学模型。设模型参数量为 (单位:字节),训练时单卡需要存储的内存由以下几部分组成:
其中各项的含义是:
- 参数 + 梯度:fp16 训练下,参数占 字节,反向传播时累积梯度同样占 字节。
- Adam 优化器状态:Adam 需要为每个参数维护一阶动量 和二阶动量 ,通常以 fp32 存储各占 字节,加上原始 fp32 master copy 一份 字节——合计 。
- 激活值(Activation):前向时保存的中间张量用于反向传播,约 字节( 是序列长度, 是层数,假设 hidden_size 与 成正比)。
代入 Llama-3.1-405B 的实际数字(,,),仅 Adam 状态一项就需要约 4.86TB——这是单张 H100(80GB)的 60 倍。
这就是为什么 DDP(朴素的 DP)跑不了大模型:DDP 每张卡都要保留完整的参数 + 梯度 + 优化器状态, 张卡就重复存 份,没有任何节省。所有大模型训练框架的核心目标,就是把这三块内存按某种策略切分到多张卡上。
三、数据并行(DP)与 ZeRO-1/2/3 的演进
ZeRO(Zero Redundancy Optimizer)是微软 DeepSpeed 团队 2019 年提出的分片式数据并行,核心思想是:DDP 不需要每张卡都保留完整的优化器状态,只需要各卡保留 1/N 就行。它分三个阶段渐进地把切分粒度推得更细:
ZeRO Stage 1:仅切分优化器状态 → 每卡内存 ≈ (12P/N + 2P + 2P)
ZeRO Stage 2:切分优化器状态 + 梯度 → 每卡内存 ≈ (12P/N + 2P/N + 2P)
ZeRO Stage 3:切分优化器状态 + 梯度 + 参数 → 每卡内存 ≈ 16P/N
ZeRO-1 的工程取舍:仅切分优化器状态时,参数和梯度仍需在每卡保留完整副本——好处是 forward/backward 计算阶段无需额外的 all-gather,通信量与朴 DDP 几乎相同;代价是参数/梯度没有省内存,所以 7B 左右的模型用 ZeRO-1 就够了,但 70B+ 必须升到 ZeRO-2 或 -3。
ZeRO-3 的代价:参数也被切分后,forward 阶段每层都要做一次 all-gather 把完整参数拉到本地算完再释放——通信量从 O(P) 涨到 O(P × layers_per_step),但内存从 O(P) 降到 O(P/N)。这是一个内存换通信的帕累托权衡,DeepSeek-V3 训练报告(2024-12)就明确提到 ZeRO-3 + 双份副本(参数 + 计算副本)的模式让 671B 模型得以在 2048 张 H800 上启动。
工程经验:当 (单卡容量),就必须上 ZeRO-2 或 -3,否则 OOM 直接挂。
四、FSDP vs DDP:完全分片 vs 参数服务器
PyTorch 1.11 引入的 **FSDP(Fully Sharded Data Parallel)**本质上是 ZeRO-3 的 PyTorch 原生实现,但它有几个工程上更友好的特性:
# FSDP 的最小可用配置(PyTorch 2.3+)
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(
model,
sharding_strategy=ShardingStrategy.FULL_SHARD, # 等价 ZeRO-3
mixed_precision=MixedPrecision(
param_dtype=torch.bfloat16,
reduce_dtype=torch.float32,
),
backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
forward_prefetch=True,
limit_all_gathers=True,
)
与 DeepSpeed ZeRO-3 相比,FSDP 的工程优势在于:
- 与 PyTorch 原生 API 兼容:不需要把模型包成
deepspeed.initialize()返回的特殊对象,标准的nn.Module写法直接可用。 forward_prefetch/backward_prefetch控制:可以预取下一层的参数到本地,重叠通信与计算——这个细节可以让万卡集群的 MFU(Model FLOPs Utilization)提升 5-10 个百分点。limit_all_gathers限制:避免 FSDP 在 forward 阶段一次性 all-gather 太多层参数导致瞬时显存峰值,对长序列训练特别重要。
反面:FSDP 的 checkpoint 保存比 DeepSpeed 复杂——因为参数是分片的,save_pretrained 不能直接用,需要先调用 summon_full_params 上下文管理器聚合完整参数(DeepSpeed 的 module.save_checkpoint() 更简单)。
选型经验:2026 年的生产级训练,DeepSpeed(ZeRO-3 + 各种 offload 扩展)依然是 MoE / 超大模型的首选;FSDP 则是从单卡快速迁移到多卡 + 标准 PyTorch 流程的最低门槛路径。
五、张量并行(TP):Megatron-LM 的列切/行切
数据并行解决的是"同一份模型在不同卡上跑不同 batch",但 单卡放不下整个模型时,DP 帮不上忙,必须把单层计算本身切到多张卡上——这就是张量并行(Tensor Parallelism,TP)。
Megatron-LM(NVIDIA 2021)的核心思想是:Transformer 的 MLP 层是 ,可以把 按列切到 N 张卡上分别算 ,再把结果拼起来算 ;而多头注意力 可以按注意力头切——每个 GPU 负责 的头。各卡算完后用 all-reduce 合并结果。
通信模式:
MLP forward (2D mesh, TP=2):
GPU0: X @ W1[:, :d/2] → a0 ─┐
├─→ all-reduce → Y
GPU1: X @ W1[:, d/2:] → a1 ─┘
数学上,TP 把单层 MLP 的参数量从 ( + )降到 ,但每一步 MLP 都要触发一次 all-reduce——TP 通信粒度细、频率高,所以 TP 几乎只能在单机 8 卡 NVLink 内部用(NVLink 900GB/s 带宽 + ~1μs 延迟),跨机 TP 性能会断崖式下跌。
2026 共识:TP 通常固定为 8(对应单机 8 卡),不跨机。
六、流水线并行(PP):GPipe 与 PipeDream 的 1F1B 调度
TP 解决了"单层切多卡",但单机 8 卡仍然不够放 405B 模型(405B / 8 = 50.6B 单卡,远超 80GB 显存 + 优化器后的可用量)。这时需要把整个模型按层切到不同机器上——这就是流水线并行(Pipeline Parallelism,PP)。
朴素的 PP 把模型分成 N 段,每段放在一台机器上。前向时数据流过 N 段,反向时梯度反流——但问题是第 1 段算完才能传第 2 段,整条 pipeline 像流水线工厂一样运转,启动阶段有大量空泡。
GPipe(朴素 PP, 4 stages, micro-batch=8)的时间线:
S1: F1 F2 F3 F4 F5 F6 F7 F8 B1 B2 B3 B4 B5 B6 B7 B8
S2: F1 F2 F3 F4 F5 F6 F7 F8 B1 B2 B3 B4 B5 B6 B7 B8
S3: F1 F2 F3 F4 F5 F6 F7 F8 B1 B2 B3 B4 B5 B6 B7 B8
S4: F1 F2 F3 F4 F5 F6 F7 F8 B1 B2 B3 B4 B5 B6 B7 B8
↑ 空泡(GPU 等待上游 forward 完成的闲置时间)
1F1B(One Forward One Backward)调度是 PipeDream(Microsoft 2019)的核心改进:稳定状态下,pipeline 保持每个 micro-batch 同时跑一个 forward + 一个 backward,空泡时间被压缩到一个常量(等于 pipeline 深度 ):
其中 是 micro-batch 数。当 时 bubble ratio 趋近于 0。但 1F1B 让所有 stage 必须持有所有 micro-batch 的激活值(用于反向),内存压力反而上升——这就是为什么后续的 Interleaved 1F1B(每个 stage 拆成多个 model chunk)成为 2024+ 的标配。
工程经验:PP 段数 通常 = 集群节点数 / TP 度数;段间通信走 NCCL P2P,比 all-reduce 便宜。
七、3D 并行 + 序列并行 + 专家并行:组合的艺术
现代大模型训练通常是 TP=8(单机 8 卡)+ PP=节点数 + DP=剩余维度 的 3D 组合,再叠加 ZeRO-1(仅切优化器状态)+ 序列并行(SP)+ 专家并行(EP for MoE):
总卡数 N = TP × PP × DP × (可能的 EP 维度)
= 8 × 16 × 16 × 1 (稠密模型 Llama-3 风格)
= 8 × 8 × 16 × 4 (MoE 模型 DeepSeek-V3 风格)
序列并行(Sequence Parallelism)是 Megatron-LM 2022 提出的补充:当 TP=8 时,每张卡只持有序列的 1/8,但 LayerNorm 和 dropout 是序列维度的全局归一化——必须跨卡 all-reduce 才能算对。SP 把这些"全局算子"放到序列维度切分上,进一步减少 TP 卡上的冗余激活值。
**专家并行(Expert Parallelism, EP)**专属于 MoE 模型:把 个专家分到不同卡上,输入 token 经过 gate 路由到对应专家所在卡——通过 all-to-all 通信完成 token 重排。DeepSeek-V3 训练报告里 EP=64 是关键调优点之一,让 671B 总参数中只有 ~37B 激活参数参与每次前向。
Mermaid:3D 并行 + EP 的数据流
DP=16 (16 路数据并行)
│
├─ TP=8 (单机内 NVLink 张量并行)
│ │
│ └─ LayerNorm/Dropout → SP (序列并行)
│
├─ PP=8 (跨机流水线并行, 1F1B)
│
└─ EP=4 (MoE 路由, all-to-all)
组合选型口诀:先按"单机装不下"上 TP=8;按"一机不够"上 PP;按"数据吞吐不够"上 DP;按"显存仍有富余想扩大 batch"叠加 ZeRO-1;MoE 模型最后加 EP。
八、生产实战:DeepSeek-V3 / Llama-3 训练报告的数字真相
DeepSeek-V3(2024-12 发布的技术报告)公开的训练基础设施数字是这一代工程的"参考标杆":
| 维度 | 数字 | 工程含义 |
|---|---|---|
| 总卡数 | 2048 × H800 | 单集群规模 |
| 总成本 | 278.8 万 H800 卡时 | 约 5.6M 美元 |
| MFU | ~52% | 通信与计算重叠优化后的水平 |
| 并行策略 | PP=16 + EP=64 + ZeRO-1 + DP | 3.5D 并行(非纯 3D) |
| 专家数 | 256 routed + 1 shared | 第一层 MoE 的设计 |
Llama-3.1-405B(Meta 2024-07)则采用了更"传统"的 TP=8 + PP=16 + CP(Context Parallelism 跨机序列切分)+ DP=剩余 路径,重点优化了 CP 的 ring attention 把长序列训练的张量并行通信开销摊薄。
两个项目共同传递的工程信号是:MFU 50%+ 是 2026 年生产级训练的及格线,低于 40% 基本说明并行策略组合有问题。
九、生产级 3D 并行调优清单
- 先算内存再选策略:用第二节的公式精确计算 、、 给定下的单卡显存,决定从 DP 起手还是必须 TP+PP 起步。
- TP 固定为 8:跨机 TP 在 NVLink 缺失时会断崖——务必把 TP 框死在单机 8 卡内。
- PP 段数 = 节点数 或其因子;段间通信走 NCCL P2P,开启 interleaved 1F1B 减少 bubble。
- DP 维度做最后的扩展轴:DP 通信只是梯度 all-reduce,可用 gradient compression / overlap 优化。
- MoE 模型必须加 EP:DeepSeek-V3 的 256 路由专家 + EP=64 让激活参数降到 1/8,不切 EP 等于浪费显存。
- 启用 FSDP
forward_prefetch/backward_prefetch:通信与计算重叠,万卡集群可提 5-10% MFU。 - 监控 NCCL 通信时间占比:用
torch.distributed.monitored_barrier+nsysprofile,all-reduce 时间 > 计算时间 30% 时优先降 PP 段数。 - Checkpoint 策略选 DeepSpeed 风格:ZeRO-3 的分片 checkpoint 比 FSDP 标准 save 简单 50% 代码量。
- 失败重启时优先恢复 DP 维度:DDP 的恢复粒度最细,TP/PP 的状态恢复代价高。
- 小模型(≤7B)别用 3D 并行:DDP + ZeRO-2 就够;强行上 PP 会因为 bubble 反而更慢。
九点五、典型事故案例与复盘模式
下面三个事故 case 来自公开技术报告或社区复盘,展示 3D 并行配置错位时如何把数百万美元算力烧成静默 OOM:
Case A:TP 跨机的 NVLink 缺失性灾难
某团队在 32 节点集群上训练 70B 模型,错误地把 TP 设成 16(每节点 4 卡参与 TP)。前 24 小时训练一切正常,MFU 35%;第 25 小时突然出现部分 step 慢 8-10 倍——NCCL 的 all-reduce 在跨机 InfiniBand 上达到 100% 带宽瓶颈。诊断显示 TP=16 强制每层 MLP 都触发跨机通信(单节点 4 卡无法组成 TP group)。复盘:TP 必须严格框在 NVLink 域内(即单机 8 卡),跨机通信留给 PP。修复后改成 TP=8 + PP=4 + DP=8,MFU 回升到 48%。
Case B:Pipeline Bubble 反噬 671B 模型
某 MoE 团队在 512 卡 H100 上训练 200B 模型,初始配置 PP=64 + micro-batch=4,看起来张量切得很细,结果 MFU 只有 28%。profile 后发现 bubble ratio = ——micro-batch 太少了,每 4 个 step 就被空泡吃光。修复:micro-batch 调到 32,bubble ratio 降到 200%,MFU 回升到 51%。教训:PP 段数 与 micro-batch 数 必须满足 ,否则 bubble 直接吞掉算力。
Case C:ZeRO-3 与 Activation Checkpointing 的内存峰值冲突
某 13B 模型训练,启用 ZeRO-3 + Activation Checkpointing 后单卡显存峰值反而比 ZeRO-2 高 20%。诊断发现 ZeRO-3 把参数切到 1/N 后,每层 forward 都需要重新 all-gather 完整参数——这个通信过程中的临时张量在 GPU 显存里驻留,叠加 activation checkpoint 的反向重算峰值,把显存打爆。修复:在 FSDP 配置中加 limit_all_gathers=True(限制同时 all-gather 的层数)+ 把 activation checkpoint 颗粒度从"按层"调成"按 2 层一组",峰值显存降回预期范围。
共同教训:3D 并行不是越激进越好——每一维并行都引入新的通信/内存/调度张力,必须用
nsysprofile + 实际 MFU 数字校准,而不是"理论上更并行就更高效"。
十、参考文献
- Rajbhandari, S., et al. (2020). ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning. SC '20.
- Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053.
- Rasley, J., et al. (2020). DeepSpeed: System Optimizations Enable Training Deep Learning Models with over 100 Billion Parameters. KDD '20.
- Zhao, Y., et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. VLDB '23.
- Harlap, A., et al. (2018). PipeDream: Generalized Pipeline Parallelism for DNN Training. SOSP '18.
- Narayanan, D., et al. (2021). Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. SC '21.
- DeepSeek-AI. (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437.
- Dubey, A., et al. (2024). The Llama 3 Herd of Models. arXiv:2407.21783.
- Korthikanti, V. A., et al. (2022). Reducing Activation Recomputation in Large Transformer Models. arXiv:2205.05198.
- Liu, H., et al. (2024). Sequence Parallelism: Long Context Training Megatron-style. GitHub: NVIDIA/Megatron-LM #988.
本文为工程实战综合分析,所有性能数字基于公开技术报告与生产环境经验估算。具体配置需以集群硬件拓扑(NVLink/IB/RoCE 带宽)与框架版本为准。