第 11 课：5D 并行

对应官方章节：5D Parallelism in a Nutshell（含 Scope and focus、Summarizing it all）
前置：第 4～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 数量 $N_{\text{gpus}}$ 分解为 DP、TP、PP、CP 四个维度的乘积，EP 不额外占用卡数：

5D 混合并行空间分解 (DP × TP × PP × CP × EP)

总卡数 N_gpus = 1024

节点内 NVLink 域:
TP (8)

跨机 RoCE / IB 域网络:
CP(2) × PP(4) × DP(16)
(MoE 层内 EP(8) 从 DP×CP 中划分)

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

$N_{\text{gpus}} = N_{\text{dp}} \times N_{\text{tp}} \times N_{\text{pp}} \times N_{\text{cp}}$

在 Megatron-LM 中，EP 是在 MoE 层内对 DP×CP 组重新划分得到的： $N_{\text{dp}} \times N_{\text{cp}} = N_{\text{ep}} \times N_{\text{edp}}$ （ $N_{\text{edp}}$ 为专家参数的数据并行度），因此要求 $N_{\text{ep}}$ 整除 $N_{\text{dp}} \times N_{\text{cp}}$ 。本课代码取 $N_{\text{ep}} = 1$ ，乘积检查与此一致。

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

并行维度 (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/N_{\text{dp}}$ 权重，层级 GEMM 必须重建完整矩阵才能计算，算后立即销毁以保住显存上限；反向链式求导再次依赖权重故需二次拉取。	中/高（适合跨机或机内）	ZeRO-3 通信量增至 $3\Psi$ ，大集群延迟累积
TP (+SP)	隐藏维 / 序列维	参数、梯度与中间激活均分至 $1/\text{TP}$	AllReduce / AllGather + ReduceScatter	时机：前向与反向每层 Attention 与 MLP 关键计算路径上（单层前向 2 次、反向 2 次）。原因：矩阵沿隐藏维切分后，各卡输出仅为局部部分和（Partial Sum）；后续残差加法与 LayerNorm 严格依赖完整全局激活，未完成 AllReduce 前计算流无法向下推进，形成硬同步阻塞。	极高（必须机内 NVLink）	通信位于计算关键路径，严禁跨节点
PP (流水线)	网络层深度 ( $L$ )	静态权重按层数均分为 $1/p$	P2P 邻居点对点发送/接收	时机：微批次（Micro-Batch）跨越 Stage 边界流转时（前向发边界激活，反向发激活梯度）。原因：网络深度被物理切断在不同卡上；下一 Stage 必须拿到前一 Stage 边界层的输出张量才能启动前向，反向时前一 Stage 必须拿到逆向回传的激活梯度才能求导。	低（仅传层边界激活）	流水线气泡（Bubble Ratio $\approx \frac{p-1}{m}$ ）
CP (上下文)	序列长度 ( $s$ )	沿序列切分整模型激活	环形 P2P / All-to-All	时机：仅在多头自注意力（Self-Attention）计算期间高频流转；MLP/Norm/Embedding 0 通信。原因：除注意力外所有算子均为逐 Token 独立计算（Token-wise），切分序列维无跨卡依赖；唯独 Attention 打破独立性，每个 Token 必须与全序列所有 Token 点积（ $Q K^T$ ），故仅在 Attention 阶段跨卡流转 KV 块或换轴。	中/高（跨节点 RDMA 网络）	自注意力多轮通信开销
EP (专家)	MoE 专家索引 ( $E$ )	稀疏 FFN 专家参数分发	双向 All-to-All (分发+汇聚)	时机：仅在 MoE 稀疏 FFN 层触发两阶段通信（前向路由后分发、专家算完后汇聚；反向对称）。原因：不同专家参数分散在不同 GPU 上；门控网络动态判定 Token 归属后，本地卡无目标专家权重，必须通过 All-to-All 投递到专家卡（Dispatch），专家算完后再将输出汇聚回 Token 原属卡以做残差连接（Combine）。	高（受跨机双向横截面带宽制约）	专家负载失衡产生木桶短板

通信时机分类

按通信能否与计算重叠，可把上表归为四类：

关键路径阻塞：TP(+SP)。行并行输出是部分和，残差相加与 LayerNorm 必须等 AllReduce 完成，单层前向 2 次、反向 2 次，无法与后续计算重叠。
逐层或边界触发：ZeRO-3 在每层前向和反向前 All-Gather 权重、算完即释放，反向再 Reduce-Scatter 梯度；PP 在微批次跨越 Stage 边界时做 P2P 收发。
局限于特定算子：CP 只在 Attention 内交换 K/V（或做 All-to-All 换轴）；EP 只在 MoE 层做 Dispatch 与 Combine 两次 All-to-All。
可与反向重叠：DP / ZeRO-1 / ZeRO-2 只在反向阶段按桶通信，可由独立通信流与剩余反向计算重叠。

2. 硬件物理拓扑与通信分层映射法则

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

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

第一层：机内拓扑（Intra-Node NVLink / NVSwitch）：
- 物理带宽： $900 \sim 1800\text{ GB/s}$ ，端到端延迟 $< 1\mu s$ 。
- 映射原则：优先部署张量并行（TP + SP）。TP 在每个 Transformer 层的注意力与 MLP 输出端都有阻塞式集合通信，跨机走网卡会让 MFU 明显下降。因此通常设定 $\text{TP} \in \{2, 4, 8\}$ ，不跨越物理主机边界。
第二层：机间二层网络（Intra-Cluster IB / RoCE RDMA）：
- 物理带宽： $400 \sim 800\text{ Gbps}$ （折合单向约 $50 \sim 100\text{ GB/s}$ ），跨交换机延迟 $\approx 2 \sim 5\mu s$ 。
- 映射策略：适合承载 流水线并行（PP）、数据并行（DP/ZeRO-1/2） 以及 上下文并行（CP）。
  - PP 仅传输边界切片的激活值，带宽需求小；
  - DP/ZeRO 的集合通信由独立通信流执行，可完全与反向求导计算异步重叠；
  - CP 的 Ring Attention 采用双缓冲流水线，通信耗时大多可被注意力计算掩盖。
第三层：数据中心全互联横截面（East-West Spine-Leaf Fabric）：
- 面临跨机柜跨交换机路由，容易发生拥塞。
- 映射策略：承载 专家并行（EP） 时必须配合 DeepSeek 提出的“节点限制路由（Node-Limited Routing）”，把 All-to-All 流量尽量限制在少数节点内，减轻核心交换机压力。

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

常用架构： $\text{TP} \times \text{PP} \times \text{DP}$ (+ ZeRO-1)：训练千亿级稠密模型（如 Llama 3 405B）的常见组合。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 \ge 32\text{k}$ ）引入；EP 仅在引入 MoE 稀疏结构时激活，两者与其他主并行维度均保持正交。

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

并行方案选型：依次回答三个问题

1 · 序列超过 32k？

是：注意力激活随 $s^2$ 增长，需要上下文并行（CP / Ring）

2 · 单卡放得下？

放得下：纯 DP 或 ZeRO-1，吞吐最高

3 · 一台 8 卡机放得下？

放得下：机内 TP=8 + 跨机 ZeRO-3 / FSDP

放不下：机内 TP=8 + 跨机 PP + ZeRO-1/2

约束

TP 只放在 NVLink 域内
PP 只在 stage 边界传激活
用 1F1B / interleaved 调度压低气泡

图 11-2：并行方案选型的判断顺序

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

第一步：显存破局（让模型装得下）：
- 检查单卡是否能容纳模型静态状态与激活。若能，首选纯 DP / ZeRO-1，系统效率最高；
- 若单卡无法容纳，评估单机 8 卡是否能容纳：单机 8 卡放得下时，首选 $\text{ZeRO-3}$ （FSDP）或机内 $\text{TP}=8$ ；
- 超过单机 8 卡显存极限时（如 70B+ 模型），立即开启 机内 $\text{TP}=8$ + 跨机 $\text{PP}=4\sim 16$ ，必要时辅以选择性激活重计算。
第二步：吞吐配平（达到目标全局批次 GBS）：
- 业务算法决定了最佳训练全局批次 $\text{GBS}_{\text{tokens}}$ 。
- 根据公式 $\text{GBS}_{\text{tokens}} = \text{MBS} \times \text{grad\_accum} \times N_{\text{dp}} \times s$ ，在给定的 GPU 规模下倒推所需的数据并行度 $N_{\text{dp}}$ 与累积步数 $\text{grad\_accum}$ 。
第三步：气泡与带宽调优（提升算力利用率 MFU）：
- 验证流水线气泡比例 $\frac{p-1}{m}$ 是否在 10% 以内；若气泡过大，适度增大梯度累积步数 $m$ ，或开启 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 B ≈ 10.7 GB/卡（代码未按 TP/SP 分片激活，偏保守）；合计约 25 + 10.7 + 2 ≈ 38 GB < 80 GB ✓。128 卡约束下纯 DP 无法同时满足 TP=8，因为 dp·tp=1024 > 128——这就是”并行度乘积必须等于 GPU 总数”约束的直观意义。）

三个问答题要点

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-way PP，并在 DP 维上用 ZeRO-1 分片优化器状态）；
若组合 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、每卡账本 < 显存。

官方主参考

← 上一课
第 10 课：专家并行 EP 下一课 →
第 12 课：配置搜索与性能诊断

--- title: "第 11 课：5D 并行" jupyter: python3 --- - 对应官方章节：5D Parallelism in a Nutshell（含 Scope and focus、Summarizing it all） - 前置：第 4～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 数量 $N_{\text{gpus}}$ 分解为 DP、TP、PP、CP 四个维度的乘积，EP 不额外占用卡数： ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-sg" data-nid="Parallel5D" data-depth="0"><div class="fd-sg-title">5D 混合并行空间分解 (DP × TP × PP × CP × EP)</div><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-rank"><div class="fd-node fd-box" data-nid="Cluster">总卡数 N_gpus = 1024</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank fd-few"><div class="fd-node fd-box" data-nid="Node">节点内 NVLink 域:<br>TP (8)</div><div class="fd-node fd-box" data-nid="Spine">跨机 RoCE / IB 域网络:<br>CP(2) × PP(4) × DP(16)<br>(MoE 层内 EP(8) 从 DP×CP 中划分)</div></div></div></div></div></div></div><i class="fd-e" data-s="Cluster" data-t="Node" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Cluster" data-t="Spine" data-k="solid" data-h="1" data-b="0"></i></div> ``` <p class="caption" align="center"><em>图 11-1：5D 并行维度分解与物理通信层次映射架构</em></p> $$N_{\text{gpus}} = N_{\text{dp}} \times N_{\text{tp}} \times N_{\text{pp}} \times N_{\text{cp}}$$ 在 Megatron-LM 中，EP 是在 MoE 层内对 DP×CP 组重新划分得到的：$N_{\text{dp}} \times N_{\text{cp}} = N_{\text{ep}} \times N_{\text{edp}}$（$N_{\text{edp}}$ 为专家参数的数据并行度），因此要求 $N_{\text{ep}}$ 整除 $N_{\text{dp}} \times N_{\text{cp}}$。本课代码取 $N_{\text{ep}} = 1$，乘积检查与此一致。各并行维度的显存切分对象、通信原语与瓶颈对比如下（核心全景表）： | 并行维度 (Dimension) | 物理切分对象 | 显存节省本质 | 核心通信算子 | 通信触发时机与必须通信的原因 (Timing & Rationale) | 网络带宽要求与放置域 | 主要瓶颈与副作用 | |---|---|---|---|---|---|---| | **DP (数据并行)** | Batch 样本 | 减小每卡本地批次（激活显存） | AllReduce (求和) | **时机**：反向阶段（每层求导完触发分桶 AllReduce；前向及优化器步无通信）。<br>**原因**：前向各卡持有独立完整权重镜像且处理不同样本，无跨卡数据依赖；反向必须将各卡局部梯度求全局平均以保证数学单批次等价；层间梯度求导无依赖，故可与前层反向计算异步重叠。 | 低/中（跨机以太网/IB） | 受限于全局目标 GBS 上限 | | **ZeRO-1/2/3** | 优化器/梯度/权重 | 沿 DP 维切分静态状态 | Reduce-Scatter + All-Gather | **时机**：ZeRO-1/2 仅反向阶段；ZeRO-3 在前向/反向进入每层**前** All-Gather 权重并在算后立即释放（反向加发 Reduce-Scatter 梯度）。<br>**原因**：ZeRO-1/2 前向保留完整权重故 0 通信；ZeRO-3 单卡常驻仅 $1/N_{\text{dp}}$ 权重，层级 GEMM 必须重建完整矩阵才能计算，算后立即销毁以保住显存上限；反向链式求导再次依赖权重故需二次拉取。 | 中/高（适合跨机或机内） | ZeRO-3 通信量增至 $3\Psi$，大集群延迟累积 | | **TP (+SP)** | 隐藏维 / 序列维 | 参数、梯度与中间激活均分至 $1/\text{TP}$ | AllReduce / AllGather + ReduceScatter | **时机**：前向与反向每层 Attention 与 MLP 关键计算路径上（单层前向 2 次、反向 2 次）。<br>**原因**：矩阵沿隐藏维切分后，各卡输出仅为**局部部分和（Partial Sum）**；后续残差加法与 LayerNorm 严格依赖完整全局激活，未完成 AllReduce 前计算流无法向下推进，形成硬同步阻塞。 | **极高（必须机内 NVLink）** | 通信位于计算关键路径，严禁跨节点 | | **PP (流水线)** | 网络层深度 ($L$) | 静态权重按层数均分为 $1/p$ | P2P 邻居点对点发送/接收 | **时机**：微批次（Micro-Batch）跨越 Stage 边界流转时（前向发边界激活，反向发激活梯度）。<br>**原因**：网络深度被物理切断在不同卡上；下一 Stage 必须拿到前一 Stage 边界层的输出张量才能启动前向，反向时前一 Stage 必须拿到逆向回传的激活梯度才能求导。 | **低（仅传层边界激活）** | 流水线气泡（Bubble Ratio $\approx \frac{p-1}{m}$） | | **CP (上下文)** | 序列长度 ($s$) | 沿序列切分整模型激活 | 环形 P2P / All-to-All | **时机**：仅在多头自注意力（Self-Attention）计算期间高频流转；MLP/Norm/Embedding 0 通信。<br>**原因**：除注意力外所有算子均为**逐 Token 独立计算（Token-wise）**，切分序列维无跨卡依赖；唯独 Attention 打破独立性，每个 Token 必须与全序列所有 Token 点积（$Q K^T$），故仅在 Attention 阶段跨卡流转 KV 块或换轴。 | 中/高（跨节点 RDMA 网络） | 自注意力多轮通信开销 | | **EP (专家)** | MoE 专家索引 ($E$) | 稀疏 FFN 专家参数分发 | 双向 All-to-All (分发+汇聚) | **时机**：仅在 MoE 稀疏 FFN 层触发两阶段通信（前向路由后分发、专家算完后汇聚；反向对称）。<br>**原因**：不同专家参数分散在不同 GPU 上；门控网络动态判定 Token 归属后，本地卡无目标专家权重，必须通过 All-to-All 投递到专家卡（Dispatch），专家算完后再将输出汇聚回 Token 原属卡以做残差连接（Combine）。 | **高（受跨机双向横截面带宽制约）** | 专家负载失衡产生木桶短板 | #### 通信时机分类按通信能否与计算重叠，可把上表归为四类： - **关键路径阻塞**：TP(+SP)。行并行输出是部分和，残差相加与 LayerNorm 必须等 AllReduce 完成，单层前向 2 次、反向 2 次，无法与后续计算重叠。 - **逐层或边界触发**：ZeRO-3 在每层前向和反向前 All-Gather 权重、算完即释放，反向再 Reduce-Scatter 梯度；PP 在微批次跨越 Stage 边界时做 P2P 收发。 - **局限于特定算子**：CP 只在 Attention 内交换 K/V（或做 All-to-All 换轴）；EP 只在 MoE 层做 Dispatch 与 Combine 两次 All-to-All。 - **可与反向重叠**：DP / ZeRO-1 / ZeRO-2 只在反向阶段按桶通信，可由独立通信流与剩余反向计算重叠。 ### 2. 硬件物理拓扑与通信分层映射法则设计混合并行策略绝非盲目组合，其第一性原理是：**将通信频率最高、对延迟最敏感的并行维度，映射到物理带宽最高、延迟最低的硬件通信互联层级上**。现代 AI 超算集群在物理上通常具备严格的三级互联拓扑： 1. **第一层：机内拓扑（Intra-Node NVLink / NVSwitch）**： - 物理带宽：$900 \sim 1800\text{ GB/s}$，端到端延迟 $< 1\mu s$。 - **映射原则**：**优先部署张量并行（TP + SP）**。TP 在每个 Transformer 层的注意力与 MLP 输出端都有阻塞式集合通信，跨机走网卡会让 MFU 明显下降。因此通常设定 $\text{TP} \in \{2, 4, 8\}$，不跨越物理主机边界。 2. **第二层：机间二层网络（Intra-Cluster IB / RoCE RDMA）**： - 物理带宽：$400 \sim 800\text{ Gbps}$（折合单向约 $50 \sim 100\text{ GB/s}$），跨交换机延迟 $\approx 2 \sim 5\mu s$。 - **映射策略**：适合承载 **流水线并行（PP）**、**数据并行（DP/ZeRO-1/2）** 以及 **上下文并行（CP）**。 - PP 仅传输边界切片的激活值，带宽需求小； - DP/ZeRO 的集合通信由独立通信流执行，可完全与反向求导计算异步重叠； - CP 的 Ring Attention 采用双缓冲流水线，通信耗时大多可被注意力计算掩盖。 3. **第三层：数据中心全互联横截面（East-West Spine-Leaf Fabric）**： - 面临跨机柜跨交换机路由，容易发生拥塞。 - **映射策略**：承载 **专家并行（EP）** 时必须配合 DeepSeek 提出的“节点限制路由（Node-Limited Routing）”，把 All-to-All 流量尽量限制在少数节点内，减轻核心交换机压力。 ### 3. 维度间的兼容性与互斥法则 - **常用架构：$\text{TP} \times \text{PP} \times \text{DP}$ (+ ZeRO-1)**：训练千亿级稠密模型（如 Llama 3 405B）的常见组合。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 \ge 32\text{k}$）引入；EP 仅在引入 MoE 稀疏结构时激活，两者与其他主并行维度均保持正交。 ### 4. 工业界并行策略选型决策树（如**图 11-2** 所示） :::: {.vfig} ::: {.vfig-title} 并行方案选型：依次回答三个问题 ::: ::::: {.vfig-grid style="--cols:2"} ::: {.vfig-card .c-blue} **1 · 序列超过 32k？** 是：注意力激活随 $s^2$ 增长，需要上下文并行（CP / Ring） ::: ::: {.vfig-card .c-purple} **2 · 单卡放得下？** 放得下：纯 DP 或 ZeRO-1，吞吐最高 ::: ::: {.vfig-card .c-teal} **3 · 一台 8 卡机放得下？** 放得下：机内 TP=8 + 跨机 ZeRO-3 / FSDP 放不下：机内 TP=8 + 跨机 PP + ZeRO-1/2 ::: ::: {.vfig-card .c-orange} **约束** - TP 只放在 NVLink 域内 - PP 只在 stage 边界传激活 - 用 1F1B / interleaved 调度压低气泡 ::: ::::: :::: <p class="caption">图 11-2：并行方案选型的判断顺序</p> 面对全新的大模型预训练任务，资深架构师遵循标准的三步推演法则： 1. **第一步：显存破局（让模型装得下）**： - 检查单卡是否能容纳模型静态状态与激活。若能，首选纯 DP / ZeRO-1，系统效率最高； - 若单卡无法容纳，评估单机 8 卡是否能容纳：单机 8 卡放得下时，首选 $\text{ZeRO-3}$（FSDP）或机内 $\text{TP}=8$； - 超过单机 8 卡显存极限时（如 70B+ 模型），立即开启 **机内 $\text{TP}=8$ + 跨机 $\text{PP}=4\sim 16$**，必要时辅以选择性激活重计算。 2. **第二步：吞吐配平（达到目标全局批次 GBS）**： - 业务算法决定了最佳训练全局批次 $\text{GBS}_{\text{tokens}}$。 - 根据公式 $\text{GBS}_{\text{tokens}} = \text{MBS} \times \text{grad\_accum} \times N_{\text{dp}} \times s$，在给定的 GPU 规模下倒推所需的数据并行度 $N_{\text{dp}}$ 与累积步数 $\text{grad\_accum}$。 3. **第三步：气泡与带宽调优（提升算力利用率 MFU）**： - 验证流水线气泡比例 $\frac{p-1}{m}$ 是否在 10% 以内；若气泡过大，适度增大梯度累积步数 $m$，或开启 Interleaved 1F1B / Zero-Bubble 调度。 - 检查网络通信是否能被单层计算完全掩盖，最终输出全局稳态训练配置。 ## 代码填空题为一个模型 + 集群设计配置：给定约束，找出满足"装得下 + 达到 gbs"的 (tp, pp, cp, dp, zero, grad_acc)。 ```{python} #| course_role: exercise 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 为何很少组合。 ::: {.callout-tip collapse="true" title="💡 点击展开参考答案与详细代码实现" icon="false"} ## 参考答案（仅 answer 分支）先完成题目再核对；核对后必须解释关键公式，并改一个规模重新计算。 ```{python} #| course_role: solution 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']})") ``` ## 补全后的代码（关键填空处） ```python # 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 的）： ```python 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 B ≈ 10.7 GB/卡（代码未按 TP/SP 分片激活，偏保守）；合计约 25 + 10.7 + 2 ≈ 38 GB < 80 GB ✓。128 卡约束下纯 DP 无法同时满足 TP=8，因为 dp·tp=1024 > 128——这就是"并行度乘积必须等于 GPU 总数"约束的直观意义。） ## 三个问答题要点 ### 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-way PP，并在 DP 维上用 ZeRO-1 分片优化器状态）； - 若组合 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、每卡账本 < 显存。 ## 官方主参考 - [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook) - [PyTorch distributed documentation](https://pytorch.org/docs/stable/distributed.html) ::: ::: {.course-pagination} [**← 上一课**<br>第 10 课：专家并行 EP](lesson10.qmd){.pagination-link .pagination-prev} [**下一课 →**<br>第 12 课：配置搜索与性能诊断](lesson12.qmd){.pagination-link .pagination-next} :::