第 11 课:5D 并行

参考答案版:包含完整代码实现与问答题深度参考解析。建议先独立完成练习题后再展开核对。

Important统一完成标准与及格线

代码 4 分、Q1~Q3 各 2 分,通过线 8/10。必须解释正确性边界、显存/通信公式中的单位与分片维度;未实际运行的内容只能标记为静态审查。

本课目标

完成后你需要能够:

  • 列出五个并行维度及其切分的对象与维度:DP(batch)、TP(hidden)、SP/CP(sequence)、PP(layers)、EP(experts),外加 ZeRO-1/2/3。
  • 说出各技术的显存节省对象与主要瓶颈(教材总表)。
  • 按教材的三步流程为一个给定的模型 + 集群设计并行配置。
  • 解释组合规则:TP 节点内、PP/ZeRO-3 跨节点、ZeRO-1/2 与 PP 兼容、ZeRO-3 与 PP 很少组合。

核心概念

1. 5D 并行体系总览与切分维度正交性

在超大规模基础大模型(如 70B 至 1T 参数)的分布式预训练中,没有任何单一的并行技术能够独立承载海量参数、超长上下文以及庞大批次的严苛要求。 现代大模型训练系统(如 Megatron-LM、DeepSeek-V3、NeMo)通过在多维几何网格上正交组合多种技术,构建起 5D 混合并行体系。

如图 11-1 所示,集群总 GPU 数量 NgpusN_{\text{gpus}} 被形式化分解为 5 个独立维度的笛卡尔积:

flowchart TD
    subgraph Parallel5D["5D 混合并行空间分解 (DP × TP × PP × CP × EP)"]
        direction TB
        Cluster["总卡数 N_gpus = 1024"]
        Cluster --> Node["节点内 NVLink 高速域: TP (8)"]
        Cluster --> Spine["机架内 / 跨机 RoCE / IB 域: CP (2) × PP (4) × EP (8) × DP (2)"]
    end

图 11-1:5D 并行维度分解与物理通信层次映射架构

Ngpus=Ndp×Ntp×Npp×Ncp×NepN_{\text{gpus}} = N_{\text{dp}} \times N_{\text{tp}} \times N_{\text{pp}} \times N_{\text{cp}} \times N_{\text{ep}}

各并行维度的显存切分对象、通信原语与瓶颈对比如下(核心全景表):

并行维度 (Dimension) 物理切分对象 显存节省本质 核心通信算子 通信触发时机与必须通信的原因 (Timing & Rationale) 网络带宽要求与放置域 主要瓶颈与副作用
DP (数据并行) Batch 样本 减小每卡本地批次(激活显存) AllReduce (求和) 时机:反向阶段(每层求导完触发分桶 AllReduce;前向及优化器步无通信)。
原因:前向各卡持有独立完整权重镜像且处理不同样本,无跨卡数据依赖;反向必须将各卡局部梯度求全局平均以保证数学单批次等价;层间梯度求导无依赖,故可与前层反向计算异步重叠。
低/中(跨机以太网/IB) 受限于全局目标 GBS 上限
ZeRO-1/2/3 优化器/梯度/权重 沿 DP 维切分静态状态 Reduce-Scatter + All-Gather 时机:ZeRO-1/2 仅反向阶段;ZeRO-3 在前向/反向进入每层前 All-Gather 权重并在算后立即释放(反向加发 Reduce-Scatter 梯度)。
原因:ZeRO-1/2 前向保留完整权重故 0 通信;ZeRO-3 单卡常驻仅 1/Ndp1/N_{\text{dp}} 权重,层级 GEMM 必须重建完整矩阵才能计算,算后立即销毁以保住显存上限;反向链式求导再次依赖权重故需二次拉取。
中/高(适合跨机或机内) ZeRO-3 通信量增至 3Ψ3\Psi,大集群延迟累积
TP (+SP) 隐藏维 / 序列维 参数、梯度与中间激活均分至 1/TP1/\text{TP} AllReduce / AllGather + ReduceScatter 时机:前向与反向每层 Attention 与 MLP 关键计算路径上(单层前向 2 次、反向 2 次)。
原因:矩阵沿隐藏维切分后,各卡输出仅为局部部分和(Partial Sum);后续残差加法与 LayerNorm 严格依赖完整全局激活,未完成 AllReduce 前计算流无法向下推进,形成硬同步阻塞。
极高(必须机内 NVLink) 通信位于计算关键路径,严禁跨节点
PP (流水线) 网络层深度 (LL) 静态权重按层数均分为 1/p1/p P2P 邻居点对点发送/接收 时机:微批次(Micro-Batch)跨越 Stage 边界流转时(前向发边界激活,反向发激活梯度)。
原因:网络深度被物理切断在不同卡上;下一 Stage 必须拿到前一 Stage 边界层的输出张量才能启动前向,反向时前一 Stage 必须拿到逆向回传的激活梯度才能求导。
低(仅传层边界激活) 流水线气泡(Bubble Ratio ≈p−1m\approx \frac{p-1}{m})
CP (上下文) 序列长度 (ss) 沿序列切分整模型激活 环形 P2P / All-to-All 时机:仅在多头自注意力(Self-Attention)计算期间高频流转;MLP/Norm/Embedding 0 通信。
原因:除注意力外所有算子均为逐 Token 独立计算(Token-wise),切分序列维无跨卡依赖;唯独 Attention 打破独立性,每个 Token 必须与全序列所有 Token 点积(QKTQ K^T),故仅在 Attention 阶段跨卡流转 KV 块或换轴。
中/高(跨节点 RDMA 网络) 自注意力多轮通信开销
EP (专家) MoE 专家索引 (EE) 稀疏 FFN 专家参数分发 双向 All-to-All (分发+汇聚) 时机:仅在 MoE 稀疏 FFN 层触发两阶段通信(前向路由后分发、专家算完后汇聚;反向对称)。
原因:不同专家参数分散在不同 GPU 上;门控网络动态判定 Token 归属后,本地卡无目标专家权重,必须通过 All-to-All 投递到专家卡(Dispatch),专家算完后再将输出汇聚回 Token 原属卡以做残差连接(Combine)。
高(受跨机双向横截面带宽制约) 专家负载失衡产生木桶短板

