【大模型基础设施工程】06：3D 并行深度——数据 / 张量 / 流水 / 序列 / ZeRO

单卡 80GB HBM 最多装 70B FP16 的参数（不含梯度、优化器、激活）。 要训 DeepSeek-V3 这样的 671B MoE，必须把模型切成几千片，让几千张 GPU 像一台机器一样工作。 切分的艺术，就是并行的艺术。

一、为什么单卡再也装不下

1.0 开胃菜：一张时间线

每一次规模跨越都伴随并行策略的一次补刀。

1.1 显存的四项开销

训练一个 Transformer，GPU 显存的消耗大体分为四块：

以 Adam 混合精度训 70B 为例：参数 + 梯度 + 优化器 ≈ 2+2+12 = 16 × 70B = 1120 GB，再加激活——单卡 80GB 连塞一角都不够。

1.2 算力的两堵墙

除了显存，还有两堵看不见的墙：

• 内存带宽墙：HBM3e 约 4.8TB/s，FP8 GEMM 的算力 / 带宽比决定你必须做更大的 batch 才能打满 Tensor Core。
• 集合通信墙：一次 AllReduce 70B 参数 FP16 = 140GB，哪怕是 NVLink 900GB/s，梯度同步也要几百毫秒。跨机 IB/RoCE 更慢。

所以并行化的目标不是”把模型切开能装下”就完事，而是在切分带来的通信代价和计算能完成的比例（MFU）之间找平衡。工业界 Pre-train 的 MFU 目标：密集模型 40%+，MoE 30%+，能稳到 50% 算极优。

二、并行的五个维度

2.0 先认清四种集合通信原语

后面所有讨论都离不开这四种 NCCL 原语，先一次性说清楚：

AllReduce(x_i)      → 每个 rank 拿到 Σ x_i         ；流量 ≈ 2(N-1)/N × |x|
ReduceScatter(x_i)  → 每个 rank 拿到 Σ x_i 的 1/P 片；流量 ≈ (N-1)/N × |x|
AllGather(s_i)      → 每个 rank 拿到全部 s_i 拼接   ；流量 ≈ (N-1)/N × |s|×P
All-to-All(x_i[j])  → rank i 的第 j 片发给 rank j  ；流量 ≈ (N-1)/N × |x|

恒等式：AllReduce = ReduceScatter + AllGather，流量减半——这正是 ZeRO-2/3 与 SP 的魔法来源。

2.1 DP：数据并行（Data Parallel）

最朴素：每张卡完整持有一份模型副本，把 global batch 切成 N 份，每卡前向 / 反向各一份，反向后对梯度做 AllReduce。

GPU0: model | batch[0:B/N]  --\
GPU1: model | batch[B/N:2B/N] --> AllReduce(grad) --> update
GPU2: model | batch[2B/N:..] --/

• 通信量：每步 2N 字节（FP16 参数规模的梯度）× AllReduce 因子 2(N-1)/N ≈ 2。
• 特点：bandwidth-bound，模型越大越痛。70B 模型一步 280GB AllReduce，跨节点极容易成为瓶颈。
• 变种：Gradient Bucketing（PyTorch DDP）、Hierarchical AllReduce、Gradient Compression（PowerSGD，精度有损，工业界很少用于 Pre-train）。

2.2 TP：张量并行（Tensor Parallel）

把单个大矩阵乘法切到多卡上，由 Megatron-LM 定型。切法两种：

• 列切（Column Parallel）：Y = XA，把 A 按列切 [A1, A2]，每卡输出 Yi = X Ai，最后需要 AllGather 合并 Y。
• 行切（Row Parallel）：Y = XA，把 A 按行切，同时 X 也按列切，每卡算 Yi = Xi Ai，最后 AllReduce 求和。

Megatron 的 MLP（Linear → GELU → Linear）采用”列切 + 行切”配对：

X --[col-split Linear]--> Y(split) --GELU--> Z(split) --[row-split Linear]--> AllReduce --> Out

Attention 的 QKV Linear 是列切（切 head 维度），Output Linear 是行切。一层 Transformer 有 2 次 AllReduce（forward）+ 2 次 AllReduce（backward，与 SP 结合可以换成 ReduceScatter/AllGather）。

• 通信量：每层 4 × batch × seq × hidden 字节（FP16）。
• 特点：通信频繁、强依赖 NVLink，一般只在 node 内（TP ≤ 8）。跨 NVSwitch 会掉 MFU。

2.3 PP：流水线并行（Pipeline Parallel）

按层切，不同 stage 放到不同 GPU 上。前向像流水一样：GPU0 算完第 1-8 层，把激活送到 GPU1 算 9-16 层……反向反着走。

经典问题是”气泡（bubble）“——头尾阶段的 GPU 必然有空闲。后文 §五 专门讲。

通信：stage 之间只发送 / 接收激活与激活梯度，是 P2P send/recv，数据量远小于 DP 的 AllReduce，跨节点友好。

2.4 SP：序列并行（Sequence / Context Parallel）

当 seq_len 达到 32K、128K 乃至 1M，激活（B × S × H）本身就爆显存，单 TP 切不动了。SP 沿 序列维度 切分：

• Megatron-SP：在 TP 的基础上，把 LayerNorm、Dropout 等”行无关”算子的激活按 seq 切，TP 通信由 AllReduce 换成 ReduceScatter + AllGather 等价变换，峰值激活降 TP 倍。
• Ring Attention（Liu et al. 2023）：把 KV 在环形拓扑上传递，每卡算局部 attention，然后传 KV 到下一个节点，几乎线性扩展到任意长序列。
• Ulysses（DeepSpeed）：在 attention 内部做 all-to-all，把 (B, S/P, H) 换成 (B, S, H/P)，走完 attention 再换回来。head 数要能被并行度整除。
• Context Parallel（CP）：Megatron-Core 的术语，通常指 Ring Attention 风格。

