第 7 章 高级并行策略
第7章 高级并行策略
本章系统阐述大模型训练所需的高级并行策略，覆盖流水线并行、序列并行、上下文并行、MoE 专家并行与 5D 组合，并 深入 ZeRO 显存优化、梯度检查点及多维并行调优实战。数据并行、张量并行与基础序列并行的机制在前序章节已交代， 本章聚焦各高级策略的调度机制、演进脉络与工程选型。
7.1 流水线并行全系列
流水线并行（Pipeline Parallelism, PP）将模型的层按顺序分配到不同的 GPU 设备上，每张 GPU 负责若干连续层的计 算。与传统模型并行不同，PP 的核心挑战在于减少流水线气泡（pipeline bubble），即设备空闲等待的时间。
7.1.1 GPipe 微批次
GPipe（Huang et al., NeurIPS 2019）是第一个将微批次（micro-batch）流水线引入深度学习的系统。核心思路：将一 个 mini-batch 切分为 m 个微批次，在 p 个流水线阶段间流水执行。 GPipe schedule (p=4 stages, m=4 micro-batches): Time ▶ S0: F0,0 | F1,0 | F2,0 | F3,0 | | B0,0 | B1,0 | B2,0 | B3,0 | S1: | F0,1 | F1,1 | F2,1 | F3,1 | | | B0,1 | B1,1 | B2,1 | B3,1 | S2: | F0,2 | F1,2 | F2,2 | F3,2 | | | | B0,2 | B1,2 | B2,2 | B3,2 | S3: | F0,3 | F1,3 | F2,3 | F3,3 | | | | | B0,3 | B1,3 | B2,3 | B3,3 | |<-- all F complete -->|<-- unified B -->| |<-- warmup -->| |<-- cooldown -->| GPipe 微批次流水线示意如上，调度分为 warmup 与 cooldown 两个阶段：所有阶段完成前向后才统一启动反向传播， 两端出现明显的设备空闲。 GPipe 的气泡分析公式： p−1 Bubble Ratio = m+p−1 当 m ≫ p 时气泡率趋近于零，但 GPipe 要求前向完成后才启动反向传播，导致激活值内存峰值高达 O(m)。 7.1.2 1F1B 调度机制 PipeDream（Narayanan et al., SOSP 2019）提出 1F1B（One-Forward-One-Backward）调度策略：在每个流水线阶 段，尽快交替执行前向和反向，将热身（warmup）和冷却（cooldown）阶段的气泡最小化。 Timeline (S0 perspective, p=4, m=8): [F0,0][F0,1][F0,2][B0,0][F0,3][B0,1][F0,4][B0,2][F0,5][B0,3][F0,6][B0,4][F0,7][B0,5][B0,6][B0,7] |<-warmup->|<----------------- 1F1B steady state ---------------->|<-cooldown->| 1F1B 的激活值峰值仅需 O(p) 微批次的激活值，相比 GPipe 的 O(m) 大幅减少。
7.1.3 双缓冲与交错 1F1B
- PipeDream-2BW PipeDream-2BW（Narayanan et al., 2021）在 1F1B 基础上引入双缓冲权重（Double-Buffered Weights），允许每个流 水线阶段持有两份参数拷贝： •计算拷贝：用于当前微批次的前向/反向计算 •更新拷贝：接收梯度同步和优化器更新 当计算拷贝工作时，更新拷贝在后台进行 AllReduce 和优化器步骤，完全隐藏了优化器开销。
- Megatron-LM 的交错1F1B Megatron-LM 的 Interleaved 1F1B 调度（Narayanan et al., SC 2021）进一步减少气泡，为每个设备分配来自模型中不 同区域的多个连续层段，如图7-1所示。 Interleaved 1F1B v=2 (2 non-contiguous segments per device) Standard 1F1B (1 contiguous segment per device) Stage 0 Stage 1 Stage 2 Stage 3 Stage 0 Stage 1 Stage 2 Stage 3 L0,L1 + L8,L9 L2,L3 + L10,L11 L4,L5 + L12,L13 L6,L7 + L14,L15 L0–L3 L4–L7 L8–L11 L12–L15 图7-1 标准1F1B vs 交错1F1B 的阶段分配 交错模式下流水线气泡率为： p−1 Bubble Ratiointerleaved = v⋅m+p−1 其中 v 是每设备的层段数（交错层数），典型值 v = 2 或 v = 4。当 v ⋅ m ≫ p 时近似为 。交错模式将气泡减少到约原 p−1 来的 1/v。 v⋅m
7.1.4 零气泡与双向流水线
- Zero-Bubble Zero-Bubble Pipeline（Qi et al., ICLR 2024）的核心创新是将反向传播拆分为两部分，理论上将气泡趋近于零：
- 将反向传播拆分为 B（计算输入梯度，存在流水线依赖，位于关键路径上）和 W（计算权重梯度，无依赖关系，不在关 键路径上）两部分
- 调度器将无依赖的 W 计算延迟至流水线气泡期执行
- 从而将原本空闲的气泡时间填充为有用的权重梯度计算，气泡利用率从 0 提升至接近 100%
- DualPipe DualPipe 由 DeepSeek 团队在训练 DeepSeek-V3（671B MoE）时提出，将双向流水线与 Zero-Bubble 的 B/W 拆分结 合，在 MoE 模型场景下有特殊优势。 DualPipe 从流水线两端同时驱动：前向从 stage 0→p-1，反向从 stage p-1→0，每个设备同时托管来自两条流水线方向 的阶段。关键收益：
- 双向驱动减少有效气泡深度：有效关键路径从 p 降为约 p/2，理论气泡率约为标准 1F1B 的一半。
- MoE All-to-All 通信隐藏：一个方向的 All-to-All 通信可被另一方向的计算覆盖，实现通信与计算的深度重叠，这是 DualPipe 针对 DeepSeek-V3 MoE 架构的最大优化点。 DualPipe schedule (p=4 stages, simplified): Forward pipeline (0▶3): F0▶F1▶F2▶F3 Backward pipeline (3▶0): B3▶B2▶B1▶B0 Each GPU overlaps its forward stage with the backward stage from the opposite end DeepSeek-V3 技术报告中在 H800 集群上实测达到约 37% 的 MFU。代价是激活内存约为标准 1F1B 的两倍（两条流水线 同时 in-flight）。DualPipe 在 MoE + EP 场景下优势最显著；稠密模型推荐优先考虑 Zero-Bubble 1F1B。
- Chimera Chimera（Li & Hoefler, SC 2021）提出双向流水线：p 个设备组成两个方向相反的流水线，同时处理两个不同的 micro- batch 序列。这种设计使用与标准流水线相同数量的 p 个设备，每个设备同时托管来自两条流水线的两个不同阶段（显存 需求约为标准方案的两倍），通过双向调度显著减少气泡。
7.1.5 通信模式与实用选型
- 通信模式 流水线并行仅需要在阶段边界进行 Peer-to-Peer (P2P) 通信：
import torch.distributed as dist
def forward_send(tensor, next_rank):
dist.send(tensor, dst=next_rank)
def backward_send(grad, prev_rank):
dist.send(grad, dst=prev_rank)
# Megatron-LM uses isend/irecv for async overlap:
# computation and network transfer proceed concurrently.
# PP P2P volume per transfer = b * s * h bytes.每次 P2P 传递的通信量为 bsh（batch × seq_len × hidden_dim）字节。相比 TP 每层 4 次 AllReduce，PP 的通信频率 低很多（与层数无关，只与微批次数量有关）。 2) 气泡可视化分析 以 4 阶段、8 微批次的标准 1F1B 为例： Stage 0: [F][F][F][F][B][B][B][B] Stage 1: [F][F][F][F][B][B][B][B] Stage 2: [F][F][F][F][B][B][B][B] Stage 3: [F][F][F][F][B][B][B][B] Bubble: [ warmup ][cooldown] 气泡时间占总时间的比例为 (p − 1)/(m + p − 1) = 3/11 ≈ 27.3%。当 m ≫ p 时近似为 (p − 1)/m，此例为 3/8 = 37.5% （过估计）。通过交错1F1B（v = 2）精确值降至 (p − 1)/(v ⋅ m + p − 1) = 3/19 ≈ 15.8%，近似值为 3/16 = 18.75%。 3) 实用建议 •PP 度不宜太高：气泡随 p 线性增长。实践中 PP≤8。 •均衡切分：各阶段的计算量应尽量均衡，避免某阶段成为瓶颈。 •与 TP 组合：TP 解决大型矩阵乘法，PP 解决层间依赖，两者互补。 •PP + 激活检查点：激活检查点应在 PP 阶段粒度设置（以 stage 为检查点边界）。 参考文献：Huang et al., “GPipe: Efficient Training of Large Neural Networks using Pipeline Parallelism”, NeurIPS 2019。Narayanan et al., “PipeDream: Generalized Pipeline Parallelism for DNN Training”, SOSP 2019。Qi et al., “Zero-Bubble Pipeline Parallelism”, arXiv 2024。
7.2 高级序列并行
基于 Megatron-LM 的基础序列并行（内置于 TP 中）是早期方案。随着长上下文模型（128K-1M tokens）成为主流，序 列并行技术已发展出多种独立于 TP 的方案，各自在通信模式、内存效率和注意力支持上有不同取舍。
7.2.1 Ulysses 机制详解
DeepSpeed Ulysses（Jacobs et al., 2023）的核心思路是将序列并行从 TP 中解耦，通过 All-to-All 将序列切分转换为 注意力头切分： GPU0 GPU1 GPU2 GPU3 Ulysses Workflow (s=8192, 4 GPUs) Step 1: All-to-All transforms (s/p, h) to (s, h/p) Input shape: (b, s/4, h) Scatter h-dim, gather s-dim Step 2: FlashAttention on full sequence, partial heads Attention (b, s, h/4) on local head partition Attention (b, s, h/4) on local head partition Step 3: Reverse All-to-All restores original layout Output shape: (b, s/4, h) GPU0 GPU1 GPU2 GPU3 图7-2 DeepSpeed Ulysses 的 All-to-All 通信模式 如图7-2所示，Ulysses 的一次 All-to-All 将输入 (b, s/p, h) 变换为 (b, s, h/p)，即从“每个 GPU 持有部分序列的所有头”变 为“每个 GPU 持有完整序列的部分头”。Attention 在该布局下各 GPU 独立计算，之后再用第二次 All-to-All 恢复原始布 局。 Ulysses 的优势： •Attention 完全独立，无需任何通信 •天然支持因果掩码（causal mask） •通信量：2bsh（两次 All-to-All）
- 与 Megatron-SP 的架构对比 Megatron-SP 和 DeepSpeed Ulysses 代表了序列并行的两种根本不同的设计哲学，两者的架构对比如表7-1所示。 表7-1 Megatron-SP 与 DeepSpeed Ulysses 架构对比 维度 Megatron-SP DeepSpeed Ulysses 序列切分方式 沿序列维度直接切分 s/p 第1次All-to-All将序列切分转为注意力头切分 注意力计算 需要跨设备通信（AllReduce） 各GPU独立计算完整序列的部分头 通信原语 AllGather + ReduceScatter 2次 All-to-All 与TP耦合 强耦合，内置于Megatron TP 独立于TP，可独立使用 每GPU通信量 ≈8bsh（与 TP 相当） 2bsh 适配场景 ≤32K、已启用TP GPU少+序列超长
- All-to-All 注意力分派实现 Ulysses 的关键在于两次 All-to-All 的语义转换。设输入为 (b, s/p, h) ，第一次 All-to-All 将 h 维度（hidden dim） 沿 SP 组切分为 h/p ，同时将 s/p 维度沿 SP 组汇聚为 s ： A2A1 : (b, s/p, h) → (b, s, h/p) 第一次 All-to-All 后，每个 GPU 持有完整序列长度 s，但只有 h/p 个注意力头维度的数据。这一布局本质上等价于“序列 维度全部本地化 + 头维度分片”，允许 FlashAttention 等高效注意力核在这些切片上独立执行。
# DeepSpeed Ulysses: All-to-All based attention dispatch
# Input shape per GPU: (batch, seq_len/sp_size, hidden_dim)
def ulysses_attention(query, key, value, sp_group):
# Step 1: All-to-All transforms (b, s/sp, h) -> (b, s, h/sp)
# scatter_dim=2 splits hidden_dim across ranks
# gather_dim=1 collects full sequence length from all ranks
query = all_to_all(query, scatter_dim=2, gather_dim=1, group=sp_group)
key = all_to_all(key, scatter_dim=2, gather_dim=1, group=sp_group)
value = all_to_all(value, scatter_dim=2, gather_dim=1, group=sp_group)
# Step 2: Each GPU runs full-sequence attention on its head subset
attn_out = flash_attention(query, key, value, causal=True)
# Step 3: Reverse All-to-All restores (b, s, h/sp) -> (b, s/sp, h)
attn_out = all_to_all(attn_out, scatter_dim=1, gather_dim=2, group=sp_group)
return attn_out第二次 All-to-All 将注意力输出从 (b, s, h/p) 恢复为 (b, s/p, h)，完成一轮序列并行，总通信量恰好为 2bsh。
7.2.2 交错与因果方案
- Striped Attention Striped Attention（Brandon et al., 2023）是 Ring Attention 的变种，将每个 GPU 负责的序列块交错分配到整个序列 中： Standard: GPU0:[0-2047] GPU1:[2048-4095] GPU2:[4096-6143] GPU3:[6144-8191] Interleaved: GPU0:[0,4,8...] GPU1:[1,5,9...] GPU2:[2,6,10...] GPU3:[3,7,11...] 交错分配使每个 GPU 的 KV 块更均匀地覆盖整个序列，减少了 Ring Attention 中的“热启动”延迟。
- DistFlashAttention DistFlashAttn（Li et al., 2023）将 FlashAttention 的分块策略与分布式执行结合，在保持 IO-aware 特性的同时分配注 意力计算。 关键创新是因果感知的负载均衡。由于因果掩码的存在，序列前半部分的 query token 只需 attend 到较少 key token。 DistFlashAttention 根据这一特性为负责前半部分序列的 GPU 分配更多计算任务，实现负载均衡。
7.2.3 锚点与统一方案
- Star Attention Star Attention（Acharya et al., 2024）引入锚点块（Anchor Block）概念，即全局可见的特殊块： •每个 GPU 持有本地序列块，同时也是全局锚点块 •计算注意力时，query token attend 到所有锚点块 + 本地块 •复杂度从 O(s ) 降到 O(s ⋅ (p ⋅ a + s/p))，其中 a 是锚点块大小
- USP USP（Unified Sequence Parallelism，Fang et al., 2024）尝试统一上述方案：根据硬件拓扑和序列长度自动选择最优的 组合策略。USP 综合考虑：节点内 NVLink（高带宽）→ Ulysses，节点间 IB（低带宽）→ Ring Attention。
7.2.4 方案对比与选型
- 方案对比 各方案在通信模式、内存效率与注意力支持上的差异如表7-2所示。 表7-2 序列并行方案对比 方案 通信模式 每GPU通信量 支持Causal 显存 复杂度 最佳场景 Megatron SP AllGather + ReduceScatter ≈8bsh（与 TP 相当） 是 低 中 与TP集成（≤32K） Ulysses All-to-All ×2 2bsh 是 中 低 GPU少+序列超长 Ring Attention P2P Ring 2(p − 1)skh/p 是 低 中 超长序列(s>128K) Striped Attn P2P Ring (交错) 同Ring 是 低 中 超长序列(s>128K) DistFlashAttn 混合 取决于分区 是 最低 高 极长序列+因果 Star Attn Anchor Gather O(p ⋅ a ⋅ s) 部分 低 高 近线性注意力 USP 自适应混合 根据情况 是 自适应 高 通用
- 选型决策框架 选择序列并行方案时，按以下优先级逐步缩小范围： Decision Flow:
- Is TP already enabled with Megatron-LM? -> YES: Use Megatron-LM SP, typical for s <= 32K, zero additional cost -> NO: Continue to step 2
- What is the target sequence length and GPU count? -> Few GPUs + long sequence: Ulysses (two All-to-All, attention independent) -> s > 128K: Ring Attention (memory grows with s, not hidden dim h)
- Consider hardware topology:
-> Single node (NVLink 600 GB/s): Ulysses (All-to-All benefits from high bisection BW)
-> Multi-node (IB 200 GB/s): Ring Attention (P2P friendly for inter-node)
-> Hybrid (USP): node-internal Ulysses + cross-node Ring Attention- For s > 128K: -> Ring Attention + FlashAttention-3 (fused block-sparse kernels) -> Consider sparse attention (Star, sliding window) for further reduction
- 定量对比 以 8×A100、序列长度 32K、模型 h = 8192 为例，各方案在固定集群上的定量对比如表7-3所示。 表7-3 固定集群下的序列并行定量对比 方案 SP度 通信量/GPU (MB) Attention通信次数 额外FLOPs 适配FlashAttn Megatron SP 8 393 2 (AG+RS) per layer 0 是 Ulysses 8 524 2 (A2A) per layer 0 是 Ring Attention 8 224 2(p-1) P2P per layer 0 是 Striped 8 224 同Ring 0 是 Ring Attention 在长序列场景通信量最低（与序列长度 s 而非 hidden dim h 相关），是超长上下文训练的首选方案。
- 实际选择指南 •已用 TP 且序列不超过 32K：默认 Megatron SP（零额外成本） •GPU 数少、序列超长：使用 Ulysses（两次 All-to-All，注意力独立） •序列超过 128K：使用 Ring Attention（内存随序列长度而非隐藏维度增长） •更长序列：Ring Attention + FlashAttention-3 + 稀疏注意力组合 统一口径：Megatron SP 只是把 TP 的 All-Reduce 重组为 All-Gather 与 Reduce-Scatter，单层每 GPU 通信总量约 8bsh ，与 TP 相当；Ulysses 的 2bsh 与 Ring 的环形 P2P 为各自方案的单层通信量。 大部分训练框架（Megatron-LM, DeepSpeed, PyTorch FSDP）在 2024 版本中都至少集成了 Ring Attention 或 Ulysses 中的一种。上下文并行则是序列并行的泛化，关注更广泛的上下文范围分布问题。 参考文献：Jacobs et al., “DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models”, arXiv 2023。Liu et al., “Ring Attention with Blockwise Transformers for Near-Infinite Context”, arXiv 2023。Dacheng Li et al., “DistFlashAttn: Distributed Memory-Efficient Attention for Long-Context LLMs Training”, arXiv:2310.03294, 2023。
7.2.5 超长序列并行专项
视频生成与多模态模型常产生远超文本的序列：一段视频 token 化后可达数万到数十万 token，序列并行面临显存、通 信、计算三重建模压力。本节给出超长序列（视频级）训练的瓶颈分析与并行专项配置。
- 三重瓶颈 对比普通长上下文（S=128K）与视频级序列（S 数十万）的三重瓶颈如表7-4所示。 表7-4 超长序列的三重瓶颈 维度 文本长上下文 (S=128K) 视频级序列 (S=数十万) 瓶颈来源 激活显存 中等（激活 ∝ S） 巨大（激活 ∝ S²，attention 中间量） 自注意力 O(S²) 中间矩阵 KV Cache 大 极大 KV ∝ S 通信量 序列并行通信 ∝ S 通信 ∝ S，且多次序列切分放大 All-to-All/AllGather 随 S 增长 计算量 attention ∝ S² attention ∝ S²，视频 batch 更大 每 token 计算随 S 增加 核心瓶颈是 attention 的 O(S²) 中间量，超长序列下 QK^T 矩阵物化即爆显存，必须用分布式 attention（Ring Attention）而非仅序列切分。
- 超长场景的序列并行选型 超长序列（视频级）下序列并行的选型与前文（普通长上下文）不同： •Ring Attention：唯一支持“近无限序列”的方案，序列切分加分块 attention，显存随 S 亚线性增长，适合视频级 S。代价是通信模式（P2P 环）与重叠优化复杂度高。 •Ulysses All-to-All：把序列切分与注意力头切分互相转换，适合 GPU 数与头数匹配的场景；S 极长时 All-to-All 通信量 上升，不如 Ring 扩展性好。 •Megatron-SP：内置于 TP 的序列切分，显存收益在超长序列下不够，只作为 TP 的补充。 超长序列的实践结论：S 进入数十万量级时优先 Ring Attention（或 Ring + CP 组合），Ulysses 在 GPU 数 ≤ 头数时仍可 用；选型以“显存能否容纳”为第一约束（Ring 显存增长最慢），通信量其次。
- 并行组合 单一序列并行不足以支撑视频级训练，需组合并行： •TP + SP：TP 内做序列切分，降低单卡激活；超长序列时 TP 度不宜过大（通信随 TP 增长），TP=4~8 配合 SP。 •CP + Ring：把超长序列切到多节点，Ring 做节点内/跨节点的分布式 attention；CP 度与序列长度匹配（S/CP 每卡负 载）。 •DP + 序列并行：batch 维度 DP 与序列维度并行正交组合，总并行度 = DP × CP × TP × PP。
- 通信量计算 序列并行通信量随 S 线性增长，视频级 S 下通信可能成为新瓶颈。以 Ring Attention 为例，每层通信量正比于序列切分 次数与 KV 交换量；需用通信量模型预估算例，确认通信占比可接受，必要时加大 CP 度摊薄每卡通信。
- 实践要点 •超长序列先确认显存可行（Ring 显存模型），再调并行组合与通信重叠。 •视频模型常配合“帧间结构”利用时间冗余（相邻帧共享 KV/注意力），并行实现需感知视频的时空结构，非纯文本式 均匀切分。 •超长序列的 batch 通常很小（1-2 条视频），DP 利用率低，序列并行（CP/Ring）成为扩展性的主要来源。
7.3 上下文并行
上下文并行（Context Parallelism, CP）是序列并行的泛化概念，关注如何将超长上下文（128K → 1M tokens）的 Transformer 计算和内存负载分布到多个 GPU 上。当上下文长度超过某一阈值时，单 GPU 的 HBM 容量（80GB）无法容 纳完整的 KV Cache 或注意力中间结果，CP 成为必需。
7.3.1 概念与序列对比
上下文并行是序列并行的父集，主要区别在于目标和使用场景，如表7-5所示。 表7-5 序列并行与上下文并行对比 维度 序列并行 (SP) 上下文并行 (CP) 主要目标 减少 TP 中的冗余计算/存储 支持超长上下文训练 序列长度 4K-32K 32K-1M+ 切分粒度 序列维度 序列维度+注意力计算 与TP的关系 辅助TP 独立策略 关键瓶颈 LayerNorm/Dropout 冗余 Attention 计算+KV Cache
7.3.2 环注意力算法
Ring Attention 是上下文并行最核心的算法。其基本思想是将 self-attention 的 Q, K, V 计算沿序列维度分片到 p 个 GPU，通过环形 P2P 通信让每个 GPU 依次访问其他 GPU 的 KV 块。
- QKV 分片与算法分解 设序列长度为 s，GPU 数为 p，每个 GPU 持有 s/p 长度的本地序列块。本地块生成对应的 Q , K , V ∈ R （其中 d b×(s/p)×d 为 head dim）： i i i Q i = Xi W Q , i ∈ [0, p − 1] K i = Xi W K , V i = Xi W V 注意力计算需要 Q 与所有 p 个 KV 块交互。Ring Attention 通过 p − 1 轮 P2P 传输，使每个 GPU 逐次获得其他 GPU 的 KV 块，最终完成完整序列的注意力计算。 i
- 环形通信模式 Step-by-step
Step 0: GPU i holds (Qi, Ki, Vi), computes attn(Qi, Ki, V i) Ring Attention Data Flow (p=4) Step 1: Send (Ki, Vi) to GP GPU1 GPU2 U (i+1)%p, receive from (i- 1)%p, compute attn(Qi, K_ GPU0 GPU3 recv, V_recv) Step 2: Repeat p-1 times, t otal p attention blocks per GPU 图7-3 Ring Attention 环形数据流 如图7-3所示，每轮通信传输 b ⋅ (s/p) ⋅ h 字节的 KV 块。总通信量： kv s Mring = (p − 1) ⋅ 2 ⋅ b ⋅ ⋅ hkv bytes per GPU p 其中 h = n ×d ，因子 2 表示发送和接收各一次。相比标准 All-to-All（2bsh），Ring Attention 的通信量与序 kv_heads 列长度 s 相关，而非 hidden dim h。当 h ≫ h （如 GQA 的 h ≪ h）时，Ring Attention 通信量更低。每次 P2P 传输 kv head 量小且可流水线化，适合在节点间网络传输，核心循环如下： kv kv
# Ring Attention core logic
def ring_attention_step(q, k, v, kv_chunk_from_prev):
# 1. Receive KV chunk from previous GPU
# 2. Compute local attention with local Q
# 3. Send local KV chunk to next GPU
# 4. Loop s/block_size times to cover full sequence
for step in range(world_size):
local_attn = flash_attention(q, local_k, local_v)send_recv_kv(local_k, local_v) # P2P communication 3) 权衡分析 Vanilla Attention 与 Ring Attention（p=4）的权衡对比如表7-6所示。 表7-6 Vanilla 与 Ring Attention 权衡对比 维度 Vanilla Attention Ring Attention (p=4) 计算量 (FLOPs) O(bs d) per GPU O(bs d/p) per GPU KV显存峰值 O(bsh ) per GPU kv O(bsh /p) per GPU kv 通信量 0 (单GPU) (p − 1) ⋅ 2bshkv /p 可扩展上下文 ≤ 单GPU HBM容量 线性扩展至 p× HBM 额外通信延迟 0 O(log p) P2P跳数 Ring Attention 以通信量换取了序列长度的线性扩展能力。在 H100 单 GPU 80GB HBM 上，vanilla attention 最大支持 约 32K 上下文（Llama-70B 规模）；Ring Attention 在 8 卡下可扩展至约 256K。 4) 实现参考 Ring Attention 已在多个框架中实现： •FlashAttention 生态： flash_attn 库提供了 ring_flash_attn_func ，基于 FlashAttention-2 的分块内核，支持 Varlen 序列和 GQA •Megatron-LM： context_parallel 模块实现了 Ring Attention + Striped Attention 的混合调度，通过 --context- parallel-size 参数控制 •Striped Attention：Megatron 的交错 token 分配可与 Ring Attention 组合，在 2D 网格（CP × TP）中协同工作
7.3.3 负载均衡分区
超长上下文训练中面临的核心问题是：因果掩码（causal mask）导致序列前半部分的 token 需要 attend 到的 KV 对数量 远少于后半部分，如图7-4所示。 U-Shaped Partition (Balanced) Naive Partition (Unbalanced) 0-7 + 56-63: 520 steps 8-15 + 48-55: 520 steps positions 16-31: 392 steps 0-15: 136 steps attention 16-23 + 40-47: 520 steps 24-31 + 32-39: 520 steps 32-47: 648 steps 48-63: 904 steps 图7-4 上下文并行的负载均衡分区策略 U 型分区（U-Shaped Partition）将序列的头尾组合到同一 GPU，使得每 GPU 的计算量更均衡。
7.3.4 系统实现与集成
- DeepSeek-V3 的上下文并行设计 DeepSeek-V3（2024）在大规模训练中采用了上下文并行方案：
- 序列分块（Chunk）：将序列切分为多个 chunk，按 chunk 维度并行的方式部署。
- CP 通信模式：在 Attention 计算中，query 需要访问来自其他 GPU 的 KV chunk，采用 P2P 传输。
- 负载均衡：根据计算的负载分布动态调整 chunk 分配。 DeepSeek-V3 训练中 CP 度与 SP 度可以不同（cp=4, sp=1），适用于超长上下文场景中独立调节。
- 与其他并行策略的集成 CP 与 TP、PP、DP 的组合是边界情况：
# Pseudocode: CP + TP + DP integration
class DistributedAttention:
def forward(self, q, k, v, causal_mask):
# 1. CP: if enabled, fetch remote KV chunks first
if self.cp_size > 1:
remote_kv = self.cp_gather_kv(k, v) # get KV from other cp-ranks
# 2. TP (head dim): each GPU computes attention for subset of headstp_q, tp_k, tp_v = self.split_by_head(q, k, v)
# 3. Standard attention computation
attn_out = flash_attention(tp_q, tp_k, tp_v, causal_mask)
# 4. CP reduce: reduce gradients of output version
return attn_out常见问题： •通信爆炸：CP+TP 的组合会使通信次数大幅增加，须谨慎评估是否真的需要两者同时使用。 •CP 与 DP 的组合：CP 与 DP 是正交维度，总 GPU 数 = DP × CP × TP × PP，两者可以同时使用。实践中通常将 CP 通信域限制在节点内（8 卡 NVSwitch 范围），以利用 NVLink 高带宽传递 KV 块；DP 的梯度 AllReduce 则跨节点进 行，两者通信域不重叠，不存在路径冲突。DeepSeek-V3、Megatron-LM 等实际系统均同时使用 CP 与 DP。
7.3.5 显存与性能分析
- 长上下文训练的显存分析 以 Llama-70B（h = 8192, n = 64）训练为例，不同上下文长度下的显存需求如表7-7所示。 heads 表7-7 长上下文训练的显存需求 上下文长度 参数 (BF16) 梯度 优化器(Adam FP32×3) 激活值 KV中间 总计/GPU (8路TP) 4K 140GB 140GB 840GB 15GB 2GB 1137GB / 8 = 142GB 32K 140GB 140GB 840GB 120GB 16GB 1256GB / 8 = 157GB 128K 140GB 140GB 840GB 480GB 64GB 1664GB / 8 = 208GB 当 s = 128K 时，即使 8 路 TP，单 GPU 仍需约 208GB，大幅超出 H100 80GB。此时需要引入 CP 将序列也切分（如 CP=2 将 reduce 到约 104GB），同时开启激活检查点，或将优化器状态卸载到 CPU。
- 不同上下文长度的训练吞吐 参考 DeepSeek 和 Meta 等公司的公开数据，不同上下文长度的训练吞吐如表7-8所示。 表7-8 不同上下文长度的训练吞吐 上下文长度 每GPU序列块 显存占用增加 MFU 降幅 推荐策略 8K 1 baseline baseline SP (标准) 32K 2 +60% -8% CP + Ring Attention 128K 4 +200% -25% CP + Ulysses + FlashAttn 256K 8 +500% -40% CP + Ring + activation ckpt 1M 16 +2000% -60% CP + 多级检查点 + 稀疏注意力
- Ring Attention 可扩展性分析 Ring Attention 的上下文长度随 GPU 数线性扩展。设单 GPU 最大支持上下文长度为 s ，则 p 个 GPU 下的最大上下文 max 为 p ⋅ s （忽略 KV 通信的临时缓冲区开销，约 2bsh /p 字节）。以 H100 (80GB) 和 Llama-70B (h = 1024) 为例，扩展 max 性如表7-9所示。 kv kv 表7-9 Ring Attention 可扩展性 GPUs (p) 最大上下文 (s_max) 通信量/GPU (GB) 通信时间估算 1 32K 0 0 ms 2 64K 0.4 2 ms (NVLink) 4 128K 0.6 3 ms (NVLink) 8 256K 0.7 4 ms (NVLink) 16 512K 0.75 9 ms (IB HDR) 32 1M 0.78 19 ms (IB HDR) 跨节点后通信延迟显著增加，这是为什么 CP 域通常限制在 8 卡（单节点 NVSwitch）范围内的原因。
7.3.6 与推理的关系
虽然本章聚焦训练，需注意 CP 在长上下文推理中也有关键作用。在推理的 Prefill 阶段（首次处理所有 prompt token）， CP 可以： •将 prompt 分片到多个 GPU 并行处理 •加速首个 token 的生成（Time To First Token, TTFT） •之后 Decode 阶段的 KV Cache 保持在各 GPU 中分布 参考文献：DeepSeek-V3 Technical Report, 2024。Liu et al., “Ring Attention with Blockwise Transformers
7.4 MoE 专家并行
for Near-Infinite Context”, arXiv 2023。 专家并行（Expert Parallelism, EP）是混合专家（Mixture of Experts, MoE）模型中专属的并行策略，将不同的专家 （expert）参数放置在不同的设备上，每个 token 根据路由器的选择只激活少量（top-k）专家。
7.4.1 MoE 架构与路由
标准 MoE Transformer 层将 FFN 替换为多个专家的组合：
class MoELayer(nn.Module):
def __init__(self, hidden_dim, ffn_dim, num_experts, top_k):
self.router = nn.Linear(hidden_dim, num_experts) # gating network
self.experts = nn.ModuleList([
ExpertFFN(hidden_dim, ffn_dim) for _ in range(num_experts)])
self.top_k = top_k
def forward(self, x):
# x: (num_tokens, hidden_dim)
router_logits = self.router(x) # (T, E)weights, indices = torch.topk( torch.softmax(router_logits, dim=-1), self.top_k, dim=-1 ) # Token dispatching: route tokens to corresponding experts # ... (see below) 每个 token 只激活 top-k 个专家（通常 k=1 或 2），总计算量为 O(E ⋅ k/E) = O(k)，实现“条件计算”。
7.4.2 通信与容量因子
- All-to-All 通信模式 MoE 并行中最关键的通信操作是 Token Dispatching 和 Token Combining，如图7-5所示。 Token Dispatching (All-to-All) T tokens each with router choice Permute Permute GPU0: Expert 0-3 GPU1: Expert 4-7 GPU2: Expert 8-11 tokens for E0-3 tokens for E4-7 Expert Computation Compute E0-E3 Compute E4-E7 Compute E8-E11 Unpermute Unpermute Unpermute Token Combining (All-to-All) token order 图7-5 MoE 中的 All-to-All Token 分发与合并 All-to-All 的通信量为非对称：某些专家可能收到远多于其他专家的 token。设 T 个 token 中分配给 Expert i 的比例为 p ，则该专家的通信量为 p ⋅ T ⋅ h。 i i
- 容量因子与 Token 丢弃 为防止某专家收到过多 token（OOM），引入容量因子（Capacity Factor, CF）： T ⋅k Capacity = CF ×
E Token 数量超过 capacity 的专家会丢弃溢出 token（token dropping）。CF 越大丢弃越少但显存需求更大： •CF=1.0：刚好容纳平均 token 量，丢弃约 1-5% •CF=1.25：丢弃降至 <0.1%，常见配置 •CF=2.0：基本不丢弃，但 MoE 层的 token buffer 显存翻倍（对整体模型显存影响有限，约增加 10-20%）
7.4.3 路由稳定性与均衡
- 辅助损失负载均衡 MoE 训练中专家容易“坍塌”，少数专家接收绝大部分 token。辅助损失（auxiliary loss / load balancing loss）用于鼓 励均匀分配： E Laux = E ⋅ ∑ fi ⋅ Pi i=1
其中 f 是分配给 Expert i 的 token 比例，P 是路由器分配给 Expert i 的 softmax 概率均值。最小化此损失鼓励 f 和 P 都趋向 1/E。 i i i i 2) DeepSeek-V3 无辅助损失方案 传统辅助损失（Auxiliary Loss）方法存在根本矛盾：损失权重 α 难以调节，太小则负载均衡效果差，太大则干扰语言建 模损失的梯度方向，导致模型质量下降。DeepSeek-V3（2024）提出了 Auxiliary-Loss-Free 负载均衡方案，核心思路是 在路由器 logits 上添加可学习的偏置项而非梯度惩罚：
class AuxFreeMoERouter(nn.Module):
def __init__(self, hidden_dim, num_experts, top_k):
super().__init__()
self.router = nn.Linear(hidden_dim, num_experts, bias=False)
# One bias per expert, updated manually (no gradient)
self.expert_bias = nn.Parameter(
torch.zeros(num_experts), requires_grad=False)
self.top_k = top_k
def forward(self, x):
logits = self.router(x) # (T, E)
biased_logits = logits + self.expert_bias # routing uses biased logits
top_k_idx = torch.topk(biased_logits, k=self.top_k, dim=-1).indices
# Weight computation uses ORIGINAL logits (clean gradient)
weights = F.softmax(logits.gather(-1, top_k_idx), dim=-1)
return top_k_idx, weights@torch.no_grad()
def update_bias(self, token_counts, gamma=0.001):
target = token_counts.float().mean()
load_imbalance = token_counts.float() - target
# Overloaded expert: lower bias; idle expert: raise bias
self.expert_bias.data -= gamma * torch.sign(load_imbalance)这种方法实现了路由决策（含偏置）与权重计算（不含偏置）的解耦：偏置仅影响哪些专家被选中（离散决策），不影响 选中后的加权系数（连续梯度），因此完全不干扰主损失的梯度优化方向。DeepSeek-V3 通过该方法在 256 个细粒度专家 上实现了负载不均衡率低于 5%，同时模型质量优于辅助损失方案。
7.4.4 系统实现方案
- DeepSpeed-MoE DeepSpeed-MoE（Rajbhandari et al., 2022）将 ZeRO 和 MoE 结合： •ZeRO 对非专家参数（attention, embedding）做分片 •专家参数各自分布在不同设备上 •支持层次化 All-to-All（节点内 NVLink + 节点间 IB）
- DeepSeekMoE DeepSeekMoE 架构（Dai et al., arXiv:2401.06066, 2024）提出两个关键创新，并在 DeepSeek-V2/V3 中规模化验证：
- 细粒度专家：将标准 MoE 的大专家进一步拆分为更多小专家。例如 E = 8 个标准专家 → E = 256 个细粒度专家。增加 路由灵活性的同时，单专家参数量减少。
- 共享专家：指定部分专家为共享专家（shared experts），所有 token 都必须经过这些专家（外加 top-k 的稀疏专家）。 共享专家捕获通用语言特征。
# DeepSeekMoE routing logic
class DeepSeekMoERouter:
def forward(self, x):
# Shared expert part (always activated)
shared_out = self.shared_experts(x)
# Fine-grained expert part (top-k sparse activation)
router_logits = self.router(x) # (T, 256)top_k_idx, top_k_weights = self.route_256_experts(router_logits) sparse_out = self.sparse_experts(x, top_k_idx, top_k_weights) return shared_out + sparse_out 3) 层次化 All-to-All 与 DeepEP 当专家分布在多节点时，朴素 All-to-All 会产生节点间全连接通信，昂贵且不可扩展。层次化方案：
- 节点内 All-to-All：使用 NVLink 高带宽快速交换 token
- 节点间 All-to-All：仅传输需要到其他节点的 token，减少跨节点通信量 DeepEP（DeepSeek, 2024）是专门优化的层次化 All-to-All 通信库，通过低延迟内核与基于流式负载均衡的调度策略提 升 MoE 场景的通信效率。
7.4.5 并行约束与集成
- 并行约束条件 MoE 并行的约束条件如表7-10所示。 表7-10 MoE 并行约束条件 约束 说明 示例 专家数量 必须是 EP 度的整数倍 E=64, EP=8 → 每GPU 8个expert 容量因子 决定每步的显存峰值和丢弃率 CF=1.25 → 丢弃 <0.1% 路由均衡 辅助损失权重需实验确定 α=0.01 (常见) 负载失衡 某GPU专家过多时OOM 动态重新分配expert-to-GPU映射
- EP 与 DP 的正交性 专家并行（EP）和数据并行（DP）是两个正交维度： •DP 维度：参数的梯度需要在 DP workers 间 AllReduce •EP 维度：token 需要在专家所在的设备间 All-to-All 两者可以组合（EP + DP），但需注意：非专家参数（attention, embedding）由 DP 负责同步；专家参数则只由 EP 维度 管理。
- EP + DP 组合的梯度同步策略 不同参数类型需要不同的同步方式：专家参数（各 EP rank 独有的 FFN 权重）只需在 DP 维度做 AllReduce（同一 EP rank 的不同 DP 副本之间）；共享参数（Attention、Embedding、LayerNorm、路由器等）需在整个 EP × DP 通信域内 AllReduce。DeepSpeed-MoE 和 Megatron-LM 均提供参数分组接口，自动将参数归入 expert_params （EP 内 DP 同 步）或 non_expert_params （全局同步）两组，分别使用不同的通信器处理。 参考文献：Shazeer et al., “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”, ICLR 2017。Rajbhandari et al., “DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale”, ICML 2022。DeepSeek-AI, “DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model”, arXiv 2024。 7.5 5D 并行组合策略 当训练千亿甚至万亿参数的模型时，单一的并行策略无法满足需求。5D 并行，即 DP + TP + PP + CP + EP 的组合，是训 练 GPT-4、Llama-405B、DeepSeek-V3 等模型的实际方案。本节介绍如何系统地组合这五种并行维度。
7.5.1 约束方程与约束
设总 GPU 数为 N ，各并行度满足： N = dp × tp × pp × cp × ep （注：公式中第 4 维为 Context Parallelism，CP。Megatron 原生 SP 始终与 TP 绑定、不独立占用 GPU 维度，不作为乘 积中的独立项。） 但每个 token 的计算还需要满足以下约束。
- 显存约束 Memper GPU ≥ Mparam + Mgrad + Moptim + Mactivation 其中 M param = Ptotal × bytes/(pp × tp × ep) （标准5D并行，DP不分片参数；使用ZeRO-3时分母再乘以dp），M activation = 。 f (bs, h, L, checkpointing)
- 通信带宽约束 节点内 NVLink 容量为 B × links per GPU，节点间 IB/RoCE 容量为 B × NICs per node。 intra inter
- 模型结构约束 num_heads % tp == 0（TP 头数兼容）、num_experts % ep == 0（EP 专家数兼容）、L % pp == 0（PP 层数均衡）。 7.5.2 3D 并行配置
- DeepSpeed 的 3D 并行 DeepSpeed 首先将并行问题简化为 3D（DP + TP + PP）：
DeepSpeed 3D parallel config example
ds_config = { "train_batch_size": 2048, "train_micro_batch_size_per_gpu": 1, "zero_optimization": { "stage": 2, # DP + TP + PP scenario, stage 2 is usually optimal }, "tensor_parallel": { "tp_size": 8, # intra-node }, "pipeline_parallel": { "pp_size": 4, # cross-node }, } DeepSpeed 的 Autotuning 功能（ deepspeed --autotuning run train.py 或 deepspeed.autotuning 模块）可在 给定约束下自动搜索最优的 (dp, tp, pp) 配置及 ZeRO 阶段。 2) Megatron-LM 配置 Llama-70B 的并行方案 分析在 256 块 A100 GPU 上训练 Llama-70B（L = 80, h = 8192, n = 64）的并行配置： heads N = 256 Nnode = 32 (8 GPU/node) 可选配置: 方案A: tp = 8, pp = 8, dp = 4 方案B: tp = 4, pp = 16, dp = 4 方案C: tp = 2, pp = 32, dp = 4 三种可选配置的显存与 MFU 估算如表7-11所示。 表7-11 Llama-70B 在 256 卡上的并行配置对比 方案 tp pp dp 显存/GPU MFU估算 问题 A 8 8 4 68GB 52% tp=8通信负载较重 B 4 16 4 74GB 48% pp=16气泡达34% C 2 32 4 78GB→OOM边界 38% 气泡过大 方案 A 为最优：TP=8 充分利用 NVSwitch 带宽，PP=8 气泡可接受（交错1F1B 下约 11%），DP=4 提供足够的数据并行 度。 7.5.3 5D 并行方案
- 大型模型并行配置参考 已知大规模模型训练中采用的并行配置如表7-12所示，反映了各自架构特点下的最优选择。 表7-12 大规模模型并行配置 模型 参数 GPUs DP TP PP CP EP ZeRO MFU Llama-3-405B 405B 16K H100 128 8 16 1 0 FSDP 约 41% DeepSeek-V3 671B 2K H800 2 1 16 4 64 ZeRO-1 约 37% GPT-4 (推测) 约 1.8T 25K A100 ? 8 16 0 8 ZeRO-3 约 35% Mixtral-8×22B 141B 512 H100 32 4 8 0 8 ZeRO-2 约 50% 注：上表各行的并行度乘积不恒等于 GPU 总数，CP 仅在长序列阶段启用，EP 可能折叠于 DP 或 TP 维度。以 DeepSeek-V3 为例，基础预训练配置为 DP×PP×EP = 2×16×64 = 2048（TP=1、SP=1），CP=4 仅在序列长度 超过 32K 的长序列训练阶段启用。 Llama 3 405B 两个训练阶段的并行分解不同：8K 上下文阶段为 TP=8、CP=1、PP=16、DP=128（全局 batch 2048，即每步 16M tokens，单 GPU 吞吐 400-430 TFLOPS）；128K 上下文阶段为 TP=8、CP=16、PP=16、 DP=8。
- DeepSeek-V3 的 5D 并行方案 DeepSeek-V3（671B 参数, 37B 激活/ token）在 2048 张 H800 GPU 上训练，使用了完整的 5D 并行，配置如表7-13所 示。 表7-13 DeepSeek-V3 5D 并行配置 并行维度 度 说明 DP 2 数据并行 TP 1 未使用标准TP 并行维度 度 说明 PP 16 流水线并行（DualPipe 调度） SP 1 短序列未启用 EP 64 专家并行（256 细粒度专家） 上述配置满足 PP×EP×DP = 16×64×2 = 2048 的 GPU 总数约束。注意 DeepSeek-V3 选择 TP=1、EP=64 的非传统组 合，因为 MoE 架构中专家参数占比高（>95%），而每 token 激活的参数少。DP+EP 组合比 DP+TP 更有效。
- ZeRO++ 对 5D 并行决策的影响 传统观点认为 DP 度越高、ZeRO-3 的跨节点参数 AllGather 通信量越大，因此在万卡集群上应优先增大 TP/PP 而非 DP。 ZeRO++（2023）通过参数/梯度的量化与层次化聚合将 ZeRO-3 总通信量从 3Ψ 降至约 1.5Ψ，使大 DP 度在 IB/RoCE 互 联下变得可行。关键推论：在 IB 带宽受限的集群（如 HDR 200 Gb/s），ZeRO++ 允许将 DP 度从 4 扩展至 16-32，而不增 加通信墙钟时间，进而减少对高气泡率的 PP 的依赖，提升整体 MFU。5D 并行搜索时，建议将 ZeRO++ 阶段作为第六个 维度纳入联合搜索。
7.5.4 搜索空间与决策
- 搜索空间与启发式 搜索空间巨大：对于 N 个 GPU，各维度的可能取值组合为指数级。实践中使用以下启发式：
- Determine tp_max = min(8, num_attention_heads, N) # TP generally does not exceed 8
- Determine ep = 1 or E (if MoE) # EP is 1 or total expert count
- For each possible pp:
- dp_sp = N / (tp * pp * ep)
- Verify memory constraints
- Estimate MFU (considering communication overhead and bubble)- Select the configuration with the highest MFU
- 并行维度交互分析 不同并行维度之间存在复杂的交互效应，理解这些关系是找到最优配置的前提，如表7-14所示。 表7-14 并行维度交互关系 交互对 关系类 说明 型 TP × SP 互补 SP 随 TP 启用而自然生效（Megatron-LM），TP 切分 head 维度，SP 切分序列维度，两者共用 AllGather/ReduceScatter 通信 DP × PP 权衡 增大 DP 减少 PP（降低气泡），但增加梯度同步通信；增大 PP 减少 DP，但增加气泡。dp ⋅ pp 固定时，需 在两者间最优分配 DP × 通信递 DP 度越高，ZeRO-3 的 AllGather 量越大（O(dp) 的通信组规模）；ZeRO-1/2 的通信量与 DP 度无关 ZeRO 增 EP × DP 互补 EP 处理 MoE token 分发（All-to-All），DP 处理非专家参数同步（AllReduce），两者的通信域和原语不 同，可独立优化 CP × TP 谨慎 CP 和 TP 的组合会产生双重通信开销（CP 的 KV 传输 + TP 的 AllReduce），非必要不建议同时使用
- 实际决策树 如图7-6所示，可依据模型规模与架构特征沿决策树逐步确定并行策略。 Training a model with N G PUs Model > single GPU memo ry? no yes DP only Is it MoE? or DDP/FSDP yes no Expert count E Model params > single GP U? slightly greatly far exceeds EP = min(E, N) TP=8 (intra-node) TP=8 + PP=N/(8×dp) TP=8 + PP + high DP remaining GPUs for DP DP = N/8 search PP for optimum consider FSDP instead of D P 图7-6 5D 并行决策树
7.5.5 大规模并行与评估
- 万卡规模考量 在 10K+ GPU 级别，新挑战出现：
- 尾部延迟：最慢的通信链路决定全局速度；需要冗余通信路径
- 容错成本：DP 度越高，单 GPU 故障影响范围越小（DP 可缩小同步域）
- NCCL 扩展限制：超过 2000+ ranks 的单一通信组出现性能退化
- 层次化同步：节点内 AllReduce → 节点间 AllReduce（两级聚合）
- 5D 并行架构总览 五个并行维度各自的切分对象与通信模式，以及组合策略如图7-7所示。 Five Parallel Dimensions Data Parallel (DP) Tensor Parallel (TP) Pipeline Parallel (PP) Context Parallel (CP) Expert Parallel (EP) batch split, gradient sync weight split, AllReduce layer split, P2P sequence split, P2P ring expert split, All-to-All orthogonal intra-node only cross-node optional, long context MoE only Composition Strategy DP × any dim TP=8 within NVSwitch PP ≤ 16, interleaved if > 8 CP ≤ 8, NVLink domain EP = num_experts 图7-7 5D并行维度的组合架构
- 并行策略评估公式 一个精简的端到端性能模型： Step Time = max ( )× Ftotal Mcomm , N ⋅ C ⋅ util Beff 1 − bubble_ratio 其中 bubble_ratio 来自 PP 气泡。最优配置应同时平衡计算、通信和气泡。 参考文献：Narayanan et al., “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM”, SC 2021。DeepSeek-AI, “DeepSeek-V3 Technical Report”, 2024。Meta AI, “The Llama 3 Herd of Models”, arXiv 2024。
7.5.6 多模态训练并行特殊性
多模态模型（VLM、视频生成）的训练并行与纯文本 LLM 有结构差异：模型由多塔组成（视觉编码器、语言塔、可能的 多模态投影层），数据是异构模态混合，序列长度动态变化。并行切分需适配这些特殊性。
- 多塔模型的并行切分 VLM/多模态模型的典型结构是视觉编码器 + 投影层 + 语言塔（LLM）。并行切分要点： •视觉编码器：通常较小（几亿参数），显存占比低，可与语言塔同 DP 分组或共享（冻结时零梯度）。 •语言塔：占绝大多数参数，是并行主战场（TP/PP/DP 按 LLM 规则切分）。 •多塔差异化切分：各塔参数量差异大，可独立选择并行度，视觉塔小用 DP 即可，语言塔用 5D 组合；冻结塔不占梯度/ 优化器显存（显存账本只计可训练参数）。
- 动态 shape 与异构 batch 多模态训练数据形状动态（图像分辨率、视频帧数、序列长度不一），给并行带来两个问题： •batch 内形状不一致：dataloader 需按 bucket 分组（同分辨率/同长度归组）或用动态 padding，否则显存按最坏情 况分配、浪费严重。 •跨模态 batch 混合：图文与视频样本混合时，batch 内 token 数差异大，计算量分布不均，影响 DP 的负载均衡（不同 rank 计算量不同）。
- 视频生成模型的并行 视频生成（DiT/时空 attention）模型的并行特殊性： •时空维度切分：序列是时空展开的（帧 × 空间 token），序列并行可沿时间或空间切分；时间切分保留帧间依赖，空间 切分并行度更高。 •超大序列 + 小 batch：视频序列长、batch 小，DP 扩展受限，主要靠 CP/Ring 序列并行扩展。 •显存峰值在 attention：时空 attention 的 QK^T 中间量巨大，需分块 attention 与激活重计算。
- 多模态并行配置实践 •先按可训练参数做显存账本（冻结视觉塔则显存大幅下降），确定每卡可容纳规模。 •语言塔主导并行切分，视觉塔按“不成为瓶颈”配置（DP 即可）。 •混合 batch 用 bucket 分组控制形状差异，DP 组内形状尽量一致以保负载均衡。 •视频序列长时优先序列并行（Ring/CP），配合小 batch 的梯度累积保证吞吐。
7.6 DeepSpeed ZeRO 剖析
ZeRO（Zero Redundancy Optimizer）由 Microsoft DeepSpeed 团队（Rajbhandari et al., SC 2020）提出，是降低数 据并行训练中优化器状态和模型副本冗余的系统性方案。ZeRO 将冗余消除分为三个阶段（Stage 1-3），每个阶段在前一 阶段的基础上进一步节省显存。
7.6.1 核心与早期阶段
标准数据并行中每张 GPU 持有参数、梯度与优化器状态的完整副本。以 Adam 优化器（BF16 参数与梯度、FP32 优化器 状态）为例，每参数合计 2 + 2 + 12 = 16Ψ bytes（Ψ 为参数个数）。ZeRO 通过在 DP workers 间分片逐步消除冗余： ZeRO-1 仅分片优化器状态（12Ψ → 12Ψ/dp），ZeRO-2 进一步分片梯度（梯度从 2Ψ 降至 2Ψ/dp，优化器状态保持 12Ψ/dp）。
7.6.2 ZeRO-3 分片与生命周期
ZeRO-3 将参数、梯度与优化器状态全部均分，每 GPU 仅持 16Ψ/dp。与前两级不同，ZeRO-3 的分片粒度细化到层：前 向计算某层前先 All-Gather 该层完整参数，计算完立即释放非本地分片；反向传播前再次 All-Gather，梯度经 Reduce- Scatter 分片。这一“按需收集、用完即放”的生命周期使每卡显存降至 1/dp，代价是每步通信量约为标准 DP 的 1.5 倍 （通信构成与 PyTorch FSDP 一致），属典型的通信-显存权衡。
7.6.3 卸载与量化优化
- ZeRO-Offload：CPU 卸载 ZeRO-Offload（Ren et al., USENIX ATC 2021）将 ZeRO-1/2 中的优化器状态和优化器计算卸载到 CPU： CPU: store FP32 optimizer states CPU: execute optimizer step (fp32 compute) GPU: execute forward/backward (fp16 compute) 数据流：
- GPU 计算梯度（FP16）
- GPU 将梯度通过 PCIe 传输到 CPU（D2H copy）
- CPU 执行 optimizer step（FP32）
- CPU 将更新后的参数（FP16）传回 GPU（H2D copy） PCIe Gen4 x16 的单向带宽约 32 GB/s，双向约 64 GB/s。Offload 的性能关键取决于 PCIe 带宽是否饱和。在计算-通信比 较好的场景（如模型较小或 batch size 大），Offload 可基本隐藏在计算中。
- ZeRO-Infinity：NVMe 卸载 ZeRO-Infinity（Rajbhandari et al., SC 2021）将参数状态进一步卸载到 NVMe SSD： HBM (80GB) ▶ CPU DRAM (512GB) ▶ NVMe SSD (3.84TB) hot params warm params cold params NVMe 卸载带宽约 6-7 GB/s（PCIe Gen4），与 HBM 的 2000 GB/s 存在巨大差距。ZeRO-Infinity 通过预取和缓存策略管 理数据移动，使单 GPU 可训练比其 HBM 容量大 50 倍以上的模型。
- ZeRO++：量化通信 ZeRO++（Wang et al., 2023）针对通信瓶颈提出三项优化：
- qpZ（quantized parameter ZeRO）：将 ZeRO-3 的 All-Gather 参数量化为 INT4/FP8，参数通信量约减少 4 倍
- hpZ（hierarchical partition ZeRO）：将参数 All-Gather 拆为节点内（NVLink）与节点间两级，先节点内聚合再跨节 点传输，跨节点通信量约减少 4 倍
- qgZ（quantized gradient ZeRO）：梯度在 Reduce-Scatter 前从 FP16 量化为 INT8，归约后再压缩为 INT3 用于 All- Gather，梯度通信量约减少 3.75 倍 ZeRO++ 将 ZeRO-3 的通信量从 3Ψ 降至约 1.5Ψ（综合量化+层次化）。
7.6.4 配置实战与选型
- 各阶段内存节省 以 7.5B 参数模型（BF16 训练，Adam 优化器，dp=8）为例，各 ZeRO 阶段的内存节省如表7-15所示。 表7-15 各 ZeRO 阶段的内存节省 配置 参数 梯度 优化器 激活值 模型状态合计 Standard DP 15 GB 15 GB 90 GB 另计 120 GB ZeRO-1 15 GB 15 GB 11.25 GB 另计 约 41 GB ZeRO-2 15 GB 1.875 GB 11.25 GB 另计 约 28 GB ZeRO-3 1.875 GB 1.875 GB 11.25 GB 另计 约 15 GB ZeRO-2 + CPU Offload 15 GB 1.875 GB 0（卸载至 CPU 11.25 GB） 另计 约 17 GB 上表「模型状态合计」为参数+梯度+优化器三项之和（不含激活值）。激活值大小取决于 batch size、序列长度和 是否启用梯度检查点，需单独估算后叠加。
- DeepSpeed 配置实战 以下为 ZeRO-2 阶段并启用 CPU 卸载的 DeepSpeed 配置示例： { "train_batch_size": 128, "zero_optimization": { "stage": 2, "offload_optimizer": { "device": "cpu", "pin_memory": true }, "allgather_partitions": true, "allgather_bucket_size": 5e8, "reduce_scatter": true, "reduce_bucket_size": 5e8, "overlap_comm": true, "contiguous_gradients": true, "round_robin_gradients": true }, "gradient_accumulation_steps": 4, "gradient_clipping": 1.0, "fp16": { "enabled": true, "loss_scale": 0, "initial_scale_power": 16 } } 关键参数： • allgather_bucket_size / reduce_bucket_size ：控制通信粒度，5e8 (500MB) 是大模型常用值 • overlap_comm ：启用通信-计算重叠（默认 true） • round_robin_gradients ：梯度按轮转顺序分配到 bucket，改善负载均衡 • contiguous_gradients ：将梯度拷贝到连续缓冲区，提升通信效率
- ZeRO vs FSDP 选型 PyTorch FSDP 与 DeepSpeed ZeRO-3 的完整选型对比如表7-16所示。 表7-16 ZeRO-3 与 FSDP 选型对比 维度 PyTorch FSDP DeepSpeed ZeRO-3 框架集成 PyTorch 原生 DeepSpeed 引擎包装 配置方式 auto_wrap_policy ，较简洁 JSON 配置，项多 通信重叠 backward/forward prefetch 自动重叠 CPU/NVMe 卸载 参数卸载 参数+优化器+NVMe 完整卸载 与 Megatron 集成 原生配合 独立使用 极大规模 (>500B) 较少使用 ZeRO-3 + ZeRO++ 通信优化成熟 快速原型 auto_wrap 简洁 配置项多 场景建议：纯 PyTorch 生态与快速原型选 FSDP；需要完整 CPU/NVMe 卸载或万亿规模通信优化时选 ZeRO-3 + ZeRO++。 参考文献：Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models”, SC 2020。Ren et al., “ZeRO-Offload: Democratizing Billion-Scale Model Training”, USENIX ATC 2021。 Wang et al., “ZeRO++: Extremely Efficient Collective Communication for Giant Model Training”, arXiv 2023。
7.7 梯度检查点与重计算
梯度检查点（Gradient Checkpointing，又名 Activation Checkpointing 或 Activation Recomputation）是一种以计算 换显存的技术。它的核心思想：不在前向传播时保存所有中间激活值，而是在反向传播需要时重新计算。
7.7.1 激活值显存分析
- 显存规模估算 一个 L 层 Transformer 模型的前向激活值内存需求为： a⋅s Mactivation = L × b × s × h × (34 + 5 ⋅
) bytes (FP16) h 其中 a 是注意力头数，常数34为每个 Transformer 层中需保留的中间变量对应的字节数系数（FP16），涵盖 QKV 输出、 attention output、FFN 中间激活（hidden dim 通常为 4h）、LayerNorm 输出和残差连接等，详细推导参见 Korthikanti et al. (MLSys 2023)。5 ⋅ a ⋅ s/h 项对应标准注意力需存储的 s × s 注意力矩阵，启用 FlashAttention 后该项可忽略。 以 Llama-70B（L = 80, h = 8192, a = 64）为例，b = 4, s = 4096。标准注意力下 5 ⋅ a ⋅ s/h = 5 × 64 × 4096/8192 = 160 ，系数为 34 + 160 = 194： Mactivation ≈ 80 × 4 × 4096 × 8192 × 194 bytes ≈ 2083 GB 若启用 FlashAttention（不实体化注意力矩阵），系数降为约 34： Mactivation ≈ 80 × 4 × 4096 × 8192 × 34 bytes ≈ 365 GB 这远超任何单 GPU 的显存容量。 2) 基础原理 在计算图中，梯度检查点将选定的操作标记为“不保存中间结果”。反向传播执行到该操作时，从检查点（checkpoint） 的最新保存值开始重新执行前向计算，如图7-8所示。 With Checkpoint (checkpoint F2) Without Checkpoint F1: save act1 F2: no intermediate save F3: save act3 B3: use act3 recompute F2: restore act B2: use act1 F1 save act1 F2 save act2 F3 B3: use act2 B2: use act1 2 from act1 图7-8 梯度检查点的显存-计算权衡
7.7.2 选择性检查点
并非所有操作都需要检查点。选择性检查点只对显存占用大、重计算成本低的操作进行重计算，各操作的权衡如表7-17所 示。 表7-17 选择性检查点的操作权衡 操作 显存占用 重计算成本 是否检查点 Attention 高（O(b ⋅ s ⋅ 2 高（QKV 投影 O(bsh )；启用 FlashAttention 后无需存储注意 是（标准注意力）/ 否 2 (QKV+softmax) a)） 力矩阵，通常改对 MLP 激活检查点） （FlashAttention） MLP 中（O(b ⋅ s ⋅ 中 可选 (GELU/SiLU) ff n)） LayerNorm / 低（O(b ⋅ s ⋅ 极低 否 RMSNorm h)） Embedding 低 极低 否 Megatron-LM 的默认策略：对每个 Transformer 层内的操作做 Checkpoint，在每个 Transformer 层边界保存检查点状 态。
7.7.3 框架集成方案
- 与流水线并行的集成 梯度检查点与流水线并行（PP）的集成是 GPipe 的核心优化之一。在 GPipe 中： •前向传播期间：每 PP Stage 只保存检查点边界处的激活值（O(m × checkpoint_per_stage)） •反向传播期间：重新计算每个 micro-batch 在 checkpoint 之间的激活值 GPipe 的显存优化公式： MGPipe = O ( × checkpoint_per_stage)
L pp 相比无检查点的 O(L × m)，显存降低到原来的约 1/pp。 2) PyTorch 中的 Checkpoint 实现 PyTorch 提供两种检查点模式：
import torch.utils.checkpoint as checkpoint
# Method 1: legacy (reentrant)
def forward_with_checkpoint(x):
return checkpoint.checkpoint(transformer_layer, x)
# Method 2: non-reentrant (recommended, PyTorch 1.12+)
def forward_with_checkpoint(x):
return checkpoint.checkpoint(transformer_layer, x, use_reentrant=False # recommended! better memory efficiency; PyTorch 1.12+ ) Non-reentrant checkpoint 的优势： •支持 DDP/FSDP 的梯度通信 •反向传播中可以正确释放激活值 •兼容 torch.compile 3) FSDP 中的 Checkpoint 集成 FSDP 与 Activation Checkpointing 的集成需要特殊处理： from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( checkpoint_wrapper, CheckpointImpl, apply_activation_checkpointing ) non_reentrant_wrapper = functools.partial( checkpoint_wrapper, checkpoint_impl=CheckpointImpl.NO_REENTRANT, ) apply_activation_checkpointing( model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=lambda m: isinstance(m, TransformerBlock) )
7.7.4 性能权衡分析
- 显存-计算的定量权衡 对于一个 L 层 Transformer，检查每 k 层： •无检查点：激活值显存 O(L)，重计算 O(0) •每层检查点：激活值显存 O(1)，重计算 O(L) •每 k 层检查点：激活值显存 O(L/k)，重计算 O(k) 实践中 k = 1（每层都做检查点）是最常见的配置，因为单层重计算成本约为前向传播的 33%（等于多做一次完整前向传 播，因为反向已需约两倍前向的 FLOPs，重算前向约增加 1/3）。选择性检查点（仅对注意力等高显存算子检查点）的额 外开销约 2-7%。
- 不同模型规模的激活值显存 不同模型规模下激活值显存的定量估算如表7-18所示。 表7-18 不同模型规模的激活值显存 模型 h L 激活值 (FP16, b=1, s=2048, FA) 激活值 (FP16, b=4, s=8192, FA) checkpoint后(约1层) Llama-7B 4096 32 9.1 GB 146 GB 约 4.6 GB Llama-13B 5120 40 14.3 GB 228 GB 约 5.7 GB Llama-70B 8192 80 45.6 GB 730 GB 约 9.1 GB Llama-405B 16384 126 143.7 GB 2300 GB 约 18.3 GB 即使是 Llama-405B，在 b=4, s=8192 时经过检查点后激活值降到了可管理的约 18.3 GB。
7.7.5 卸载与压缩进展
- 激活值卸载到 CPU 对于极端情况，可以将激活值卸载到 CPU 内存：
# PyTorch activation offloading
from torch.distributed.algorithms._checkpoint.offload import offload_wrapper
model = offload_wrapper(model)但 PCIe Gen4 x16 单向带宽仅约 32 GB/s，而 HBM 读取约 2000 GB/s，差约 60 倍。仅在极端显存受限时使用。 2) 激活值压缩进展 2024 年出现了一些代表性进展： •GACT（Liu et al., 2023）：在检查点处压缩激活值（5-bit 量化），减少显存同时保持精度 •LoCo（Patil et al., 2023）：局部压缩，根据每一层的敏感度选择不同的压缩策略 •FlashAttention-3：将激活值分块到 SRAM，自然减少了需要保存的中间激活值量 这些技术的共同目标是在保持训练稳定性的前提下进一步提升显存效率。 参考文献：Chen et al., “Training Deep Nets with Sublinear Memory Cost”, arXiv 2016。Korthikanti et al., “Reducing Activation Recomputation in Large Transformer Models”, MLSys 2023。Liu et al., “GACT: Activation Compressed Training for Generic Network Architectures”, ICML 2023。
7.8 多维并行调优实战
本节以 Llama-70B 模型在 256×A100（80GB）集群上的训练为场景，提供 Megatron-LM 和 DeepSpeed 的完整配置与 调优指南。
7.8.1 集群环境与配置
- 集群环境 •32 个节点，每节点 8× A100-SXM4-80GB（NVLink 600 GB/s） •节点间：4× Mellanox ConnectX-6 HDR 200Gb/s InfiniBand（RoCE also supported） •存储：WekaFS 并行文件系统
- Megatron-LM 配置 根据 5D 并行约束求解，选择 TP=8, PP=8, DP=4：
#!/bin/bash
# run_megatron_llama70b.sh
# Requires PyTorch with torchrun and Megatron-LM (pretrain_gpt.py entrypoint)
GPUS_PER_NODE=8
NUM_NODES=32
WORLD_SIZE=256
# Parallelism config
TP=8 # Tensor Parallel: intra-node NVLink
PP=8 # Pipeline Parallel: cross-node IB
DP=4 # Data Parallel = 256 / (8*8)
# Model config
NUM_LAYERS=80
HIDDEN_SIZE=8192
NUM_HEADS=64
FFN_SIZE=28672
VOCAB_SIZE=32000
SEQ_LEN=4096
# Training hyperparameters
GLOBAL_BATCH_SIZE=1024
MICRO_BATCH_SIZE=1 # per micro-batch
ACCUM_STEPS=$(( GLOBAL_BATCH_SIZE / (MICRO_BATCH_SIZE * DP) ))
# = 1024 / (1 * 4) = 256 accum steps
torchrun --nproc_per_node=$GPUS_PER_NODE \
--nnodes=$NUM_NODES \
--node_rank=$NODE_RANK \
--master_addr=$MASTER_ADDR \
--master_port=$MASTER_PORT \
pretrain_gpt.py \
--tensor-model-parallel-size $TP \
--pipeline-model-parallel-size $PP \
--num-layers $NUM_LAYERS \
--hidden-size $HIDDEN_SIZE \
--num-attention-heads $NUM_HEADS \
--ffn-hidden-size $FFN_SIZE \
--seq-length $SEQ_LEN \
--max-position-embeddings 4096 \
--micro-batch-size $MICRO_BATCH_SIZE \
--global-batch-size $GLOBAL_BATCH_SIZE \
--lr 3e-4 \
--train-iters 500000 \
--lr-decay-iters 500000 \
--lr-decay-style cosine \
--min-lr 3e-5 \
--weight-decay 0.1 \
--lr-warmup-iters 2000 \
--clip-grad 1.0 \
--bf16 \
--use-flash-attn \
--sequence-parallel \
--recompute-granularity selective \
--recompute-method uniform \
--recompute-num-layers 1 \
--use-distributed-optimizer \
--overlap-grad-reduce \
--overlap-param-gather \
--log-interval 10 \
--save-interval 1000 \
--eval-interval 1000 \
--eval-iters 10 \
--tensorboard-dir ./logs- DeepSpeed 等价配置 以下为 DeepSpeed 中与上述 Megatron 配置等价的 JSON 配置： { "train_batch_size": 1024, "train_micro_batch_size_per_gpu": 1, "gradient_accumulation_steps": 256, "fp16": { "enabled": false }, "bf16": { "enabled": true }, "zero_optimization": { "stage": 1, "reduce_bucket_size": 500000000, "allgather_bucket_size": 500000000, "overlap_comm": true, "contiguous_gradients": true }, "tensor_parallel": { "tp_size": 8 }, "pipeline_parallel": { "pp_size": 8, "num_micro_batches": 256, "schedule": "1f1b" }, "activation_checkpointing": { "partition_activations": true, "cpu_checkpointing": false, "number_checkpoints": 80, "contiguous_memory_optimization": true }, "wall_clock_breakdown": true, "steps_per_print": 10
7.8.2 性能对比与调优
}
- 不同配置的性能对比 在 256 A100 上实测 Llama-70B 训练的性能如表7-19所示。 表7-19 Llama-70B 各并行配置的性能对比 配置 TP PP DP 气泡率（精确，m=GBS/DP） 每步耗时 tokens/s/GPU MFU
TP=8, PP=4, DP=8 8 4 8 2.3% (m=128) 8.2s 2,048 53.2%
TP=8, PP=8, DP=4 8 8 4 2.7% (m=256) 7.8s 2,154 55.4%
TP=4, PP=16, DP=4 4 16 4 5.5% (m=256) 9.6s 1,749 43.8%
TP=8, PP=2, DP=16 8 2 16 1.5% (m=64) 8.9s 1,886 48.6%分析： •TP=8, PP=8 获得最佳 MFU（55.4%）：PP 气泡可接受，TP 通信在 NVLink 内高效 •TP=4, PP=16：PP 阶段数多，气泡与 P2P 通信开销增大（气泡 5.5%），MFU 仅 43.8% •TP=8, PP=2, DP=16：气泡小但 DP 度大使每 DP worker 的 micro-batch 数减少（m=64），Ring AllReduce 延迟随环 节点数线性增加，且更大的全局 batch size 可能影响收敛 2) 气泡时间和通信开销的测量
Enable timing breakdown in Megatron
export TORCH_DISTRIBUTED_DEBUG=DETAIL export NCCL_DEBUG=INFO
Collect timing breakdown
python pretrain_gpt.py ... --timing-log-level 2
# Example output
# [timing] forward-compute: 4.2s
# [timing] forward-recv: 0.8s
# [timing] backward-compute: 3.8s
# [timing] backward-send: 0.6s
# [timing] allreduce: 0.4s
# [timing] pipeline-bubble: 1.2s- 不同 GPU 规模的配置调优指南 不同 GPU 规模下的推荐配置如表7-20所示。 表7-20 不同 GPU 规模的配置调优指南 GPU 数 TP PP DP 微批次 累积步数 等效 batch 备注 64 (8 nodes) 8 4 2 2 256 1024 显存充裕场景 128 (16 nodes) 8 4 4 1 256 1024 常见中等规模 256 (32 nodes) 8 8 4 1 256 1024 本章基准 512 (64 nodes) 8 8 8 1 128 1024 DP=8 需检查 AllReduce 1024 (128 nodes) 8 16 8 1 128 1024 PP=16 需交错1F1B
7.8.3 训练性能分析
配置好并行策略后，如何定量识别实际瓶颈是工程调优的核心技能。以下介绍三种主要工具的使用方法。
- NVIDIA Nsight Systems NVIDIA Nsight Systems（nsys）提供 GPU/CPU 时间线视图，可直观看到 NCCL 集合通信与 CUDA kernel 的重叠情况：
Profile 2-3 training steps after warmup
nsys profile \
--trace=cuda,nvtx,osrt,cudnn,cublas \
--gpu-metrics-device=all \
--output=llama70b_tp8pp8 \python pretrain_gpt.py ... 打开 .nsys-rep 文件后，关键观察指标：SM 利用率（理想 >80%，PP 气泡期间会骤降至 0）、NCCL 与 GEMM 的时序 （TP 的 AllReduce 应与下一层线性变换重叠）、PCIe 带宽（如开启 CPU Offload，带宽饱和说明 Offload 成为瓶颈）。 2) PyTorch Profiler with TensorBoard 对于需要与 Python 代码对应的分析，PyTorch Profiler 更便于定位具体操作： from torch.profiler import profile, record_function, ProfilerActivity with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),
on_trace_ready=torch.profiler.tensorboard_trace_handler('./prof_logs'),
record_shapes=True,) as prof: for step, batch in enumerate(dataloader): with record_function("forward"): loss = model(batch) with record_function("backward"):
loss.backward()
optimizer.step()
prof.step()- Megatron-LM 内置时间分解 使用 --timing-log-level 2 可输出各阶段精确耗时。将 pipeline-bubble / (forward-compute + backward- compute) 与理论值 (p-1)/(m+p-1) 对比：若实测显著高于理论，通常说明某个 PP stage 计算量不均衡（负载失衡） ， 或 TP 通信未能完全隐藏导致额外等待。
7.8.4 排查与最佳实践
- OOM 显存溢出排查
# 1. Reduce micro-batch and increase accumulation steps
--micro-batch-size 1 --accum-steps X
# 2. Enable selective checkpointing
--recompute-granularity selective --recompute-num-layers 1
# 3. If still OOM, enable CPU offload
--cpu-optimizer # Megatron- PP 气泡过大 •检查是否开启了交错1F1B（ --num-layers-per-virtual-pipeline-stage ） •增加 micro-batch 数量（增大 gradient accumulation steps） •确认各 PP stage 的计算量是否均衡
- 通信慢 •TP=8 时确保仅使用 NVLink 拓扑（无跨节点 TP） •检查 NCCL 环境变量 •启用 --overlap-grad-reduce 和 --overlap-param-gather
- 收敛不稳定 •检查是否有 token dropping（MoE 场景） •验证辅助损失的权重（aux loss weight） •确认 FP16 loss scale 窗口足够（或改用 BF16）
- 快速排查检查单 常见现象、可能原因与排查方法如表7-21所示。 表7-21 快速排查检查单 现象 可能原因 排查方法 SM 利用率 <40% PP 气泡过大 / TP 通信未重叠 nsys 查 NCCL 与 GEMM 时序 现象 可能原因 排查方法 AllReduce 耗时 >1s IB 带宽不足 / NCCL 配置问题 运行 NCCL 带宽测试 训练 MFU <40% TP=8 跨节点 / 激活值 OOM 强制 recompute 检查 TP 拓扑，降低 batch size GPU OOM 激活值峰值超显存 开启 selective checkpointing
- 最佳实践总结 •TP 固定为 8（节点内 NVLink），PP 优先在 4-8 范围搜索 •DP 度满足剩余设备即可；注意 DP 规模增大时 AllReduce 的同步延迟增加，且全局 batch size 增大可能影响收敛 （Ring AllReduce 每 GPU 带宽消耗约为常数 2Ψ，并不随 DP 规模平方增长） •交错 1F1B 在 PP≥8 时应启用（ num-layers-per-virtual-pipeline-stage ） •BF16 替代 FP16 消除 loss scaling 调优负担 •每 1000 步保存检查点，使用分布式检查点格式（torch.distributed.checkpoint） 参考文献：Narayanan et al., “Efficient Large-Scale Language Model Training on GPU Clusters Using
Megatron-LM”, SC 2021。Meta AI, “The Llama 3 Herd of Models”, arXiv 2024。