核心通信时机与生命周期时序推演(Forward / Backward / Step)

理解通信时机是把握并行系统性能重叠(Overlap)与网络瓶颈的关键:

  1. 计算关键路径阻塞型(In-line Blocking):
    • TP (+SP):通信深度嵌套在单层 Transformer 的前向与反向计算图核心。RowParallel 矩阵乘法输出的是局部部分和,必须经过 AllReduce 求和后才能输入后续的残差相加与 LayerNorm,无法与后续计算异步重叠。单层前向 2 次、反向 2 次,属于最高频、对延迟最敏感的阻塞通信。
  2. 前后向阶段触发型(Per-Layer / Boundary Event):
    • ZeRO-3:将通信平摊到每一层前后。前向算到第 ll 层前必须执行 All-Gather 拿全量权重,算完立即释放显存;反向算到第 ll 层前再次 All-Gather 权重,求导完成后立即通过 Reduce-Scatter 分发该层梯度分片并再次释放权重。
    • PP:通信仅发生在微批次(Micro-Batch)跨越流水线切分点(Stage 边界)的时刻。Stage kk 算完边界层输出立即向 Stage k+1k+1 发送 P2P 激活消息,反向则回传激活梯度。
  3. 特定算子局部触发型(Operator-Specific):
    • CP:通信完全局限于自注意力(Self-Attention)模块内部。无论是 Ring Attention 的环形 KV P2P 轮转,还是 DeepSpeed Ulysses 的 Attention 前后 All-to-All 换轴,网络中的 Embedding、MLP、LayerNorm 等逐 Token 算子均在各卡本地完全独立执行,产生 0 跨卡通信。
    • EP:通信完全局限于稀疏 MoE 门控与 FFN 模块内部。在 Router 计算出各 Token 目标专家后,触发第 1 次 All-to-All(Token Dispatch);专家执行完 GEMM 计算后,触发第 2 次 All-to-All(Token Combine)汇聚输出。
  4. 反向计算异步重叠型(Backward Overlapped):
    • DP / ZeRO-1 / ZeRO-2:前向传播与优化器更新阶段完全无跨卡网络流量。通信仅在反向传播过程中触发,由独立的 CUDA 通信流(NCCL Stream)负责。通过分桶(Bucket)机制,在当前层计算反向梯度的同时,后台并发对已完成求导的前序层梯度执行 AllReduce 或 Reduce-Scatter,能够近乎 100% 掩盖通信耗时。

深度剖析:各并行维度“为什么必须在该时机通信”的数学因果链