实操上：Ulysses 在 TP 已经切了 head 的情况下不好叠加，Ring/CP 正交性更好；DeepSeek-V3 和训练 128K+ 上下文的主流是 CP + TP + PP + EP + DP。

2.5 EP：专家并行（Expert Parallel，MoE 专用）

MoE 每层有几十到几百个 expert，但 token 只路由到 top-k 个。EP 做的就是把 expert 分散到 N 张卡：

token routing -> all-to-all dispatch -> local experts -> all-to-all combine

• 通信：两次 all-to-all，网络消耗 = 激活 × 2 × top-k × (1-1/EP)。
• 痛点：路由抖动（expert 负载不均）、all-to-all 拖尾、drop-token vs capacity factor。
• DeepSeek-V3 的 EP = 64，两阶段路由（node-limited routing）把跨节点 all-to-all 限制在可控范围。

五个维度的关系可以用一张图表达：

三、ZeRO 与 FSDP：DP 的进化形态

DP 的痛点是”每张卡都完整保存参数 / 梯度 / 优化器”，冗余极大。ZeRO（Zero Redundancy Optimizer，DeepSpeed 2019）把这三份状态切分到 DP 各卡上。

3.1 ZeRO 三阶段

ZeRO-3 的通信量 ≈ 标准 DDP 的 1.5 倍（AllGather + ReduceScatter + AllGather 反向），但显存降到 1/Ndp。

3.2 FSDP（Fully Sharded Data Parallel）

PyTorch 原生实现的 ZeRO-3：

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3
    mixed_precision=mp_policy,
    device_id=torch.cuda.current_device(),
)

关键细节：

• Unit 切分：以 Transformer Block 为粒度 wrap，前向进入 block 时 AllGather、退出时 free。
• Prefetch：下一个 block 的 AllGather 与当前 block 计算 overlap。
• FSDP2（PyTorch 2.4+）：基于 DTensor，取消 FlatParameter，与 TP / PP / TorchTitan 组合更干净。
• HSDP（Hybrid Shard）：节点内 FullShard + 节点间 Replicate，减少跨节点 AllGather。

3.3 ZeRO-Offload / CPU / NVMe

• ZeRO-Offload：优化器状态放 CPU，step 在 CPU 上算（Adam）。
• ZeRO-Infinity：参数也能下放到 NVMe，做到单机训百亿以上。适合”穷人 / 实验室”场景，吞吐会显著下降。

工业预训练一般不用 Offload，SFT / LoRA 经常用。

3.4 ZeRO-R：激活与临时 buffer

ZeRO 家族里还有常被忽视的一档：

• P_a（partitioned activations）：把激活按 TP 切到各卡，与 --sequence-parallel 等价。
• C_B（constant buffer size）：把 AllReduce buffer 固定，避免大张量 OOM。
• M_D（memory defragmentation）：周期整理显存碎片。

这些在纯 DeepSpeed 栈里是 ZeRO-R；Megatron / FSDP 都内化实现了类似机制。

3.5 FSDP vs DeepSpeed 实战差异

选型不是二选一：国内头部团队常”Megatron 主干 + DeepSpeed 的优化器/Offload + FSDP2 做 SFT/RL”混用。

四、3D 并行组合与通信量估算

单用任何一种都不够，工业训练都是 TP × PP × DP (× SP × EP) 的笛卡儿积。

4.1 经典配置

DeepSeek-V3 的一个设计亮点：放弃 TP，用 EP + PP + ZeRO-1 组合，因为 MoE 的 FFN 已被 expert 天然切分，不再需要 TP 切 MLP；这样避免了 TP 带来的每层 AllReduce，换来更大的 PP 窗口。

4.2 Group 布局与通信域

以 TP=8, PP=4, DP=16 共 512 卡为例，NCCL 里会创建三类 ProcessGroup：

rank = tp_rank + tp_size * (dp_rank + dp_size * pp_rank)

TP group：同一 (pp_rank, dp_rank) 的 8 张卡 → 放 node 内 NVLink
PP group：同一 (tp_rank, dp_rank) 的 4 张卡 → 跨节点，但通信量小
DP group：同一 (tp_rank, pp_rank) 的 16 张卡 → 跨节点 AllReduce

Rail-optimized 拓扑：ibN 网卡与 GPUn 一一配对，DP AllReduce 走自己的 rail，避免 cross-rail。这是 NCCL 2.18+ NCCL_CROSS_NIC=0 的默认行为。

4.3 通信量粗算（每 step）

记：a = batch × seq × hidden，L 层，N 参数。

经验规则：若 DP AllReduce 超过 20% 单步时间，考虑升 ZeRO 或砍 DP；若 TP AllReduce 跑出节点，直接掉 MFU 30%+。

4.4 Rank 布局的可视化

以 TP=4, PP=2, DP=2，共 16 GPU 为例：

                 PP stage 0              PP stage 1
            ┌──────────────┐       ┌──────────────┐
DP group 0  │ R0 R1 R2 R3  │──PP──▶│ R8 R9 R10 R11│
            │ ◀── TP ────▶ │       │ ◀── TP ────▶ │
            └──────┬───────┘       └──────┬───────┘
                   │ DP AllReduce         │ DP AllReduce
            ┌──────┴───────┐       ┌──────┴───────┐