面试中高频考查:“为什么前向有/无通信?为什么某维度必须在这个节点强同步?” 深度回答应紧扣数学公式与显存状态流转:

  • DP 为什么前向 0 通信、反向必须通信?
    • 前向:输入样本沿批次切分 B=∑rBrB = \sum_{r} B_r。由于每张卡预先同步好了完全一致的模型权重 WW,前向损失函数 Lr(W)=1|Br|∑x∈Brℓ(x,W)L_r(W) = \frac{1}{|B_r|}\sum_{x \in B_r} \ell(x, W) 在各卡本地纯闭门计算,各卡之间数学上不存在任何数据交叉。
    • 反向:全局优化目标为最小化全数据集期望损失,权重单步更新公式为 W(t+1)=W(t)−η⋅1Ndp∑r=1NdpgrW^{(t+1)} = W^{(t)} - \eta \cdot \frac{1}{N_{\text{dp}}} \sum_{r=1}^{N_{\text{dp}}} g_r。由于每张卡算出的局部梯度 grg_r 仅反映本地样本分布,必须通过 AllReduce(SUM) 求全局加和并除以 NdpN_{\text{dp}},否则各卡参数将在第一步更新后立刻发生分化(Model Divergence)。
  • TP 为什么死锁在计算关键路径(In-line Blocking)?
    • 在 RowParallel 线性层中(如自注意力输出投影 WOW_O 或 MLP 第二层 WdownW_{\text{down}}),输入被切分至各卡 XiX_i,权重按行切分 WiW_i。各卡在本地完成局部矩阵乘法:Yi=XiWiY_i = X_i W_i。
    • 数学真实输出为内积和:Y=∑i=1TPYiY = \sum_{i=1}^{\text{TP}} Y_i。此时每张卡持有的 YiY_i 仅为一个未完备的代数部分和(Partial Sum)。
    • 由于紧随其后的算子是残差相加 Z=Xresidual+YZ = X_{\text{residual}} + Y 与 LayerNorm 归一化 LN(Z)\text{LN}(Z),LayerNorm 的均值与方差统计量在数学上必须依赖完备的全局激活向量。如果不对 YiY_i 执行 AllReduce 求和,后续所有算子都将输入完全错误的数据。因此,AllReduce 必须插入在计算图关键路径中,计算流在此被强行阻塞。
  • ZeRO-3 为什么必须在前后向“反复拉取与销毁权重”?
    • 显存与计算的不可调和矛盾:ZeRO-3 的核心命题是让每张卡仅常驻 1/Ndp1/N_{\text{dp}} 的模型权重(消除 2Ψ2\Psi 冗余)。然而,GPU Tensor Core 执行矩阵乘 GEMM 时,底层指令(如 HMMA/WGMMA)物理上必须加载当前层的完整权重矩阵。
    • 时机设计:前向计算到第 ll 层时,各卡通过 All-Gather 临时拼凑出第 ll 层的全量参数,完成该层前向;算完若不立即销毁,显存将瞬间退化为 ZeRO-2。因此算完必须立即释放(Free)。
    • 反向为何再次拉取:根据链式求导法则,损失对前一层输入的梯度为 ∂ℒ∂X=∂ℒ∂Y⋅WT\frac{\partial \mathcal{L}}{\partial X} = \frac{\partial \mathcal{L}}{\partial Y} \cdot W^T。这一步对输入激活求导依然强依赖当前层的完整权重矩阵 WW!因为前向算完已将全量权重销毁,反向回传到该层时不得不第二次执行 All-Gather 重新拉取权重;求导完毕后再立即释放,并执行 Reduce-Scatter 将计算出的参数梯度分片送往各自主控卡。
  • CP 为什么能够把通信限制在 Attention 内部?
    • Transformer 的架构特性决定了:Token 间的空间交互仅发生在自注意力层(通过注意力权重矩阵 A=Softmax(QKT/d)A = \text{Softmax}(Q K^T / \sqrt{d}) 连接全序列)。
    • 一旦脱离 Attention 算子,无论是 FFN(MLP)、激活函数(SiLU/GELU)、残差连接还是 LayerNorm,全部是对单个 Token 独立进行的映射(yt=f(xt)y_t = f(x_t))。在序列切分下,每个 GPU 独立持有 SNcp\frac{S}{N_{\text{cp}}} 个 Token,计算逐 Token 算子天然完全隔离,因此通信被 100% 局限在注意力计算内部。
  • EP 为什么必须两阶段 All-to-All(Dispatch + Combine)?
    • 专家的物理空间孤岛:在混合专家模型(MoE)中,专家参数按卡切分(例如 GPU 0 存放 Expert 0/1,GPU 1 存放 Expert 2/3)。但输入 Token 是按数据并行或序列并行分布在各个卡上的。
    • 阶段一(Dispatch 分发):门控网络(Router)根据当前卡上的 Token 特征计算出路由决策(例如 Token AA 命中 Expert 2)。由于本地 GPU 0 根本没有 Expert 2 的模型权重,必须通过第 1 次 All-to-All 将 Token AA 的特征向量从 GPU 0 打包网络发送给承载 Expert 2 的 GPU 1。
    • 阶段二(Combine 汇聚):GPU 1 上的 Expert 2 执行完前向前馈计算后,产出的特征必须送回原卡 GPU 0,因为原卡保留着该微批次对应的残差连接分支 XresidualX_{\text{residual}} 和注意力激活,后续层必须在原卡上执行 X+∑Expert(X)X + \sum \text{Expert}(X)。故必须触发第 2 次 All-to-All 将专家计算结果精准送回起点。

2. 硬件物理拓扑与通信分层映射法则(面试核心红线)

设计混合并行策略绝非盲目组合,其第一性原理是:将通信频率最高、对延迟最敏感的并行维度,映射到物理带宽最高、延迟最低的硬件通信互联层级上。

现代 AI 超算集群在物理上通常具备严格的三级互联拓扑:

  1. 第一层:机内拓扑(Intra-Node NVLink / NVSwitch):
    • 物理带宽:900∼1800 GB/s900 \sim 1800\text{ GB/s},端到端延迟 <1μs< 1\mu s。
    • 映射铁律:必须优先且仅部署张量并行(TP + SP)。由于 TP 在每个 Transformer 层的注意力与 MLP 输出端都嵌有阻塞式集合通信,如果跨机走网卡,系统 MFU 将遭受毁灭性打击。因此通常设定 TP∈{2,4,8}\text{TP} \in \{2, 4, 8\},绝不跨越物理主机边界。
  2. 第二层:机间二层网络(Intra-Cluster IB / RoCE RDMA):
    • 物理带宽:400∼800 Gbps400 \sim 800\text{ Gbps}(折合单向约 50∼100 GB/s50 \sim 100\text{ GB/s}),跨交换机延迟 ≈2∼5μs\approx 2 \sim 5\mu s。
    • 映射策略:适合承载 流水线并行(PP)、数据并行(DP/ZeRO-1/2) 以及 上下文并行(CP)。
      • PP 仅传输边界切片的激活值,带宽需求小;
      • DP/ZeRO 的集合通信由独立通信流执行,可完全与反向求导计算异步重叠;
      • CP 的 Ring Attention 采用双缓冲流水线,通信耗时被注意力 GEMM 完美掩盖。
  3. 第三层:数据中心全互联横截面(East-West Spine-Leaf Fabric):
    • 面临跨机柜跨交换机路由,容易发生拥塞。
    • 映射策略:承载 专家并行(EP) 时必须配合 DeepSeek 提出的“节点限制路由(Node-Limited Routing)”,将绝大多数 All-to-All 流量限制在局部机架内,避免打崩核心交换机。

3. 维度间的兼容性与互斥法则

  • 黄金标准架构:TP×PP×DP\text{TP} \times \text{PP} \times \text{DP} (+ ZeRO-1): 这是训练千亿稠密模型(如 Llama 3 405B、GPT-4)最稳健的基石体系。TP 在机内消化单层大矩阵,PP 跨机横跨数十台服务器切分层深,DP 在最外层支撑海量批次,叠加 ZeRO-1 将优化器状态切分,达成吞吐与显存的终极平衡。

  • ZeRO-3 与 PP 极少同时使用: 两者在架构定位上高度重叠(均旨在分担模型权重)。若在同一个 PP Stage 内部再开启 ZeRO-3,会导致微批次流经该 Stage 时反复对该层权重发射 All-Gather,引发严重的冗余通信放大;工程上PP 通常只与 ZeRO-1 或 ZeRO-2 结合。

  • CP 与 EP 的即插即用性: CP 仅在输入 Context 超过单机容量时(s≥32ks \ge 32\text{k})引入;EP 仅在引入 MoE 稀疏结构时激活,两者与其他主并行维度均保持正交。

4. 工业界并行策略选型决策树(如图 11-2 所示)

flowchart LR
    Start{"模型参数量 N 与序列长度 S"} --> Branch1{"S > 32k ?"}
    Branch1 -->|是| AddCP["引入 CP (上下文并行)"]
    Branch1 -->|否| Branch2{"单卡装得下模型吗?"}
    Branch2 -->|装得下| PureDP["纯 DP / ZeRO-1 (最高吞吐)"]
    Branch2 -->|装不下| CheckMem{"NVLink 域装得下吗?"}
    CheckMem -->|是| TP_ZeRO["TP (≤8) + ZeRO-3 / FSDP"]
    CheckMem -->|否| AddPP["TP (8) + PP + ZeRO-1/2 (跨节点流水)"]

图 11-2:工业界大模型并行训练方案选型与切分决策树