DP group 1  │ R4 R5 R6 R7  │──PP──▶│R12 R13 R14 R15│
            └──────────────┘       └──────────────┘

• TP 组：{0,1,2,3}、{4,5,6,7}、{8,9,10,11}、{12,13,14,15}。
• PP 组：{0,8}、{1,9}、…、{7,15}。
• DP 组：{0,4}、{1,5}、…、{11,15}。

Megatron 按 tp × dp × pp 的顺序把 rank 扁平化，这也是”TP 放同 node”自然成立的原因——同一个 node 上的 8 个 rank 天然落入同一 TP 组。

五、流水线气泡：从 GPipe 到 DualPipe

5.1 气泡公式

设 PP 度 P，micro-batch 数 M：

• GPipe（先全前向再全反向）：bubble = (P-1) / M。若 P=8, M=16，气泡 47%！
• 1F1B（PipeDream-Flush / Megatron）：也是 (P-1) / M，但峰值激活显存降为 P 份而非 M 份。
• Interleaved 1F1B（Virtual Pipeline）：每卡持有 v 个非相邻层段，气泡 = (P-1) / (v × M)，代价是 P2P 次数 ×v。
• Zero Bubble（Sun et al. 2024）：把反向拆成 B（input grad）和 W（weight grad），重排 F / B / W 让气泡降到理论 0。
• DualPipe（DeepSeek-V3）：双向流水，两个方向上交错做前 / 反向，with 计算 / 通信重叠。实测 bubble ≈ 0% 且把 all-to-all 也藏到计算后面。

5.2 时序对比（Mermaid）

gantt
    title 1F1B vs Zero Bubble vs DualPipe（P=4，单位 = micro-batch 步）
    dateFormat  X
    axisFormat  %s

    section 1F1B
    GPU0 F1 :a1, 0, 1
    GPU0 F2 :a2, after a1, 1
    GPU0 F3 :a3, after a2, 1
    GPU0 F4 :a4, after a3, 1
    GPU0 B1 :b1, after a4, 1
    GPU0 B2 :b2, after b1, 1
    GPU0 B3 :b3, after b2, 1
    GPU0 B4 :b4, after b3, 1

    section ZeroBubble
    GPU0 F1 :c1, 0, 1
    GPU0 F2 :c2, after c1, 1
    GPU0 B1 :c3, after c2, 1
    GPU0 W1 :c4, after c3, 1
    GPU0 F3 :c5, after c4, 1
    GPU0 B2 :c6, after c5, 1

    section DualPipe
    GPU0 F1+ :d1, 0, 1
    GPU0 F2+B1- :d2, after d1, 1
    GPU0 F3+B2- :d3, after d2, 1
    GPU0 F4+B3- :d4, after d3, 1
    GPU0 B4- :d5, after d4, 1

5.3 Zero Bubble 关键：B / W 拆分

反向传播中：

• B（grad_input = grad_output @ W.T）：下游需要 → 必须尽早算。
• W（grad_weight = X.T @ grad_output）：只自己用 → 可以延后到流水尾巴。

把 W 往后推，中间空档让其他 micro-batch 的 F/B 填进来。

5.4 DualPipe 的双向 + 重叠

DualPipe 核心思想：

1. 双向调度：把模型分成两份对称流水，一份从头到尾、一份从尾到头，两股数据流反向流动，共享 GPU 计算单元。
2. 细粒度 chunk：将 attention / MLP / all-to-all dispatch / all-to-all combine 切为 4 块，按 ATTN → MLP → combine → dispatch 重排，让反向计算遮蔽前向 all-to-all。
3. 不需要 TP，天然适配 EP-heavy MoE。

代价：显存需要保存 2 份参数副本（因为两个方向同时在算），适合 MoE 这种参数相对 sparse 的结构；Dense 大模型用 DualPipe 不划算。

六、通信与计算 overlap

MFU 上 50% 的关键：让每一次集合通信都藏到 GEMM 后面。

6.1 NCCL stream 与 CUDA graph

# PyTorch 内部约定：通信走单独的 NCCL stream
comm_stream = torch.cuda.Stream()
with torch.cuda.stream(comm_stream):
    dist.all_reduce(grad, async_op=True)
# 计算继续在默认 stream

DDP 的 bucket 机制：梯度按桶（默认 25MB）触发 AllReduce，反向还没结束、前面桶的 AllReduce 已经在跑。

6.2 Chunk / micro-batch

• TP AllGather-MatMul Fusion（Megatron --sequence-parallel + FlashAttention 2.x + cuBLASLt）：把 AllGather 拆 chunk 与 MatMul pipeline 重叠，近年的”tensor parallelism with sequence parallel communication overlap”即此。
• NCCL SendRecv + Compute：PP 的 stage 通信可以先 issue send，再做下一 micro-batch 前向。

6.3 DualPipe 的 all-to-all 重叠

MoE 的 all-to-all 最痛苦。DualPipe 把 dispatch / combine 拆成 4 个 warp group，与本 chunk 的反向 compute 用不同 SM：

Compute SMs:  | attn_bwd | mlp_bwd |
Comm SMs:     |          dispatch(next) / combine(prev) |

这是 DeepSeek-V3 能在 H800（NVLink 带宽被砍一半）上做到 MFU 40%+ 的秘密武器。

七、激活重计算与梯度累积

7.1 Activation Checkpointing

反向用到的中间激活占用极大（B × S × H × L × 12 级别）。checkpoint：前向时只存 block 入口激活，反向时重新跑一次前向。换显存与 1/3 额外算力。