面对全新的大模型预训练任务,资深架构师遵循标准的三步推演法则:

  1. 第一步:显存破局(让模型装得下):
    • 检查单卡是否能容纳模型静态状态与激活。若能,首选纯 DP / ZeRO-1,系统效率最高;
    • 若单卡无法容纳,评估单机 8 卡是否能容纳:单机 8 卡放得下时,首选 ZeRO-3\text{ZeRO-3}(FSDP)或机内 TP=8\text{TP}=8;
    • 超过单机 8 卡显存极限时(如 70B+ 模型),立即开启 机内 TP=8\text{TP}=8 + 跨机 PP=4∼16\text{PP}=4\sim 16,必要时辅以选择性激活重计算。
  2. 第二步:吞吐配平(达到目标全局批次 GBS):
    • 业务算法决定了最佳训练全局批次 GBStokens\text{GBS}_{\text{tokens}}。
    • 根据公式 GBStokens=MBS×grad_accum×Ndp×s\text{GBS}_{\text{tokens}} = \text{MBS} \times \text{grad\_accum} \times N_{\text{dp}} \times s,在给定的 GPU 规模下倒推所需的数据并行度 NdpN_{\text{dp}} 与累积步数 grad_accum\text{grad\_accum}。
  3. 第三步:气泡与带宽调优(提升算力利用率 MFU):
    • 验证流水线气泡比例 p−1m\frac{p-1}{m} 是否在 10% 以内;若气泡过大,适度增大梯度累积步数 mm,或开启 Interleaved 1F1B / Zero-Bubble 调度。
    • 检查网络通信是否能被单层计算完全掩盖,最终输出全局稳态训练配置。

代码填空题

为一个模型 + 集群设计配置:给定约束,找出满足”装得下 + 达到 gbs”的 (tp, pp, cp, dp, zero, grad_acc)。

Code
def validate_config(
    cfg: dict,                  # {"tp","pp","cp","ep","zero","mbs","grad_acc"}
    num_params: int,
    gpu_memory: int,            # bytes
    num_nodes: int,
    gpus_per_node: int,
    num_heads: int,
    target_gbs_tokens: int,
    seq_len: int,
    num_layers: int,            # 用于粗估激活(每层 ~2·seq·mbs·h,h 需另行传入或简化)
    hidden: int,
) -> tuple[bool, list[str]]:
    """检查配置是否合法。返回 (是否可行, 违规原因列表)。"""
    problems = []
    num_gpus = num_nodes * gpus_per_node

    # 1) 并行度乘积必须等于(或不超过)GPU 总数
    used = cfg["dp"] * cfg["tp"] * cfg["pp"] * cfg["cp"] * cfg["ep"]
    if used != num_gpus:
        problems.append(f"并行度乘积 {used} != GPU 总数 {num_gpus}")

    # 2) TP 必须在节点内,且不超过注意力头数
    if cfg["tp"] > gpus_per_node:
        problems.append("TP 超过单节点 GPU 数")
    if cfg["tp"] > num_heads:
        problems.append("TP 超过注意力头数")

    # 3) gbs(token)达标:mbs * grad_acc * dp * seq
    gbs = cfg["mbs"] * cfg["grad_acc"] * cfg["dp"] * seq_len
    if gbs < target_gbs_tokens:
        problems.append(f"gbs={gbs} < 目标 {target_gbs_tokens}")

    # 4) 每卡静态显存(BF16+Adam,无 FP32 梯度累积)
    # 注意:参数/梯度/优化器状态先被 TP 分片(/tp),ZeRO 项再沿 DP 分片(/dp)。
    #   zero=0:全部 /tp;zero=1:参数和梯度 /tp,优化器状态 /(dp·tp);
    #   zero=2:参数 /tp,梯度和优化器状态 /(dp·tp);zero=3:全部 /(dp·tp)
    k, psi, tp = 12, num_params, cfg["tp"]
    if cfg["zero"] == 0:
        static = (2 * psi + 2 * psi + k * psi) / tp
    elif cfg["zero"] == 1:
        static = ______                     # 填空
    elif cfg["zero"] == 2:
        static = ______                     # 填空
    else:
        static = ______                     # 填空

    # 5) 激活粗估(每层两处 s·mbs·h,BF16 2 字节;忽略 attention 二次项)
    act_per_gpu = (2 * num_layers * seq_len // cfg["cp"] * cfg["mbs"] * hidden * 2)
    total = static + act_per_gpu + 2 * 1024**3    # 预留 2 GiB 运行时/CUDA context
    if total > gpu_memory:
        problems.append(f"估算显存 {total/1e9:.1f} GB > 单卡 {gpu_memory/1e9:.0f} GB")

    return (not problems), problems

def search_configs(
    num_params: int, gpu_memory: int, num_nodes: int, gpus_per_node: int,
    num_heads: int, num_layers: int, hidden: int,
    target_gbs_tokens: int, seq_len: int, mbs: int,
    zero_stages=(0, 1, 2, 3), tp_choices=(1, 2, 4, 8), cp_choices=(1, 2, 4, 8),
    pp_choices=(1, 2, 4, 8),
):
    """在 TP/PP/CP/ZeRO 网格上搜索可行配置(dp 由总数反推,ep=1)。"""
    num_gpus = num_nodes * gpus_per_node
    found = []
    for tp in tp_choices:
        for pp in pp_choices:
            for cp in cp_choices:
                ep = 1
                if num_gpus % (tp * pp * cp * ep) != 0:
                    continue
                dp = num_gpus // (tp * pp * cp * ep)
                for zero in zero_stages:
                    # grad_acc 取满足 gbs 的最小整数
                    grad_acc = max(1, -(-(target_gbs_tokens) // (mbs * dp * seq_len)))
                    cfg = dict(tp=tp, pp=pp, cp=cp, ep=ep, zero=zero,
                               mbs=mbs, grad_acc=grad_acc, dp=dp)
                    ok, problems = validate_config(cfg, num_params, gpu_memory,
                                                   num_nodes, gpus_per_node,
                                                   num_heads, target_gbs_tokens,
                                                   seq_len, num_layers, hidden)
                    if ok:
                        found.append((cfg, problems))
    return found

if __name__ == "__main__":
    # 案例:70B 模型(h=8192, L=80, 头 64),16 节点 × 8×80GB H100,
    # 目标 gbs = 4M tokens、seq = 4096、mbs = 1
    found = search_configs(
        num_params=70_000_000_000, gpu_memory=80 * 1024**3,
        num_nodes=16, gpus_per_node=8, num_heads=64, num_layers=80,
        hidden=8192, target_gbs_tokens=4_000_000, seq_len=4096, mbs=1,
    )
    print(f"找到 {len(found)} 个可行配置;前 5 个:")
    for cfg, _ in found[:5]:
        print(f"  dp={cfg['dp']:3d} tp={cfg['tp']} pp={cfg['pp']:2d} "
              f"cp={cfg['cp']} zero={cfg['zero']} grad_acc={cfg['grad_acc']} "
              f"(并行度乘积 {cfg['dp']*cfg['tp']*cfg['pp']*cfg['cp']})")

三个问答题

Q1

为什么 TP 组必须留在节点内,而 PP/ZeRO-3 可以跨节点?请分别从”通信内容/频率”与”能否重叠”两个角度回答。CP 与 EP 在跨节点问题上各处于什么位置?

Q2

用教材三步流程为以下场景设计配置:70B 稠密模型、16 节点 × 8×H100(80 GB)、seq=4096、目标 gbs=4M tokens、mbs=1。写出你选择的 (dp, tp, pp, cp, zero, grad_acc) 与理由;至少给出一个”装得下但吞吐差”的反例配置并说明为什么差。

Q3

为什么 ZeRO-3 与 PP 很少组合,而 ZeRO-1/2 与 PP 互补(DeepSeek-V3 即 PP+ZeRO-1)?如果坚持组合 ZeRO-3+PP,ZeRO-3 应该做哪一项关键配置来减少通信浪费?

检查与通过标准

总分 10 分:代码正确 4 分(约束检查完整、gbs 公式、搜索能产出合理配置)、三题各 2 分、通过线 8 分。

一票否决项:

  • 五个维度的切分对象/维度搞混(如把 CP 说成切 batch)。
  • 认为 TP 可以跨节点扩展而没有明显代价。
  • gbs 公式漏乘 dp 或 seq。
  • 说不出 ZeRO-3 与 PP 为何很少组合。

参考答案(仅 answer 分支)

先完成题目再核对;核对后必须解释关键公式,并改一个规模重新计算。

Code
def validate_config(
    cfg: dict,                  # {"tp","pp","cp","ep","zero","mbs","grad_acc"}
    num_params: int,
    gpu_memory: int,            # bytes
    num_nodes: int,
    gpus_per_node: int,
    num_heads: int,
    target_gbs_tokens: int,
    seq_len: int,
    num_layers: int,            # 用于粗估激活(每层 ~2·seq·mbs·h,h 需另行传入或简化)
    hidden: int,
) -> tuple[bool, list[str]]:
    """检查配置是否合法。返回 (是否可行, 违规原因列表)。"""
    problems = []
    num_gpus = num_nodes * gpus_per_node

    # 1) 并行度乘积必须等于(或不超过)GPU 总数
    used = cfg["dp"] * cfg["tp"] * cfg["pp"] * cfg["cp"] * cfg["ep"]
    if used != num_gpus:
        problems.append(f"并行度乘积 {used} != GPU 总数 {num_gpus}")

    # 2) TP 必须在节点内,且不超过注意力头数
    if cfg["tp"] > gpus_per_node:
        problems.append("TP 超过单节点 GPU 数")
    if cfg["tp"] > num_heads:
        problems.append("TP 超过注意力头数")

    # 3) gbs(token)达标:mbs * grad_acc * dp * seq
    gbs = cfg["mbs"] * cfg["grad_acc"] * cfg["dp"] * seq_len
    if gbs < target_gbs_tokens:
        problems.append(f"gbs={gbs} < 目标 {target_gbs_tokens}")

    # 4) 每卡静态显存(BF16+Adam,无 FP32 梯度累积)
    # 注意:参数/梯度/优化器状态先被 TP 分片(/tp),ZeRO 项再沿 DP 分片(/dp)。
    #   zero=0:全部 /tp;zero=1:参数和梯度 /tp,优化器状态 /(dp·tp);
    #   zero=2:参数 /tp,梯度和优化器状态 /(dp·tp);zero=3:全部 /(dp·tp)
    k, psi, tp = 12, num_params, cfg["tp"]
    if cfg["zero"] == 0:
        static = (2 * psi + 2 * psi + k * psi) / tp
    elif cfg["zero"] == 1:
        static = (2 * psi + 2 * psi) / tp + (k * psi) / (cfg["dp"] * tp)                     # 填空
    elif cfg["zero"] == 2:
        static = (2 * psi) / tp + (2 * psi + k * psi) / (cfg["dp"] * tp)                     # 填空
    else:
        static = (2 * psi + 2 * psi + k * psi) / (cfg["dp"] * tp)                     # 填空

    # 5) 激活粗估(每层两处 s·mbs·h,BF16 2 字节;忽略 attention 二次项)
    act_per_gpu = (2 * num_layers * seq_len // cfg["cp"] * cfg["mbs"] * hidden * 2)
    total = static + act_per_gpu + 2 * 1024**3    # 预留 2 GiB 运行时/CUDA context
    if total > gpu_memory:
        problems.append(f"估算显存 {total/1e9:.1f} GB > 单卡 {gpu_memory/1e9:.0f} GB")

    return (not problems), problems

def search_configs(
    num_params: int, gpu_memory: int, num_nodes: int, gpus_per_node: int,
    num_heads: int, num_layers: int, hidden: int,
    target_gbs_tokens: int, seq_len: int, mbs: int,
    zero_stages=(0, 1, 2, 3), tp_choices=(1, 2, 4, 8), cp_choices=(1, 2, 4, 8),
    pp_choices=(1, 2, 4, 8),
):
    """在 TP/PP/CP/ZeRO 网格上搜索可行配置(dp 由总数反推,ep=1)。"""
    num_gpus = num_nodes * gpus_per_node
    found = []
    for tp in tp_choices:
        for pp in pp_choices:
            for cp in cp_choices:
                ep = 1
                if num_gpus % (tp * pp * cp * ep) != 0:
                    continue
                dp = num_gpus // (tp * pp * cp * ep)
                for zero in zero_stages:
                    # grad_acc 取满足 gbs 的最小整数
                    grad_acc = max(1, -(-(target_gbs_tokens) // (mbs * dp * seq_len)))
                    cfg = dict(tp=tp, pp=pp, cp=cp, ep=ep, zero=zero,
                               mbs=mbs, grad_acc=grad_acc, dp=dp)
                    ok, problems = validate_config(cfg, num_params, gpu_memory,
                                                   num_nodes, gpus_per_node,
                                                   num_heads, target_gbs_tokens,
                                                   seq_len, num_layers, hidden)
                    if ok:
                        found.append((cfg, problems))
    return found

if __name__ == "__main__":
    # 案例:70B 模型(h=8192, L=80, 头 64),16 节点 × 8×80GB H100,
    # 目标 gbs = 4M tokens、seq = 4096、mbs = 1
    found = search_configs(
        num_params=70_000_000_000, gpu_memory=80 * 1024**3,
        num_nodes=16, gpus_per_node=8, num_heads=64, num_layers=80,
        hidden=8192, target_gbs_tokens=4_000_000, seq_len=4096, mbs=1,
    )
    print(f"找到 {len(found)} 个可行配置;前 5 个:")
    for cfg, _ in found[:5]:
        print(f"  dp={cfg['dp']:3d} tp={cfg['tp']} pp={cfg['pp']:2d} "
              f"cp={cfg['cp']} zero={cfg['zero']} grad_acc={cfg['grad_acc']} "
              f"(并行度乘积 {cfg['dp']*cfg['tp']*cfg['pp']*cfg['cp']})")

补全后的代码(关键填空处)

# 4) 每卡静态显存(BF16+Adam,无 FP32 梯度累积)——先 TP 分片,ZeRO 项再沿 DP 分片
    if cfg["zero"] == 0:
        static = (2 * psi + 2 * psi + k * psi) / tp
    elif cfg["zero"] == 1:
        static = (2 * psi + 2 * psi) / tp + (k * psi) / (dp * tp)
    elif cfg["zero"] == 2:
        static = (2 * psi) / tp + (2 * psi + k * psi) / (dp * tp)
    else:
        static = (2 * psi + 2 * psi + k * psi) / (dp * tp)

    # 3) gbs(token)达标
    gbs = cfg["mbs"] * cfg["grad_acc"] * cfg["dp"] * seq_len

示例输出(70B、16 节点 × 8×80GB = 128 卡,gbs=4M、seq=4096、mbs=1)部分可行配置(共 60 个;纯 ZeRO-3 类配置也大量可行,这里只列 TP=8 的):

dp= 16 tp=8 pp=1 cp=1 zero=2 grad_acc=62   并行度乘积 128
 dp= 8 tp=8 pp=1 cp=2 zero=2 grad_acc=123  并行度乘积 128
 dp= 8 tp=8 pp=2 cp=1 zero=3 grad_acc=123  并行度乘积 128

(每种配置下静态账本:zero=2 时 (2Ψ)/8 + (2+12)Ψ/(dp·8)。70B、dp=16、tp=8:140/8 + 980e9/(16·8) = 17.5 + 7.7 ≈ 25 GB/卡;激活 ≈ 2·80·4096·1·8192·2/8 ≈ 1.3 GB/卡量级;加 2 GiB 余量 < 80 GB ✓。128 卡约束下纯 DP 无法同时满足 TP=8,因为 dp·tp=1024 > 128——这就是”并行度乘积必须等于 GPU 总数”约束的直观意义。)

(每种配置下静态账本:zero=2 时 (2Ψ)/8 + (2+12)Ψ/(dp·8)。70B、dp=128、tp=8:140/8 + 980e9/1024 ≈ 18.5 GB/卡;激活 ≈ 2·80·(4096)·1·8192·2/8? ——TP 分片后每层激活 (s, h/tp)≈1.3 GB 量级;加 2 GiB 余量 < 80 GB ✓。)

三个问答题要点

Q1

  • TP:部分和同步嵌在每层的计算链里(all-reduce 之后才有 LayerNorm/下一层),无法重叠 → 必须最高带宽(NVLink)→ 节点内;
  • PP:只传层边界激活,带宽需求低 → 跨节点可接受;ZeRO-3:参数 all-gather 可 prefetch 隐藏 → 跨节点可接受(但 DP>512 时隐藏不住);
  • CP:K/V 交换集中在 attention,可与计算重叠,跨节点可行(带宽需求居中,Ring 模式与实现相关);
  • EP:all-to-all 路由,跨节点可行,但通信量随路由分散度增长(DeepSeek 用 M 节点限制约束)。

Q2

推荐:dp=16, tp=8, pp=1, cp=1, zero=2, grad_acc=62(或 dp=8, tp=8, pp=2 变体):

  • TP=8 吃满单节点 NVLink(128 卡 = 16 节点,TP=8 → dp 最多 16);静态账本 (2Ψ)/8 + (2+12)Ψ/(dp·8) = 17.5 + 7.7 ≈ 25 GB/卡;
  • 激活:mbs=1、seq=4096、80 层 → 每层 ~2·4096·1·8192·2B ≈ 134 MB(TP 分片后更小),合计 ~10 GB;总账本 < 80 GB ✓;
  • gbs = 1×62×16×4096 ≈ 4.06M ≥ 4M ✓(grad_acc 偏大是 128 卡下同时满足 gbs 与 TP=8 的代价;若换 zero=3 可减少静态内存但通信更重);
  • 反例:dp=128, tp=1, pp=1, zero=0 纯 DP——70B×16B ≈ 1120 GB/卡装不下;或 tp=16(两节点)——跨节点 TP,吞吐骤降(“装得下但慢”)。

Q3

  • ZeRO-3 与 PP 都沿模型深度做通信(传权重 vs 传激活),叠加后每步通信翻倍,需要极大 gbs 才摊得开 → 很少组合;ZeRO-3 本身可以覆盖 PP 的分片需求(模型无关),PP 的收益(跨节点低带宽)对 ZeRO-3 是减法不是加法;
  • ZeRO-1/2 只分片优化器状态/梯度,与 PP 的激活流正交 → 互补;DeepSeek-V3 = PP + ZeRO-1(16 节点流水线 + 128 卡 DP 分片优化器状态);
  • 若组合 ZeRO-3+PP:ZeRO-3 应在 PP 的一串微批次期间保持权重驻留内存(不要逐 micro-batch all-gather/drop),把通信次数从 2·gas 压到 2(教材原话)。

通过要点

  • 五维口诀:DP(batch)、TP(hidden)、SP/CP(sequence)、PP(layers)、EP(experts),ZeRO 沿 DP 维分片状态;
  • 每卡静态账本公式要同时含 TP 与 DP/ZeRO 因子:先 /tp(TP 分片),ZeRO 项再 /dp;
  • 组合优先级:先装下(TP/PP/ZeRO-3)→ 再凑 batch(dp/grad_acc/CP)→ 最后调吞吐;
  • 配置合法性:并行度乘积 = GPU 总数、TP ≤ 节点卡数且 ≤ 头数、gbs = mbs·grad_acc·dp·seq、每卡账本 < 显存。

官方主参考