• Full Recompute：每个 Transformer Block 都重算，显存 ↓ ~5×，算力 ↑ ~33%。
• Selective Recompute（Megalite / Megatron-Core）：只 recompute 占内存大但算力便宜的部分（LayerNorm、Dropout、激活函数），FlashAttention 的 softmax 中间态本身就不存，效果极好。
• Offload Activations：把激活换到 CPU，适合长序列。

7.2 梯度累积

for step, batch in enumerate(loader):
    loss = model(batch) / accum_steps
    loss.backward()                       # 累在 .grad 上
    if (step + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

梯度累积把有效 batch size 放大而不增显存，但会降低 PP / DP AllReduce 的 overlap 机会（因为最后一次累积后必须同步）。Megatron 的做法是把 global_batch = DP × micro_batch × accum，accum 等于 PP micro-batch 数。

7.3 激活显存粗算

一个 Transformer Block，FlashAttention 2/3 下保存的激活量大约是：

act ≈ B × S × H × (34 / TP)  bytes（BF16，含 norm / residual / proj 输出）

L 层全保留：L × B × S × H × 34 / TP。以 L=96, B=2, S=4096, H=12288, TP=8 估算： 96 × 2 × 4096 × 12288 × 34 / 8 ≈ 41 GB / micro-batch，1F1B 再 × P = 几百 GB，不开 recompute 铁定 OOM。

开 selective recompute 后约 10 / TP，激活减到 12GB/micro-batch，完全可控。

7.4 与 CPU/GPU 传输的激活 offload

对超长序列（256K+），就算 recompute 仍然很痛。可以把 block 前向输出 offload 到 CPU，反向时异步拷回：

act = act.to("cpu", non_blocking=True)        # 前向尾部
...                                           # 其他 block 继续算
act = act.to("cuda", non_blocking=True)       # 反向前预取

PyTorch 的 torch.utils.checkpoint + offload_to_cpu=True、Megatron 的 --cpu-offloading 自动处理。代价是 PCIe 带宽（64GB/s for Gen5 ×16）会成为新瓶颈，典型折算 MFU 降 5–10%。

八、负载均衡的工程痛点

8.1 MoE 路由抖动

MoE 训练最头疼：某些 expert 成”网红”，其他 expert 饥饿。解决：

• Load-balancing loss（Switch Transformer）：softly 惩罚负载不均。
• Capacity factor：每个 expert 最多收 cf × avg_tokens，超出 drop。
• Auxiliary-free balancing（DeepSeek-V3）：取消 aux loss，用 bias 调节路由，既不污染主 loss 又保证负载。
• Node-limited routing：top-k 里限定跨的 node 数 ≤ M，把跨节点 all-to-all 压到可控。

8.2 PP 切分不均

每层 TFLOPS 不同：embedding、LayerNorm、FinalHead 比中间层轻。静态均分会让 stage0、stage_last 变慢成木桶短板。

• DP 切分法：测每层耗时 → 动态规划最小化 max(sum_stage)。Megatron 的 --pipeline-model-parallel-split-rank 支持自定义。
• Embedding 复制：头尾两个 stage 都放 embedding + head，拉近 workload。

8.3 动态 batch / seq

长短样本混批时，seq=4K 和 seq=128K 同 batch 等于短样本陪跑。做法：

• Sequence packing：把短样本塞进长 context，用 attention mask 隔离。
• Dynamic micro-batch：按 tokens 而非 samples 定义 batch。
• Sorted batching：同长度样本聚在一起（SFT 常用）。

8.4 PP 切分的动态规划伪码

给每一层一个 profile 耗时 t[i]，要切成 P 段使 max(Σt) 最小：

def split_pp(t, P):
    L = len(t)
    # dp[i][p]: 前 i 层切 p 段的最小瓶颈耗时
    INF = float("inf")
    dp = [[INF] * (P + 1) for _ in range(L + 1)]
    cut = [[0] * (P + 1) for _ in range(L + 1)]
    dp[0][0] = 0
    prefix = [0]
    for x in t:
        prefix.append(prefix[-1] + x)
    for i in range(1, L + 1):
        for p in range(1, P + 1):
            for j in range(p - 1, i):
                seg = prefix[i] - prefix[j]
                val = max(dp[j][p - 1], seg)
                if val < dp[i][p]:
                    dp[i][p] = val
                    cut[i][p] = j
    # 回溯 cut 得到每 stage 的层区间
    return dp[L][P], cut

生产上（Megatron-Core PipelinePartitioner，Colossal-AI PipelineGPTPolicy）用这种思路自动均衡。

九、通信优化细节

9.1 FP8 AllReduce / FP8 训练

NVIDIA Transformer Engine（TE）把 GEMM 输入做动态 scaling 转到 FP8（E4M3 / E5M2），AllReduce 本身一般仍走 BF16 / FP16（Hopper 的 NCCL 现已支持 FP8 collective，但数值稳定性需谨慎）。

DeepSeek-V3 自研了一套 E4M3-only 的 FP8 混合精度，对 GEMM 使用 FP8，对累加 / 权重更新保持 FP32，在 H800 上把算力吃干。

9.2 NCCL 拓扑感知

• Ring：所有 rank 串成环，带宽 = 最慢链路。简单但大规模下 latency 高。
• Tree：reduction 走树，延迟 O(log N)，大规模 DP AllReduce 默认。
• CollNet / SHARP（InfiniBand in-network reduction）：交换机里做加法，再降一半流量。
• Rail：每卡配专属网卡，避免 PCIe / NVSwitch 跨行。

Megatron 启动脚本常见：

export NCCL_IB_HCA=mlx5
export NCCL_IB_GID_INDEX=3
export NCCL_SOCKET_IFNAME=eth0
export NCCL_ALGO=Tree
export NCCL_CROSS_NIC=0
export NCCL_IB_SL=1

9.3 ZeRO Offload / CPU / NVMe

真显存不够时的兜底。一般组合：

• 训练：ZeRO-3 + Activation Checkpointing + FP8 足够了，不用 Offload。
• SFT / LoRA：ZeRO-2 + CPU Adam 就能在 A100 40G × 4 上跑 70B。
• 推理：不用 ZeRO，用 TP + PP + KV cache 分页（见 §12）。

十、实战栈

10.1 Megatron-LM（NVIDIA）

• 标配 TP + SP + PP + DP，社区最成熟。
• transformer_engine 集成 FP8。
• 新版 Megatron-Core 把并行化拆成可组合的 ParallelConfig，支持 Mamba 等 hybrid arch。
• 中国：百川、智谱、MiniMax、Qwen 的 pretrain 主干多为 Megatron fork。

10.2 DeepSpeed（微软）

• ZeRO 发源地，3 阶段 + Infinity + Offload。
• 流水线偏教学向，工业 pretrain 少用。
• 生态：HF Accelerate / Trainer 深度集成；国内火山引擎、阿里 PAI 都有 DS 预置镜像。

10.3 FSDP2 / HSDP（PyTorch 官方）

• 纯 Python，代码可读。
• PyTorch 2.4 起推荐 FSDP2 + DTensor + TorchTitan，官方给的 Llama 训练 demo。
• 与 torch.compile 搭配越来越丝滑。

10.4 Colossal-AI / TorchTitan / veScale / MindSpeed

• Colossal-AI（潞晨）：国内团队，支持 2D/2.5D/3D TP 等花式切法，中小规模适合。
• TorchTitan（Meta）：官方极简 pretrain 模板，TP + FSDP2 + PP 正在补齐。
• veScale（字节）：nD 并行 + DTensor，豆包训练栈。
• MindSpeed（华为）：昇腾生态的 Megatron 移植。
• Megatron-DeepSpeed（微软 + NV 联合分支）：BLOOM 训练所用。

10.5 DualPipe（DeepSeek 2024 开源）

DeepSeek 在 v3 后开源了 DualPipe 调度器（github.com/deepseek-ai/DualPipe），以及 EPLB（Expert Parallelism Load Balancer）和 profile-data。对做 MoE 预训练的团队是一手材料。

十一、代码示例：PyTorch FSDP2 + TP 最小样例

下面给出一个可运行骨架（需要 PyTorch 2.4+，实际训练请结合 TorchTitan）。

import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    ColwiseParallel,
    RowwiseParallel,
    parallelize_module,
)
from torch.distributed._composable.fsdp import fully_shard


class MLP(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.w1 = nn.Linear(d, 4 * d, bias=False)
        self.w2 = nn.Linear(4 * d, d, bias=False)

    def forward(self, x):
        return self.w2(torch.nn.functional.gelu(self.w1(x)))


class Block(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.attn = nn.MultiheadAttention(d, 8, batch_first=True)
        self.mlp = MLP(d)
        self.n1 = nn.LayerNorm(d)
        self.n2 = nn.LayerNorm(d)

    def forward(self, x):
        x = x + self.attn(self.n1(x), self.n1(x), self.n1(x), need_weights=False)[0]
        x = x + self.mlp(self.n2(x))
        return x


def build_model(num_layers=12, d=1024):
    return nn.Sequential(*[Block(d) for _ in range(num_layers)])


def main():
    dist.init_process_group("nccl")
    world = dist.get_world_size()
    rank = dist.get_rank()
    torch.cuda.set_device(rank % torch.cuda.device_count())

    # 2D mesh：行是 DP（FSDP），列是 TP
    tp_size = 2
    dp_size = world // tp_size
    mesh = init_device_mesh(
        "cuda",
        mesh_shape=(dp_size, tp_size),
        mesh_dim_names=("dp", "tp"),
    )

    model = build_model().cuda()

    # === TP：把每个 Block 的 MLP 列切 + 行切 ===
    tp_mesh = mesh["tp"]
    for block in model:
        parallelize_module(
            block.mlp,
            tp_mesh,
            {
                "w1": ColwiseParallel(),
                "w2": RowwiseParallel(),
            },
        )

    # === FSDP2：按 Block 粒度做 ZeRO-3 切分 ===
    dp_mesh = mesh["dp"]
    for block in model:
        fully_shard(block, mesh=dp_mesh)
    fully_shard(model, mesh=dp_mesh)

    model = torch.compile(model)

    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
    x = torch.randn(4, 256, 1024, device="cuda")

    for step in range(10):
        loss = model(x).square().mean()
        loss.backward()
        opt.step()
        opt.zero_grad()
        if rank == 0:
            print(f"step {step} loss {loss.item():.4f}")


if __name__ == "__main__":
    main()

启动：

torchrun --nproc_per_node=8 fsdp_tp_demo.py

11.1 扩展到 PP

PyTorch 2.4+ 提供 torch.distributed.pipelining（PP），基本接口：

from torch.distributed.pipelining import pipeline, SplitPoint, ScheduleGPipe, Schedule1F1B

stage_mod = pipeline(model, mb_args=(x,),
                     split_spec={"layer4": SplitPoint.BEGINNING})
schedule = Schedule1F1B(stage_mod, n_microbatches=8, loss_fn=loss_fn)
schedule.step(x, target=y)

TorchTitan 里有 TP + PP + FSDP2 的完整 llama3 demo，可以直接参考。

11.2 Megatron-LM 启动片段

作为对比，Megatron 的启动参数：

torchrun --nproc_per_node=8 --nnodes=128 \
  pretrain_gpt.py \
  --tensor-model-parallel-size 8 \
  --pipeline-model-parallel-size 16 \
  --num-layers 96 --hidden-size 12288 --num-attention-heads 96 \
  --seq-length 2048 --max-position-embeddings 2048 \
  --micro-batch-size 2 --global-batch-size 1024 \
  --sequence-parallel \
  --use-distributed-optimizer \
  --num-layers-per-virtual-pipeline-stage 4 \
  --recompute-activations --recompute-granularity selective \
  --bf16 --use-flash-attn \
  --transformer-impl transformer_engine \
  --fp8-format hybrid --fp8-amax-history-len 1024 --fp8-amax-compute-algo max

逐项意义：

• --sequence-parallel：把 TP 不切分的张量（LayerNorm 输入等）按 seq 维切，激活 /= TP。
• --use-distributed-optimizer：等价于 ZeRO-1，优化器状态按 DP 切。
• --num-layers-per-virtual-pipeline-stage 4：Interleaved 1F1B 的 v，进一步压气泡。
• --recompute-granularity selective：只 recompute 内存性价比高的 op。
• --fp8-format hybrid：E4M3 前向 + E5M2 反向（Transformer Engine 默认策略）。

11.3 DeepSpeed 配置片段

{
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 16,
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e8,
    "stage3_prefetch_bucket_size": 5e8,
    "stage3_param_persistence_threshold": 1e6,
    "offload_optimizer": {"device": "cpu", "pin_memory": true}
  },
  "bf16": {"enabled": true},
  "activation_checkpointing": {
    "partition_activations": true,
    "contiguous_memory_optimization": true
  }
}

十二、工程经验：选型与调参

12.1 决策树：到底用哪种并行？

模型能塞进单卡（含优化器）？
├─ 是 → DDP / FSDP1 即可
└─ 否 → 看参数规模
         ├─ < 20B：FSDP2（或 ZeRO-3） + Activation Ckpt，单节点搞定
         ├─ 20B–200B dense：TP(node 内 ≤ 8) + PP(跨节点) + DP + SP
         ├─ > 200B dense：TP + PP + DP + SP + Interleaved 或 Zero Bubble
         └─ MoE：EP(node 内或 node-limited) + PP + DP (+ 小 TP 或 无 TP) + DualPipe

并结合序列长度：

seq ≤ 8K    : 不需要 SP
seq 8K–32K  : Megatron-SP
seq ≥ 64K   : Context Parallel / Ring Attention
seq ≥ 256K  : CP + Activation Offload + FlashAttention

12.2 Batch size 怎么定

• Global batch = DP × micro_batch × grad_accum。Pre-train 常见 2M–4M tokens / step。
• Micro batch 受单卡显存限制，通常按”能塞满但不 OOM”取；先用 8 起步。
• Critical batch size（McCandlish et al.）：超过某 GBS，继续扩大收敛收益递减，需要调 LR 或换 optimizer（LAMB、AdamW 的 warmup+cosine）。

12.3 MFU 目标

MFU 30% 以下，先查：TP 是否跨节点；PP 气泡是否过大；ZeRO-3 是否 AllGather 打爆 IB；FlashAttention 版本；FP8 是否真开。

12.4 常见翻车

1. TP 跨节点：立刻 MFU 腰斩。规矩是 TP ≤ 节点 GPU 数。
2. PP micro-batch 太小：气泡吃掉一切。M ≥ 4P 是最低线。
3. Activation Ckpt 开太猛：浪费 30% 算力。用 selective + FlashAttention 后，full recompute 通常可以关掉。
4. DP 维度巨大：几千路 DP 做 AllReduce 跨机网络会饱和，用 HSDP 分层 reduce。
5. MoE expert 不够大 / EP 太粗：experts_per_rank = 1 时 all-to-all 代价 > 计算；建议 2–4。
6. CP + Ulysses 混用：路由和 head 维度冲突，debug 极难，选一个就行。
7. 混合精度 NaN：FP8 scaling 窗口设太长 / loss scaler 没跟上。Transformer Engine 的 fp8_autocast 默认窗口 1 常常够用。

十三、深入：一个 175B 训练的完整账本

把上面所有维度放到一次真实训练上算一笔账。假设：

• 模型：175B dense（类似 GPT-3），L=96，H=12288，heads=96，seq=2048。
• 硬件：1024 × H100 80GB，8 卡 / node，共 128 node。节点内 NVLink 900GB/s，跨节点 8×400G IB（单机 3.2Tbps）。
• 目标：global_batch = 4M tokens，MFU ≥ 45%。

13.1 切分配置

选 TP=8, PP=16, DP=8。

• TP=8 贴节点 GPU 数，MLP/Attn 的 AllReduce 全在 NVLink。
• PP=16 跨 16 个节点（每节点只属于一个 PP stage 的一片），stage 间 P2P 走 IB。
• DP=8 跨剩下 8 组节点做 AllReduce，用 HSDP 即可。

单卡参数：175B / (TP × PP) = 175B / 128 ≈ 1.37B，FP16 权重 2.74GB，加梯度 / 优化器：2.74 + 2.74 + 16.4 ≈ 22GB。

激活（全存）：L/PP × B × S × H × 12 bytes ≈ 6 × 2 × 2048 × 12288 × 12 ≈ 3.6GB / micro-batch（FlashAttention 下）。若 M=32 micro-batch，1F1B 峰值 P × 3.6GB ≈ 58GB——加权重已 80GB，必须开 selective recompute。

13.2 通信预算

单步数据流：

• 每层 TP AllReduce：4 × B × S × H × 2 = 4 × 2 × 2048 × 12288 × 2 ≈ 400MB / layer，NVLink 900GB/s 下 ≈ 0.45ms。96 层 × 4 次 ≈ 170ms（纯通信）。
• PP P2P：每 micro-batch 激活 B × S × H × 2 ≈ 100MB，IB 400Gbps（有效 ~40GB/s）下 ≈ 2.5ms / 次。
• DP ReduceScatter + AllGather（HSDP）：每 step 2 × 1.37B × 2 bytes ≈ 5.5GB，IB 下 150ms 级，必须被反向 overlap。

13.3 时序估算

单 micro-batch 纯前向计算时间约 6 × N × B × S / (FLOPS × MFU) = 6 × 175e9 × 2 × 2048 / (989e12 × 0.5) ≈ 8.7s / step（全模型）。分到 PP 每 stage ≈ 540ms；B 与 F 比 ≈ 2:1，所以 1F1B 一个 micro-batch 单 stage 耗时约 1.6s。

气泡 (P-1)/M = 15/32 ≈ 47%，太高！方案：

• 加 M：M=96（grad_accum=12）→ 气泡 15/96 ≈ 15%。
• 改 Interleaved 1F1B，v=4 → 气泡 15/(4×32) ≈ 12%。
• 改 Zero Bubble，气泡趋近 0，但要拆 B/W。

生产上 Megatron 默认 Interleaved，DeepSeek V3 用 DualPipe。

13.4 成本曲线

1024 × H100 × 一年 ≈ $15–25M（含电 + 折旧），所以每 1% MFU 提升 = 每年省 15–25 万美金。这就是为什么所有大厂要自研调度器、死磕 FP8 和零气泡——算力涨 1% 都是真金白银。

十四、国产与异构

14.1 华为昇腾（910B/910C）

• MindSpeed = Megatron 的昇腾移植；CANN 提供 HCCL（对标 NCCL）。
• 不支持 FP8，主用 BF16 / FP16。
• TP 通信走 HCCS（对标 NVLink，但带宽与拓扑不同），TP ≤ 8 同样适用。
• MoE + 3D 并行在昇腾上实现成熟度略落后于 N 卡，但盘古、Qwen 昇腾版已跑通。

14.2 寒武纪、天数、沐曦、摩尔线程

• 大多通过 PyTorch plugin + 自己的集合通信库做对接。
• 实际训练规模以中小（百亿级）为主；大规模训练的生态尚在补齐。

14.3 AMD MI300X / MI325X

• ROCm + RCCL（NCCL API 兼容）。
• PyTorch FSDP2 / Megatron 的 ROCm fork 可跑，Meta、微软公开过 MI300X 训练报告。
• FP8（OCP 规范）已支持；生态距 CUDA 还有 1–2 年。

14.4 TPU v5p / Trillium

• 完全不同体系：XLA + JAX + pjit。
• 并行通过 Mesh + PartitionSpec 声明，编译期决定通信。
• Google PaLM、Gemini 走这条路，外部团队使用较少，但 JAX + jax.experimental.shard_map 的模型定义极简洁。

十五、调试与性能分析

15.1 工具

• PyTorch Profiler + HTA（Holistic Trace Analysis）：画出每张卡的 kernel / comm timeline。
• Nsight Systems / nsys：GPU 侧最细粒度。
• NCCL_DEBUG=INFO：通信拓扑、算法选择一目了然。
• Megatron-Core --profile flag：自动标注每段耗时。
• torch.cuda.memory._snapshot()：查显存泄漏 / 碎片。

15.2 一个典型 MFU 排查流程

1. 看 tokens/sec/GPU，算 MFU。低于目标 5% 以上就要查。
2. nsys 看一个 step：
   • GEMM 占比 < 50%？→ 算子 / 精度问题（没开 FP8、没用 FlashAttention）。
   • 通信 kernel 与 GEMM 同时在跑？→ overlap OK。否则看是不是 async_op=False、或 bucket 太小。
   • PP 气泡是否如预期？→ M 调大或换调度。
3. 单独跑单卡 benchmark，对比理论峰值（989 TFLOPS × 0.7 为单 GEMM 经验线）。
4. 把 TP 调成 1（只 DP + PP）看 MFU，隔离 TP 跨节点问题。

15.3 数值稳定性坑

• FP8 溢出：检查 amax_history，必要时换 E5M2 / 回 BF16。
• MoE loss spike：aux-loss 权重太大、capacity 太紧、router 初始化。
• 长上下文 NaN：attention softmax 的数值在 128K 以上很敏感，需要 FlashAttention-3 + scale trick。
• Gradient AllReduce 溢出：ZeRO-3 反向多步累积，换成 FP32 reduce。

15.4 真实世界的”诡异故障”案例

工业训练的大部分时间都在和一些”看不见的坑”做斗争。下面几例来自公开的故障复盘（OpenAI、Meta、DeepSeek、字节、百度等技术博客）：

1. 静默数据损坏（SDC）：某 GPU 的 SM 偶发算错一个 bit，loss 缓慢偏离。定位靠 replay 两个不同 rank 的同批次数据比对。解决：周期性 checksum；故障卡打标，调度绕开。
2. HBM ECC 错误累积：BF16 训到后期突然 NaN，日志查到 ECC 纠正计数暴增。换卡 + 重 load checkpoint。
3. NCCL 死锁：某个 rank 因为磁盘满/dataset 读超时，集合通信全链卡住。方案：全 rank watchdog + 超时自动 dump stack。
4. 慢节点（straggler）：一张卡因散热问题降频 20%，PP/DP 全阻塞在 AllReduce。方案：周期性 ping 各 rank 的 step time，异常值剔除。
5. NVLink 一条挂了：节点内 8 卡变 7 卡有效，TP 全锅。nvidia-smi nvlink -s 自动巡检。
6. IB 抖动：丢包导致 NCCL 超时。方案：NCCL_IB_TIMEOUT/NCCL_IB_RETRY_CNT 调大 + 网络侧 BER 监控。
7. MoE 路由爆炸：某 step top-1 命中 > 80% 集中到单 expert，capacity overflow 丢 token，loss 直接跳 10×。方案：auxiliary-free balancing + bias EMA。
8. Checkpoint 写崩：torch.save 单机写 1TB 权重写到一半掉线。方案：分片并行写 + 原子 rename + 异步（见第 10 篇）。

这些坑没有哪一条能靠”看论文”避开，全是血泪。一套成熟训练系统的差异就体现在这里。

十六、FAQ：十个被频繁问到的问题

Q1：既然 ZeRO-3 能省到 1/N 显存，为什么还要 TP/PP？ ZeRO-3 的代价是前 / 反向前要 AllGather 完整权重，跨节点带宽一被打爆，MFU 就崩。TP/PP 切的是”永久切”，权重常驻本地，通信量反而小。超 50B 后纯 ZeRO-3 不现实。

Q2：FSDP2 和 Megatron-LM 二选一怎么选？ Python 代码可读性、快速试错、与 HF 生态打通 → FSDP2 / TorchTitan。 追求极限 MFU、成熟 3D 并行、FP8、稳定支持 671B 级 → Megatron-Core。

Q3：MoE 非得用 EP 吗？ 小 MoE（< 8 expert）可以用 TP 切。但 expert 多起来（DeepSeek 256 个），TP 切不动，必须 EP。EP + all-to-all 的瓶颈正是当下的主战场。

Q4：PP 的 micro-batch 越多越好？ 一定范围内是的。M ≥ 4P 气泡才可接受。但 M 太大会让激活存不下（1F1B 下激活 = P 份，不受 M 影响；GPipe 下 = M 份，会爆）。

Q5：ZeRO-Offload 到 CPU 真的能训 70B 吗？ 能，但吞吐常是纯 GPU 方案的 1/5–1/3。更适合 SFT / LoRA / 消融实验。

Q6：TP 为什么不能跨节点？ TP 每层 AllReduce 几百 MB，一步几十次。NVLink 900GB/s 和 IB 50GB/s 差一个数量级，跨节点 TP 的通信时间会盖住 GEMM。

Q7：DualPipe 是银弹吗？ 不是。显存要多一份权重副本，Dense 大模型不划算；对 MoE 的 all-to-all 重叠收益大。对上下文极长（激活巨大）场景也未必适合。

Q8：Sequence Parallel 和 Context Parallel 有什么差别？ Megatron 术语里，SP 特指与 TP 合体的版本（切非线性算子的 seq 维，通信量不变）。CP（Context Parallel）切真正的 attention 部分 seq 维，跟 Ring Attention 同类。两者通常合起来用。

Q9：FlashAttention 与并行怎么搭？ FlashAttention 本身不引入额外通信，与 TP/PP/DP 正交；与 CP 搭配时用 Flash-Attn 的 ring 变体（FA3 已内置）。

Q10：万卡训练什么时候必须上？ 当 training tokens > 10T、model > 300B、或你想 3–4 周内出 checkpoint。小于这个规模，千卡 + 更好的工程更划算。

十七、一页速查表

通用口诀：

• 带宽密集 → 留在 NVLink 内（TP/EP/SP）。
• 延迟不敏感、量小 → 放跨节点（PP）。
• 每 step 只一次 → DP/ZeRO 可以全局跨机。
• 有 MoE → EP 优先，配合 DualPipe + node-limited routing。
• 长上下文 → SP + CP，激活显存才是真敌人。

十八、小结

并行不是把模型切开就完事，真正的工程在于：谁切、切到哪个通信域、与什么计算 overlap、在哪个阶段同步。

• DP/ZeRO/FSDP 是基础盘——省显存，代价是 AllReduce。
• TP 切 hidden，吃 NVLink，严格节点内。
• PP 切 layer，跨节点友好，核心斗争是气泡——1F1B → Zero Bubble → DualPipe。
• SP/CP 切 seq，长上下文时代的必备。
• EP 是 MoE 的天然并行，配合 DualPipe 可以在 H800 上做出 MFU 40%+ 的 MoE 训练。
• 没有银弹，3D 乃至 5D 并行的组合 + 通信 overlap + 高质量算子库 才是万卡训练的基石。

下一篇我们拆解最具代表性的两套实现：Megatron-LM 与 DeepSpeed。

参考资料

• Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019
• Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters (PTD-P), 2021
• Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, 2020
• PyTorch FSDP / FSDP2 官方文档，torchtitan 仓库
• Huang et al., GPipe, 2019；Fan et al., DAPPLE / PipeDream, 2020
• Qi et al., Zero Bubble Pipeline Parallelism, 2024
• DeepSeek-AI, DeepSeek-V3 Technical Report, 2024；DualPipe / EPLB 开源仓库
• Liu et al., Ring Attention with Blockwise Transformers, 2023
• Jacobs et al., DeepSpeed Ulysses, 2023
• NVIDIA Transformer Engine 与 NCCL 官方文档
• Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models (Selective Recompute), 2022

上一篇：训练全景：Pre-train、SFT、RLHF、DPO、蒸馏 下一篇：Megatron-LM 与 DeepSpeed