第 4 课：数据并行 DP 与 DDP

对应官方章节：Data Parallelism（含 Revisiting global batch size、Our journey up to now）
前置：第 2 课（梯度累积）、第 3 课（all-reduce）
状态：未开始

本课目标

完成后你需要能够：

把 DP 表述为”并行版梯度累积”，写出 gbs = mbs × grad_acc × dp。
解释为什么梯度 all-reduce 必须在 optimizer step 之前完成，但可以与 backward 重叠。
实现并解释两个优化：逐参数 hook 重叠、梯度 bucketing。
解释 no_sync 在”DP + 梯度累积”中的正确性作用。
说明 DP 的边界：不降低每卡静态显存，规模过大（512+）时效率下降。

核心概念

1. 数据并行（DP/DDP）的数学本质与并行语义

数据并行（Data Parallelism）在数学本质上是空间维度展开的梯度累积。如图 4-1 与图 4-2 所示，整个集群拥有 $N$ 个并行工作进程（Rank 0 至 Rank $N-1$ ），每个 GPU 显存内都完整存放了一份同构的模型参数副本 $W$ ：

Ring AllReduce = Reduce-Scatter + All-Gather

阶段 1 · Reduce-Scatter

梯度切成 $P$ 块，沿环走 $P-1$ 步
每步发出 1 块、收到 1 块并累加
结束时每卡持有 1 块的全局和

阶段 2 · All-Gather

再走 $P-1$ 步，转发已求和的块
每步用收到的块覆盖本地旧值
结束时每卡都有完整的全局和

单卡发送量 $2\frac{P-1}{P}M < 2M$ ，几乎不随卡数 $P$ 增长；耗时 $T \approx 2(P-1)\alpha + 2\frac{P-1}{P}\cdot\frac{M}{\beta}$ （ $\alpha$ 为单步延迟， $\beta$ 为带宽）

图 4-1：Ring AllReduce 的两个阶段与单卡通信量

数据并行架构 (DDP)

全局 Batch 训练数据

分片 0 (Rank 0)

分片 1 (Rank 1)

分片 2 (Rank 2)

卡 0: 模型完整副本 W

卡 1: 模型完整副本 W

卡 2: 模型完整副本 W

本地梯度 g_0

本地梯度 g_1

本地梯度 g_2

Ring AllReduce 全局同步

各卡同步更新同构参数 W_new

图 4-2：数据并行（DDP）同构副本前向与全量梯度归约架构

当一个全局 Batch（Global Batch Size, GBS）的数据输入系统时，各 Rank 分别取走互不重叠的数据子集（Micro Batch, MBS）。由于各 Rank 上的前向传播与反向传播计算完全基于各自的本地数据进行，彼此之间没有任何数据依赖，因此计算阶段具备天然的完全并行性。

然而，优化器更新全局参数时，必须基于全局样本的真实平均梯度。设 $\text{GBS} = \text{MBS} \times \text{grad\_accum} \times \text{DP}$ ，全局目标损失对参数 $\theta$ 的梯度定义为：

$\nabla_\theta \mathcal{L}_{\text{global}} = \frac{1}{\text{GBS}} \sum_{i=1}^{\text{GBS}} \nabla_\theta \ell(x_i; \theta) = \frac{1}{\text{DP}} \sum_{r=0}^{\text{DP}-1} \left( \frac{1}{\text{MBS} \times \text{grad\_accum}} \sum_{m=1}^{\text{grad\_accum}} \sum_{j=1}^{\text{MBS}} \nabla_\theta \ell(x_{r, m, j}; \theta) \right)$

由此可见，每个 Rank 独立求导得到的仅是其本地分片的梯度之和。在执行优化器更新（optimizer.step()）之前，必须通过第 3 课所讲的 AllReduce（求和后归一化除以 DP 规模） 算子，在所有参与数据并行的节点间同步梯度。这保证了所有副本在步结时取得完全相同的新权重 $W_{\text{new}}$ 。

训练规模的核心计量维度公式为： $\text{GBS} = \text{MBS} \times \text{grad\_accum} \times \text{DP}$ $\text{Batch Tokens} = \text{GBS} \times \text{seq\_len} = \text{MBS} \times \text{grad\_accum} \times \text{DP} \times \text{seq\_len}$

在维持全局 Batch 大小不变的前提下，增大 $\text{DP}$ 规模可以直接等比削减串行累积步数 $\text{grad\_accum}$ ，从而大幅缩短每步训练的墙钟时间。

2. 朴素实现的瓶颈：计算与通信完全串行

在最原始的 DistributedDataParallel 朴素实现中，执行流程遵循严格的阶段式阻塞顺序：

前向与反向完全执行完毕：所有网络层反向计算结束，显存中产生所有参数的梯度张量。
一次性阻塞式 AllReduce：启动 NCCL 通信，将全模型参数量 $\Psi$ 对应的全部梯度张量进行同步传输。
优化器更新：等待通信全部完成后，统一更新权重。

这种串行逻辑存在巨大的算力浪费。其一步总耗时为： $T_{\text{step}} = T_{\text{fwd}} + T_{\text{bwd}} + T_{\text{comm}} + T_{\text{opt}}$

在 Ring AllReduce 模式下，传输全模型梯度的网络通信时间为 $T_{\text{comm}} \approx 2 \times \frac{\Psi \times \text{sizeof(dtype)}}{B_{\text{bus}}}$ 。在万卡集群或跨节点低速网络（如普通以太网）下， $T_{\text{comm}}$ 往往占据整步耗时的 30%~60%。在通信期间，数十万 TFLOPS 的 Tensor Core 处于完全空闲状态，吞吐极度低下。

3. 优化一：梯度后向钩子与通信隐藏（Post-Accumulate Grad Hook）

要消除空闲等待，核心思路是实现计算与通信的异步重叠（Overlap）。

注意到自动微分（Autograd）的反向传播是反向有向无环图（DAG）的倒序执行过程：从最后一层（输出层 Layer $L$ ）逐层反向推进到第一层（输入层 Layer $1$ ）。当输出层 Layer $L$ 的参数求导计算完毕时，其梯度 $\nabla W_L$ 的数值已经完全确定，绝不会再受到后续浅层网络求导计算的影响。

现代 PyTorch DDP 内部通过为每个可训练参数注册底层回调（如 register_post_accumulate_grad_hook）实现通信重叠：

当反向计算图推进并生成参数 $W_k$ 的梯度时，Hook 立即被触发。
DDP 自动将该梯度的 AllReduce 异步投递到独立的 CUDA 通信流（nccl_stream）中执行。
此时，默认的计算流（compute_stream）无须任何等待，继续向前推进计算 Layer $k-1$ 、Layer $k-2$ 的反向求导。
只要浅层的反向计算耗时大于等于深层梯度的网络传输耗时，绝大部分通信时间就被反向计算完全掩盖（ $T_{\text{comm}}$ 隐藏在 $T_{\text{bwd}}$ 背后）。

如图 4-3 所示，整个反向计算与网络传输在时间线上呈现高度交叠的双轨并发执行：

反向求导 (Backward Compute)

梯度分桶 (Bucket 0 / Bucket 1)

异步通信流 (Ring AllReduce)

PyTorch DDP 计算与通信重叠流水线

1逐层反向求导 (Layer N)

2梯度写入，达到阈值 (如 25MB)

3异步触发 AllReduce (Bucket 0)

4继续求导上一层 (Layer N-1, N-2)

环形通信与计算并行执行

5Bucket 0 同步完成

6异步触发 AllReduce (Bucket 1)

7全部 Bucket 同步完毕

图 4-3：PyTorch DDP 梯度分桶与异步 AllReduce 计算通信重叠流水

4. 优化二：梯度分桶策略（Gradient Bucketing）

虽然逐参数挂载 Hook 理论上能实现最大程度的重叠，但在实际大模型中直接对每个小参数张量单独发起 AllReduce 会引发严重的性能倒退：

小包延迟开销爆炸：神经网络中包含大量极小的权重和偏置张量（例如 LayerNorm 的 scale/bias 仅数 KB）。根据第 3 课的网络延迟模型 $T = 2(N-1)\alpha + 2\frac{N-1}{N}\frac{K}{B}$ ，当张量尺寸 $K$ 很小时，耗时完全由每步通信固有的延迟基数 $\alpha$ 主导。成百上千次小包 AllReduce 会导致网络传输极其低效。
CUDA 算子发射开销（Kernel Launch Overhead）：CPU 端频繁向 GPU 发射小张量通信算子会迅速跑满 CPU，导致 GPU 发生发射等待饥饿。

PyTorch DDP 采用分桶机制（Bucketing）进行折中处理：

按照网络模型反向求导完成的逆序方向，预先分配若干块连续的扁平内存切片（称为 Bucket，默认上限通常为 bucket_cap_mb = 25MB）。
各层的梯度产生后，先被直接拷贝（Copy）到对应的连续 Bucket 内存块中。
当一个 Bucket 被填满，或者倒序遍历到了该桶管辖的最后一层梯度时，系统才对这块规整的 25MB 大张量集中发射一次异步 Ring AllReduce。
这种策略既规避了小包通信的延迟惩罚，又避免了一次性同步全量模型带来的漫长等待，实现了网络总线利用率与重叠潜力的最优权衡。

5. 优化三：与梯度累积深度协同（`no_sync` 机制）

在 DP + 梯度累积时，DDP 默认会在每个 Micro-Batch 的反向阶段都触发分桶 AllReduce。以 $\text{grad\_accum} = 8$ 为例，前 7 次同步的结果只会被下一次本地累加覆盖，每步多付出 $7 \times 2\Psi$ 的通信量。

由于梯度求和满足加法结合律： $\sum_{r=0}^{\text{DP}-1} \left( \sum_{m=1}^{\text{grad\_accum}} g_{r, m} \right) = \sum_{m=1}^{\text{grad\_accum}} \left( \sum_{r=0}^{\text{DP}-1} g_{r, m} \right)$

因此可以先在每张卡本地累加各 Micro-Batch 的梯度，只在最后一次反向传播时做一次 AllReduce。PyTorch DDP 用上下文管理器 model.no_sync() 实现这一点：在 with model.no_sync(): 块内反向只做本地累加、不触发 AllReduce；最后一个 Micro-Batch 在块外执行，触发一次同步。通信次数因此降为原来的 $1 / \text{grad\_accum}$ ，这是多卡 DP + 梯度累积的标准做法。

6. 工程细节与 DDP 的物理扩展边界

理解 DDP 的边界是掌握后续进阶并行架构（ZeRO、TP、PP）的前提：

通信内存开销：bucket_cap_mb = 25MB 是单个桶的上限，而全部 Bucket 合起来要容纳所有参数的梯度，即额外一份完整梯度副本（ $\Psi \times \text{sizeof(dtype)}$ ，7B BF16 约 14 GB），并非只多 25MB～50MB。设置 gradient_as_bucket_view=True 后，param.grad 直接作为 Bucket 的视图，可以省掉这份复制。
DDP 完全无法降低每卡的静态显存：这是纯数据并行最致命的物理局限。每个 DDP 副本的 GPU 显存内都必须完整驻留 $1\times$ 模型参数 $\Psi$ 、 $1\times$ 梯度 $\Psi$ 以及 $1\times$ 优化器状态（FP32 Master Weights + Momentum + Variance 共 $12\Psi$ 字节）。
- 对于一个 70B 参数的 BF16 大模型，其静态状态显存高达 $70 \times 16\text{ GB} = 1120\text{ GB}$ 。
- 即使单卡显存配置到顶级的 80GB/192GB，并将 Micro-Batch 降为 1 且开启全激活重计算，单卡也根本无法容纳模型静态权重，程序将瞬间触发 OOM 崩溃。
网络扩展性极限（Ring Latency 惩罚）：
- 当节点规模扩展到数百乃至数千卡（如 $\text{DP} \ge 512$ ）时，Ring AllReduce 的延迟项 $2(N-1)\alpha$ 随着环形步数呈线性膨胀。跨交换机（Spine-Leaf 网络二跳/三跳）的延迟显著拉长。
- 当通信延迟超过了反向计算耗时时，异步重叠完全失效，算力利用率暴跌。
- 此时必须联合张量并行（TP，将计算切分到机内 NVLink 超高带宽域）以及流水线并行（PP，将层深切分），构成多维混合并行系统。

具体演示与数值推演

假设训练目标设定为全局吞吐 $\text{BST} = 4\text{M tokens}$ （即 $4,194,304$ tokens），上下文长度 $\text{seq} = 4096$ ，则每步所需全局样本数为： $\text{GBS} = \frac{4,194,304}{4096} = 1024 \text{ samples}$

单卡在兼顾激活显存与算力饱和度的情况下，测定最大单卡批次 $\text{MBS} = 2$ ：

128 卡集群（ $\text{DP} = 128$ ）： $\text{grad\_accum} = \frac{\text{GBS}}{\text{MBS} \times \text{DP}} = \frac{1024}{2 \times 128} = 4$ 每卡连续计算 4 个 Micro-Batch（前 3 个使用 no_sync 纯本地累加，第 4 个触发异步 AllReduce），然后执行一次全局权重更新。
512 卡集群（ $\text{DP} = 512$ ）： $\text{grad\_accum} = \frac{1024}{2 \times 512} = 1$ 此时不需要任何梯度累加，每个 Micro-Batch 直接就是全局步，训练速度提升近 4 倍。
扩展性瓶颈突显：在 512 卡场景下，反向计算时间缩短为单个 MBS=2 的计算耗时（假设为 80ms），但 512 卡跨节点 Ring AllReduce 的延迟开销与通信排队时间大幅增加（例如通信耗时上升至 120ms）。由于 $T_{\text{comm}} > T_{\text{bwd}}$ ，异步重叠出现暴露窗口，GPU 产生 40ms 的通信气泡（Bubble）。继续无脑堆叠 DP 规模将陷入严重的边际收益递减。

代码填空题

补全以下程序：梯度分桶 + 朴素/重叠两种方案的完成时间模拟。

Code

from dataclasses import dataclass

@dataclass
class Bucket:
    names: list[str]      # 桶内梯度名（按反向完成顺序）
    total_bytes: int

def bucket_gradients(
    grads: list[tuple[str, int]],   # [(参数名, 梯度字节数)]，按反向完成顺序（最后一层在前）
    bucket_bytes: int = 25 * 1024 * 1024,   # 默认桶大小 25 MB
) -> list[Bucket]:
    """
    贪心分桶：尽量把连续的梯度装进一个桶，直到装不下就开新桶。
    单个梯度超过桶上限时，单独成桶（不拆分）。
    """
    buckets: list[Bucket] = []
    cur_names: list[str] = []
    cur_bytes = 0

    for name, size in grads:
        if size > bucket_bytes:
            # 超大梯度：先把当前桶收尾，再单独成桶
            if cur_names:
                buckets.append(Bucket(cur_names, cur_bytes))
                cur_names, cur_bytes = [], 0
            buckets.append(Bucket([name], size))
        elif ______:      # 填空：判断当前桶是否还装得下
            cur_names.append(name)
            cur_bytes += size
        else:
            # 装不下：封桶，开新桶
            buckets.append(Bucket(cur_names, cur_bytes))
            cur_names, cur_bytes = ______   # 填空：新桶初始状态
            cur_names.append(name)
            cur_bytes += size

    if cur_names:
        buckets.append(Bucket(cur_names, cur_bytes))
    return buckets

def simulate_dp_step(
    layer_backward_s: list[float],    # 每层 backward 时长（顺序：最后一层 → 第一层）
    bucket_comm_s: list[float],       # 每桶 all-reduce 时长（顺序与分桶一致）
    overlap: bool,
) -> float:
    """
    模拟一次训练步的时间。

    - 反向是串行的：层 i 的 backward 必须等层 i-1 完成。
    - 每桶的通信可以在"该桶最后一层 backward 完成"后立即开始；
      多个桶的通信在同一通信流上串行排队。
    - naive（overlap=False）：先全部 backward，再全部通信。
    - overlapped（overlap=True）：通信在释放后立即开始，尽量与后续 backward 重叠。
    """
    # 每层 backward 的累计完成时刻
    finish = []
    acc = 0.0
    for t in layer_backward_s:
        acc += t
        finish.append(acc)

    if not overlap:
        # naive：总时间 = 全部 backward + 全部通信
        return ______

    # overlapped：通信任务按释放时间排队
    # 释放时刻 = 该桶所属最后一层 backward 完成时刻。
    # 简化模型：把层分成 len(bucket_comm_s) 段，每段一层或均分；
    # 桶 b 的释放时刻 = 该段最后一层的 finish。
    release = []
    seg = max(1, len(layer_backward_s) // len(bucket_comm_s))
    for b in range(len(bucket_comm_s)):
        end_layer = min(len(layer_backward_s) - 1, (b + 1) * seg - 1)
        release.append(finish[end_layer])

    # 通信流上串行执行：t = max(当前时间, 释放时刻) + 通信时长
    t_comm = 0.0
    for b in range(len(bucket_comm_s)):
        t_comm = ______(t_comm, release[b]) + bucket_comm_s[b]

    # 总时间 = max(backward 全部完成时刻, 通信流结束时刻)
    return max(finish[-1], t_comm)

def gbs_tokens(mbs: int, grad_acc: int, dp: int, seq: int) -> int:
    """全局 batch 大小（token 数）。"""
    return ______

if __name__ == "__main__":
    # 例：30 层，每层 backward 10ms；分 6 桶，每桶通信 12ms
    layer_t = [0.010] * 30
    bucket_comm = [0.012] * 6

    t_naive = simulate_dp_step(layer_t, bucket_comm, overlap=False)
    t_over = simulate_dp_step(layer_t, bucket_comm, overlap=True)
    print(f"naive = {t_naive * 1e3:.0f} ms, overlapped = {t_over * 1e3:.0f} ms")

    # 分桶统计
    grads = [(f"p{i}", 2 * 1024 * 1024) for i in range(40)]  # 每梯度 2MB
    buckets = bucket_gradients(grads, bucket_bytes=10 * 1024 * 1024)
    print("桶数:", len(buckets), "（不分组会是 40 次通信）")
    print("每桶字节:", [b.total_bytes // (1024 * 1024) for b in buckets])

    # 全局 batch 示例
    print("gbs_tokens:", gbs_tokens(mbs=2, grad_acc=4, dp=128, seq=4096))

三个问答题

Q1

为什么梯度同步必须在 optimizer step 之前完成？为什么它可以（部分）与 backward 重叠，而第 6 课将看到 TP 的通信”在关键路径上、难以重叠”？两者有什么本质区别？

Q2

DP + 梯度累积时，为什么可以用 no_sync() 只在最后一次 backward 后 all-reduce？请从运算性质（结合律）说明正确性，并定量比较：grad_acc=8、dp=64 时，“每次都同步”与”只同步一次”的通信次数差多少倍？

Q3

教材说 DP 在 512+ 卡时效率明显下降，且 DP 不降低每卡显存。请解释这两个限制分别来自哪里，并说明：一个 70B 模型在 8×80GB 节点上训练，为什么纯 DP（含 grad_acc）无法解决”模型装不下”的问题？此时应该转向哪些技术（给出名字即可）？

检查与通过标准

总分 10 分：代码正确 4 分（分桶条件、时间模拟、gbs 公式）、三题各 2 分、通过线 8 分。

一票否决项：

认为 DP 可以降低每卡参数/优化器显存。
认为”每次 backward 都同步”与”最后一次同步”在结果上不等价。
不知道 all-reduce 是求和（平均需要再除以 dp）或同步时机错误。
声称”DP 规模可以无限扩展”或”DP 能解决模型放不下单卡的问题”。

💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

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

Code

from dataclasses import dataclass

@dataclass
class Bucket:
    names: list[str]      # 桶内梯度名（按反向完成顺序）
    total_bytes: int

def bucket_gradients(
    grads: list[tuple[str, int]],   # [(参数名, 梯度字节数)]，按反向完成顺序（最后一层在前）
    bucket_bytes: int = 25 * 1024 * 1024,   # 默认桶大小 25 MB
) -> list[Bucket]:
    """
    贪心分桶：尽量把连续的梯度装进一个桶，直到装不下就开新桶。
    单个梯度超过桶上限时，单独成桶（不拆分）。
    """
    buckets: list[Bucket] = []
    cur_names: list[str] = []
    cur_bytes = 0

    for name, size in grads:
        if size > bucket_bytes:
            # 超大梯度：先把当前桶收尾，再单独成桶
            if cur_names:
                buckets.append(Bucket(cur_names, cur_bytes))
                cur_names, cur_bytes = [], 0
            buckets.append(Bucket([name], size))
        elif cur_bytes + size <= bucket_bytes:      # 填空：判断当前桶是否还装得下
            cur_names.append(name)
            cur_bytes += size
        else:
            # 装不下：封桶，开新桶
            buckets.append(Bucket(cur_names, cur_bytes))
            cur_names, cur_bytes = ([], 0)   # 填空：新桶初始状态
            cur_names.append(name)
            cur_bytes += size

    if cur_names:
        buckets.append(Bucket(cur_names, cur_bytes))
    return buckets

def simulate_dp_step(
    layer_backward_s: list[float],    # 每层 backward 时长（顺序：最后一层 → 第一层）
    bucket_comm_s: list[float],       # 每桶 all-reduce 时长（顺序与分桶一致）
    overlap: bool,
) -> float:
    """
    模拟一次训练步的时间。

    - 反向是串行的：层 i 的 backward 必须等层 i-1 完成。
    - 每桶的通信可以在"该桶最后一层 backward 完成"后立即开始；
      多个桶的通信在同一通信流上串行排队。
    - naive（overlap=False）：先全部 backward，再全部通信。
    - overlapped（overlap=True）：通信在释放后立即开始，尽量与后续 backward 重叠。
    """
    # 每层 backward 的累计完成时刻
    finish = []
    acc = 0.0
    for t in layer_backward_s:
        acc += t
        finish.append(acc)

    if not overlap:
        # naive：总时间 = 全部 backward + 全部通信
        return sum(layer_backward_s) + sum(bucket_comm_s)

    # overlapped：通信任务按释放时间排队
    # 释放时刻 = 该桶所属最后一层 backward 完成时刻。
    # 简化模型：把层分成 len(bucket_comm_s) 段，每段一层或均分；
    # 桶 b 的释放时刻 = 该段最后一层的 finish。
    release = []
    seg = max(1, len(layer_backward_s) // len(bucket_comm_s))
    for b in range(len(bucket_comm_s)):
        end_layer = min(len(layer_backward_s) - 1, (b + 1) * seg - 1)
        release.append(finish[end_layer])

    # 通信流上串行执行：t = max(当前时间, 释放时刻) + 通信时长
    t_comm = 0.0
    for b in range(len(bucket_comm_s)):
        t_comm = max(t_comm, release[b]) + bucket_comm_s[b]

    # 总时间 = max(backward 全部完成时刻, 通信流结束时刻)
    return max(finish[-1], t_comm)

def gbs_tokens(mbs: int, grad_acc: int, dp: int, seq: int) -> int:
    """全局 batch 大小（token 数）。"""
    return mbs * grad_acc * dp * seq

if __name__ == "__main__":
    # 例：30 层，每层 backward 10ms；分 6 桶，每桶通信 12ms
    layer_t = [0.010] * 30
    bucket_comm = [0.012] * 6

    t_naive = simulate_dp_step(layer_t, bucket_comm, overlap=False)
    t_over = simulate_dp_step(layer_t, bucket_comm, overlap=True)
    print(f"naive = {t_naive * 1e3:.0f} ms, overlapped = {t_over * 1e3:.0f} ms")

    # 分桶统计
    grads = [(f"p{i}", 2 * 1024 * 1024) for i in range(40)]  # 每梯度 2MB
    buckets = bucket_gradients(grads, bucket_bytes=10 * 1024 * 1024)
    print("桶数:", len(buckets), "（不分组会是 40 次通信）")
    print("每桶字节:", [b.total_bytes // (1024 * 1024) for b in buckets])

    # 全局 batch 示例
    print("gbs_tokens:", gbs_tokens(mbs=2, grad_acc=4, dp=128, seq=4096))

补全后的代码（关键填空处）

def bucket_gradients(grads, bucket_bytes=25 * 1024 * 1024):
    ...
    for name, size in grads:
        if size > bucket_bytes:
            if cur_names:
                buckets.append(Bucket(cur_names, cur_bytes))
                cur_names, cur_bytes = [], 0
            buckets.append(Bucket([name], size))
        elif cur_bytes + size <= bucket_bytes:     # 填空：还装得下
            cur_names.append(name)
            cur_bytes += size
        else:
            buckets.append(Bucket(cur_names, cur_bytes))
            cur_names, cur_bytes = [], 0           # 填空：新桶初始状态
            cur_names.append(name)
            cur_bytes += size
    ...

时间模拟：

if not overlap:
        return sum(layer_backward_s) + sum(bucket_comm_s)   # naive

    # overlapped：通信释放后立即排队（同一通信流）
    t_comm = 0.0
    for b in range(len(bucket_comm_s)):
        t_comm = max(t_comm, release[b]) + bucket_comm_s[b]
    return max(finish[-1], t_comm)

gbs：return mbs * grad_acc * dp * seq。

运行示例（30 层 × 10ms、6 桶 × 12ms）：naive = 372 ms，overlapped = 312 ms（部分通信被隐藏；全部隐藏时 = 300 ms）。40 个 2MB 梯度、10MB 桶 → 8 个桶（8 次通信而非 40 次）。

三个问答题要点

Q1

正确性：优化器 step 必须用”所有副本梯度的平均”，所以 all-reduce 是 step 前的同步点。
可重叠性：梯度是按层产生的（反向从后往前），每层的同步只依赖该层梯度就绪，不依赖前面层的 backward——可以用 hook 逐参数触发，通信与剩余 backward 并行；
与 TP 的本质区别：TP 的 all-reduce 之后紧接着就有必须等待该结果的计算（如 LayerNorm/下一层 matmul），同步点嵌在计算链内部、每层多次；DP 的同步点在一个训练步中只有一批、且之后没有立即依赖它的计算（等的是整个 backward 完成后的 step）。

Q2

正确性：all-reduce（SUM）是结合律运算——Σ_mb Σ_rank 与先局部累加再一次性 reduce 结果相同；平均 = SUM/(dp·grad_acc) 在最后一次统一做即可；
定量：每次都同步 → 每步 grad_acc 次 all-reduce（每次 2Ψ 通信量）；只最后一次同步 → 1 次。grad_acc=8 时通信次数/通信量差 8 倍（每次同步的通信量相同）。

Q3

效率下降：ring latency 随 N 累积（每步 2(N−1) 次消息的固定开销），网络需求（总带宽 = N·2Ψ/步）随 N 增长，超过一定规模后通信无法再被 backward 完全隐藏；
显存不变：每个 DP 副本都持有完整参数/梯度/优化器状态，dp 增加只分摊 batch，不分摊模型；
70B 放不下：静态账本 70B×16B ≈ 1120 GB ≫ 80 GB，mbs=1 也 OOM → 需要 ZeRO（第 5 课）、TP（第 6 课）、PP（第 9 课）等。

通过要点

教材配方顺序：gbs → seq → 单卡最大 mbs → dp → grad_acc = gbs/(mbs·dp·seq)。
“DP 是并行版梯度累积”：gbs = mbs × grad_acc × dp；优先 dp 后 grad_acc。

官方主参考

← 上一课
第 3 课：分布式通信原语 下一课 →
第 5 课：ZeRO 与 FSDP

--- title: "第 4 课：数据并行 DP 与 DDP" jupyter: python3 --- - 对应官方章节：Data Parallelism（含 Revisiting global batch size、Our journey up to now） - 前置：第 2 课（梯度累积）、第 3 课（all-reduce） - 状态：未开始 ## 本课目标完成后你需要能够： - 把 DP 表述为"并行版梯度累积"，写出 `gbs = mbs × grad_acc × dp`。 - 解释为什么梯度 all-reduce 必须在 optimizer step 之前完成，但可以与 backward 重叠。 - 实现并解释两个优化：逐参数 hook 重叠、梯度 bucketing。 - 解释 `no_sync` 在"DP + 梯度累积"中的正确性作用。 - 说明 DP 的边界：不降低每卡静态显存，规模过大（512+）时效率下降。 ## 核心概念 ### 1. 数据并行（DP/DDP）的数学本质与并行语义数据并行（Data Parallelism）在数学本质上是**空间维度展开的梯度累积**。如**图 4-1 与图 4-2** 所示，整个集群拥有 $N$ 个并行工作进程（Rank 0 至 Rank $N-1$），每个 GPU 显存内都完整存放了一份同构的模型参数副本 $W$： :::: {.vfig} ::: {.vfig-title} Ring AllReduce = Reduce-Scatter + All-Gather ::: ::::: {.vfig-grid .vfig-flow style="--cols:2"} ::: {.vfig-card .c-blue} **阶段 1 · Reduce-Scatter** - 梯度切成 $P$ 块，沿环走 $P-1$ 步 - 每步发出 1 块、收到 1 块并累加 - 结束时每卡持有 1 块的全局和 ::: ::: {.vfig-card .c-green} **阶段 2 · All-Gather** - 再走 $P-1$ 步，转发已求和的块 - 每步用收到的块覆盖本地旧值 - 结束时每卡都有完整的全局和 ::: ::::: ::: {.vfig-key} 单卡发送量 $2\frac{P-1}{P}M < 2M$，几乎不随卡数 $P$ 增长；耗时 $T \approx 2(P-1)\alpha + 2\frac{P-1}{P}\cdot\frac{M}{\beta}$（$\alpha$ 为单步延迟，$\beta$ 为带宽） ::: :::: <p class="caption">图 4-1：Ring AllReduce 的两个阶段与单卡通信量</p> ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-sg" data-nid="DataParallel" data-depth="0"><div class="fd-sg-title">数据并行架构 (DDP)</div><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-rank"><div class="fd-node fd-box" data-nid="Data">全局 Batch 训练数据</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="Split1">分片 0 (Rank 0)</div><div class="fd-node fd-box" data-nid="Split2">分片 1 (Rank 1)</div><div class="fd-node fd-box" data-nid="Split3">分片 2 (Rank 2)</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="Model0">卡 0: 模型完整副本 W</div><div class="fd-node fd-box" data-nid="Model1">卡 1: 模型完整副本 W</div><div class="fd-node fd-box" data-nid="Model2">卡 2: 模型完整副本 W</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="Grad0">本地梯度 g_0</div><div class="fd-node fd-box" data-nid="Grad1">本地梯度 g_1</div><div class="fd-node fd-box" data-nid="Grad2">本地梯度 g_2</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="AR">Ring AllReduce 全局同步</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Update">各卡同步更新同构参数 W_new</div></div></div></div></div></div></div><i class="fd-e" data-s="Data" data-t="Split1" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Data" data-t="Split2" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Data" data-t="Split3" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Split1" data-t="Model0" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Split2" data-t="Model1" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Split3" data-t="Model2" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Model0" data-t="Grad0" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Model1" data-t="Grad1" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Model2" data-t="Grad2" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Grad0" data-t="AR" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Grad1" data-t="AR" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Grad2" data-t="AR" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="AR" data-t="Update" data-k="solid" data-h="1" data-b="0"></i></div> ``` <p class="caption" align="center"><em>图 4-2：数据并行（DDP）同构副本前向与全量梯度归约架构</em></p> 当一个全局 Batch（Global Batch Size, GBS）的数据输入系统时，各 Rank 分别取走互不重叠的数据子集（Micro Batch, MBS）。由于各 Rank 上的前向传播与反向传播计算完全基于各自的本地数据进行，彼此之间没有任何数据依赖，因此计算阶段具备天然的完全并行性。然而，优化器更新全局参数时，必须基于全局样本的真实平均梯度。设 $\text{GBS} = \text{MBS} \times \text{grad\_accum} \times \text{DP}$，全局目标损失对参数 $\theta$ 的梯度定义为： $$\nabla_\theta \mathcal{L}_{\text{global}} = \frac{1}{\text{GBS}} \sum_{i=1}^{\text{GBS}} \nabla_\theta \ell(x_i; \theta) = \frac{1}{\text{DP}} \sum_{r=0}^{\text{DP}-1} \left( \frac{1}{\text{MBS} \times \text{grad\_accum}} \sum_{m=1}^{\text{grad\_accum}} \sum_{j=1}^{\text{MBS}} \nabla_\theta \ell(x_{r, m, j}; \theta) \right)$$ 由此可见，每个 Rank 独立求导得到的仅是其本地分片的梯度之和。在执行优化器更新（`optimizer.step()`）之前，必须通过第 3 课所讲的 **AllReduce（求和后归一化除以 DP 规模）** 算子，在所有参与数据并行的节点间同步梯度。这保证了所有副本在步结时取得完全相同的新权重 $W_{\text{new}}$。训练规模的核心计量维度公式为： $$\text{GBS} = \text{MBS} \times \text{grad\_accum} \times \text{DP}$$ $$\text{Batch Tokens} = \text{GBS} \times \text{seq\_len} = \text{MBS} \times \text{grad\_accum} \times \text{DP} \times \text{seq\_len}$$ 在维持全局 Batch 大小不变的前提下，增大 $\text{DP}$ 规模可以直接等比削减串行累积步数 $\text{grad\_accum}$，从而大幅缩短每步训练的墙钟时间。 ### 2. 朴素实现的瓶颈：计算与通信完全串行在最原始的 DistributedDataParallel 朴素实现中，执行流程遵循严格的阶段式阻塞顺序： 1. **前向与反向完全执行完毕**：所有网络层反向计算结束，显存中产生所有参数的梯度张量。 2. **一次性阻塞式 AllReduce**：启动 NCCL 通信，将全模型参数量 $\Psi$ 对应的全部梯度张量进行同步传输。 3. **优化器更新**：等待通信全部完成后，统一更新权重。这种串行逻辑存在巨大的算力浪费。其一步总耗时为： $$T_{\text{step}} = T_{\text{fwd}} + T_{\text{bwd}} + T_{\text{comm}} + T_{\text{opt}}$$ 在 Ring AllReduce 模式下，传输全模型梯度的网络通信时间为 $T_{\text{comm}} \approx 2 \times \frac{\Psi \times \text{sizeof(dtype)}}{B_{\text{bus}}}$。在万卡集群或跨节点低速网络（如普通以太网）下，$T_{\text{comm}}$ 往往占据整步耗时的 30%~60%。在通信期间，数十万 TFLOPS 的 Tensor Core 处于完全空闲状态，吞吐极度低下。 ### 3. 优化一：梯度后向钩子与通信隐藏（Post-Accumulate Grad Hook）要消除空闲等待，核心思路是实现**计算与通信的异步重叠（Overlap）**。注意到自动微分（Autograd）的反向传播是反向有向无环图（DAG）的倒序执行过程：从最后一层（输出层 Layer $L$）逐层反向推进到第一层（输入层 Layer $1$）。当输出层 Layer $L$ 的参数求导计算完毕时，其梯度 $\nabla W_L$ 的数值已经完全确定，绝不会再受到后续浅层网络求导计算的影响。现代 PyTorch DDP 内部通过为每个可训练参数注册底层回调（如 `register_post_accumulate_grad_hook`）实现通信重叠： - 当反向计算图推进并生成参数 $W_k$ 的梯度时，Hook 立即被触发。 - DDP 自动将该梯度的 AllReduce 异步投递到独立的 CUDA 通信流（`nccl_stream`）中执行。 - 此时，默认的计算流（`compute_stream`）无须任何等待，继续向前推进计算 Layer $k-1$、Layer $k-2$ 的反向求导。 - 只要浅层的反向计算耗时大于等于深层梯度的网络传输耗时，绝大部分通信时间就被反向计算完全掩盖（$T_{\text{comm}}$ 隐藏在 $T_{\text{bwd}}$ 背后）。如**图 4-3** 所示，整个反向计算与网络传输在时间线上呈现高度交叠的双轨并发执行： ```{=html} <div class="fd fd-seq" role="img"><div class="sd-grid" style="--n:3"><div class="sd-life" style="grid-column:1;grid-row:1/11"></div><div class="sd-life" style="grid-column:2;grid-row:1/11"></div><div class="sd-life" style="grid-column:3;grid-row:1/11"></div><div class="sd-actor" style="grid-column:1;grid-row:1">反向求导 (Backward Compute)</div><div class="sd-actor" style="grid-column:2;grid-row:1">梯度分桶 (Bucket 0 / Bucket 1)</div><div class="sd-actor" style="grid-column:3;grid-row:1">异步通信流 (Ring AllReduce)</div><div class="sd-note sd-note-over" style="grid-column:1/4;grid-row:2;--k:3">PyTorch DDP 计算与通信重叠流水线</div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:1/3;grid-row:3;--k:2"><div class="sd-text"><span class="sd-num">1</span>逐层反向求导 (Layer N)</div><div class="sd-line"></div></div><div class="sd-msg sd-self sd-dash" style="grid-column:2/span 2;grid-row:4;--k:2"><div class="sd-text"><span class="sd-num">2</span>梯度写入，达到阈值 (如 25MB)</div><div class="sd-loop"></div></div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:2/4;grid-row:5;--k:2"><div class="sd-text"><span class="sd-num">3</span>异步触发 AllReduce (Bucket 0)</div><div class="sd-line"></div></div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:1/3;grid-row:6;--k:2"><div class="sd-text"><span class="sd-num">4</span>继续求导上一层 (Layer N-1, N-2)</div><div class="sd-line"></div></div><div class="sd-note sd-note-over" style="grid-column:3/4;grid-row:7;--k:1">环形通信与计算并行执行</div><div class="sd-msg sd-self sd-dash" style="grid-column:3/span 1;grid-row:8;--k:1"><div class="sd-text"><span class="sd-num">5</span>Bucket 0 同步完成</div><div class="sd-loop"></div></div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:2/4;grid-row:9;--k:2"><div class="sd-text"><span class="sd-num">6</span>异步触发 AllReduce (Bucket 1)</div><div class="sd-line"></div></div><div class="sd-msg sd-self sd-dash" style="grid-column:3/span 1;grid-row:10;--k:1"><div class="sd-text"><span class="sd-num">7</span>全部 Bucket 同步完毕</div><div class="sd-loop"></div></div></div></div> ``` <p class="caption" align="center"><em>图 4-3：PyTorch DDP 梯度分桶与异步 AllReduce 计算通信重叠流水</em></p> ### 4. 优化二：梯度分桶策略（Gradient Bucketing）虽然逐参数挂载 Hook 理论上能实现最大程度的重叠，但在实际大模型中直接对每个小参数张量单独发起 AllReduce 会引发严重的性能倒退： 1. **小包延迟开销爆炸**：神经网络中包含大量极小的权重和偏置张量（例如 LayerNorm 的 scale/bias 仅数 KB）。根据第 3 课的网络延迟模型 $T = 2(N-1)\alpha + 2\frac{N-1}{N}\frac{K}{B}$，当张量尺寸 $K$ 很小时，耗时完全由每步通信固有的延迟基数 $\alpha$ 主导。成百上千次小包 AllReduce 会导致网络传输极其低效。 2. **CUDA 算子发射开销（Kernel Launch Overhead）**：CPU 端频繁向 GPU 发射小张量通信算子会迅速跑满 CPU，导致 GPU 发生发射等待饥饿。 PyTorch DDP 采用**分桶机制（Bucketing）**进行折中处理： - 按照网络模型反向求导完成的**逆序方向**，预先分配若干块连续的扁平内存切片（称为 Bucket，默认上限通常为 `bucket_cap_mb = 25MB`）。 - 各层的梯度产生后，先被直接拷贝（Copy）到对应的连续 Bucket 内存块中。 - 当一个 Bucket 被填满，或者倒序遍历到了该桶管辖的最后一层梯度时，系统才对这块规整的 25MB 大张量集中发射一次异步 Ring AllReduce。 - 这种策略既规避了小包通信的延迟惩罚，又避免了一次性同步全量模型带来的漫长等待，实现了网络总线利用率与重叠潜力的最优权衡。 ### 5. 优化三：与梯度累积深度协同（`no_sync` 机制）在 DP + 梯度累积时，DDP 默认会在每个 Micro-Batch 的反向阶段都触发分桶 AllReduce。以 $\text{grad\_accum} = 8$ 为例，前 7 次同步的结果只会被下一次本地累加覆盖，每步多付出 $7 \times 2\Psi$ 的通信量。由于梯度求和满足**加法结合律**： $$\sum_{r=0}^{\text{DP}-1} \left( \sum_{m=1}^{\text{grad\_accum}} g_{r, m} \right) = \sum_{m=1}^{\text{grad\_accum}} \left( \sum_{r=0}^{\text{DP}-1} g_{r, m} \right)$$ 因此可以先在每张卡本地累加各 Micro-Batch 的梯度，只在最后一次反向传播时做一次 AllReduce。PyTorch DDP 用上下文管理器 `model.no_sync()` 实现这一点：在 `with model.no_sync():` 块内反向只做本地累加、不触发 AllReduce；最后一个 Micro-Batch 在块外执行，触发一次同步。通信次数因此降为原来的 $1 / \text{grad\_accum}$，这是多卡 DP + 梯度累积的标准做法。 ### 6. 工程细节与 DDP 的物理扩展边界理解 DDP 的边界是掌握后续进阶并行架构（ZeRO、TP、PP）的前提： 1. **通信内存开销**：`bucket_cap_mb = 25MB` 是单个桶的上限，而全部 Bucket 合起来要容纳所有参数的梯度，即额外一份完整梯度副本（$\Psi \times \text{sizeof(dtype)}$，7B BF16 约 14 GB），并非只多 25MB～50MB。设置 `gradient_as_bucket_view=True` 后，`param.grad` 直接作为 Bucket 的视图，可以省掉这份复制。 2. **DDP 完全无法降低每卡的静态显存**：这是纯数据并行最致命的物理局限。每个 DDP 副本的 GPU 显存内都必须完整驻留 $1\times$ 模型参数 $\Psi$、$1\times$ 梯度 $\Psi$ 以及 $1\times$ 优化器状态（FP32 Master Weights + Momentum + Variance 共 $12\Psi$ 字节）。 - 对于一个 70B 参数的 BF16 大模型，其静态状态显存高达 $70 \times 16\text{ GB} = 1120\text{ GB}$。 - 即使单卡显存配置到顶级的 80GB/192GB，并将 Micro-Batch 降为 1 且开启全激活重计算，单卡也根本无法容纳模型静态权重，程序将瞬间触发 OOM 崩溃。 3. **网络扩展性极限（Ring Latency 惩罚）**： - 当节点规模扩展到数百乃至数千卡（如 $\text{DP} \ge 512$）时，Ring AllReduce 的延迟项 $2(N-1)\alpha$ 随着环形步数呈线性膨胀。跨交换机（Spine-Leaf 网络二跳/三跳）的延迟显著拉长。 - 当通信延迟超过了反向计算耗时时，异步重叠完全失效，算力利用率暴跌。 - 此时必须联合张量并行（TP，将计算切分到机内 NVLink 超高带宽域）以及流水线并行（PP，将层深切分），构成多维混合并行系统。 ## 具体演示与数值推演假设训练目标设定为全局吞吐 $\text{BST} = 4\text{M tokens}$（即 $4,194,304$ tokens），上下文长度 $\text{seq} = 4096$，则每步所需全局样本数为： $$\text{GBS} = \frac{4,194,304}{4096} = 1024 \text{ samples}$$ 单卡在兼顾激活显存与算力饱和度的情况下，测定最大单卡批次 $\text{MBS} = 2$： - **128 卡集群（$\text{DP} = 128$）**： $$\text{grad\_accum} = \frac{\text{GBS}}{\text{MBS} \times \text{DP}} = \frac{1024}{2 \times 128} = 4$$ 每卡连续计算 4 个 Micro-Batch（前 3 个使用 `no_sync` 纯本地累加，第 4 个触发异步 AllReduce），然后执行一次全局权重更新。 - **512 卡集群（$\text{DP} = 512$）**： $$\text{grad\_accum} = \frac{1024}{2 \times 512} = 1$$ 此时不需要任何梯度累加，每个 Micro-Batch 直接就是全局步，训练速度提升近 4 倍。 - **扩展性瓶颈突显**：在 512 卡场景下，反向计算时间缩短为单个 MBS=2 的计算耗时（假设为 80ms），但 512 卡跨节点 Ring AllReduce 的延迟开销与通信排队时间大幅增加（例如通信耗时上升至 120ms）。由于 $T_{\text{comm}} > T_{\text{bwd}}$，异步重叠出现暴露窗口，GPU 产生 40ms 的通信气泡（Bubble）。继续无脑堆叠 DP 规模将陷入严重的边际收益递减。 ## 代码填空题补全以下程序：梯度分桶 + 朴素/重叠两种方案的完成时间模拟。 ```{python} #| course_role: exercise from dataclasses import dataclass @dataclass class Bucket: names: list[str] # 桶内梯度名（按反向完成顺序） total_bytes: int def bucket_gradients( grads: list[tuple[str, int]], # [(参数名, 梯度字节数)]，按反向完成顺序（最后一层在前） bucket_bytes: int = 25 * 1024 * 1024, # 默认桶大小 25 MB ) -> list[Bucket]: """ 贪心分桶：尽量把连续的梯度装进一个桶，直到装不下就开新桶。单个梯度超过桶上限时，单独成桶（不拆分）。 """ buckets: list[Bucket] = [] cur_names: list[str] = [] cur_bytes = 0 for name, size in grads: if size > bucket_bytes: # 超大梯度：先把当前桶收尾，再单独成桶 if cur_names: buckets.append(Bucket(cur_names, cur_bytes)) cur_names, cur_bytes = [], 0 buckets.append(Bucket([name], size)) elif ______: # 填空：判断当前桶是否还装得下 cur_names.append(name) cur_bytes += size else: # 装不下：封桶，开新桶 buckets.append(Bucket(cur_names, cur_bytes)) cur_names, cur_bytes = ______ # 填空：新桶初始状态 cur_names.append(name) cur_bytes += size if cur_names: buckets.append(Bucket(cur_names, cur_bytes)) return buckets def simulate_dp_step( layer_backward_s: list[float], # 每层 backward 时长（顺序：最后一层 → 第一层） bucket_comm_s: list[float], # 每桶 all-reduce 时长（顺序与分桶一致） overlap: bool, ) -> float: """ 模拟一次训练步的时间。 - 反向是串行的：层 i 的 backward 必须等层 i-1 完成。 - 每桶的通信可以在"该桶最后一层 backward 完成"后立即开始；多个桶的通信在同一通信流上串行排队。 - naive（overlap=False）：先全部 backward，再全部通信。 - overlapped（overlap=True）：通信在释放后立即开始，尽量与后续 backward 重叠。 """ # 每层 backward 的累计完成时刻 finish = [] acc = 0.0 for t in layer_backward_s: acc += t finish.append(acc) if not overlap: # naive：总时间 = 全部 backward + 全部通信 return ______ # overlapped：通信任务按释放时间排队 # 释放时刻 = 该桶所属最后一层 backward 完成时刻。 # 简化模型：把层分成 len(bucket_comm_s) 段，每段一层或均分； # 桶 b 的释放时刻 = 该段最后一层的 finish。 release = [] seg = max(1, len(layer_backward_s) // len(bucket_comm_s)) for b in range(len(bucket_comm_s)): end_layer = min(len(layer_backward_s) - 1, (b + 1) * seg - 1) release.append(finish[end_layer]) # 通信流上串行执行：t = max(当前时间, 释放时刻) + 通信时长 t_comm = 0.0 for b in range(len(bucket_comm_s)): t_comm = ______(t_comm, release[b]) + bucket_comm_s[b] # 总时间 = max(backward 全部完成时刻, 通信流结束时刻) return max(finish[-1], t_comm) def gbs_tokens(mbs: int, grad_acc: int, dp: int, seq: int) -> int: """全局 batch 大小（token 数）。""" return ______ if __name__ == "__main__": # 例：30 层，每层 backward 10ms；分 6 桶，每桶通信 12ms layer_t = [0.010] * 30 bucket_comm = [0.012] * 6 t_naive = simulate_dp_step(layer_t, bucket_comm, overlap=False) t_over = simulate_dp_step(layer_t, bucket_comm, overlap=True) print(f"naive = {t_naive * 1e3:.0f} ms, overlapped = {t_over * 1e3:.0f} ms") # 分桶统计 grads = [(f"p{i}", 2 * 1024 * 1024) for i in range(40)] # 每梯度 2MB buckets = bucket_gradients(grads, bucket_bytes=10 * 1024 * 1024) print("桶数:", len(buckets), "（不分组会是 40 次通信）") print("每桶字节:", [b.total_bytes // (1024 * 1024) for b in buckets]) # 全局 batch 示例 print("gbs_tokens:", gbs_tokens(mbs=2, grad_acc=4, dp=128, seq=4096)) ``` ## 三个问答题 ### Q1 为什么梯度同步必须在 optimizer step 之前完成？为什么它可以（部分）与 backward 重叠，而第 6 课将看到 TP 的通信"在关键路径上、难以重叠"？两者有什么本质区别？ ### Q2 DP + 梯度累积时，为什么可以用 `no_sync()` 只在最后一次 backward 后 all-reduce？请从运算性质（结合律）说明正确性，并定量比较：grad_acc=8、dp=64 时，"每次都同步"与"只同步一次"的通信次数差多少倍？ ### Q3 教材说 DP 在 512+ 卡时效率明显下降，且 DP 不降低每卡显存。请解释这两个限制分别来自哪里，并说明：一个 70B 模型在 8×80GB 节点上训练，为什么纯 DP（含 grad_acc）无法解决"模型装不下"的问题？此时应该转向哪些技术（给出名字即可）？ ## 检查与通过标准总分 10 分：代码正确 4 分（分桶条件、时间模拟、gbs 公式）、三题各 2 分、通过线 8 分。一票否决项： - 认为 DP 可以降低每卡参数/优化器显存。 - 认为"每次 backward 都同步"与"最后一次同步"在结果上不等价。 - 不知道 all-reduce 是求和（平均需要再除以 dp）或同步时机错误。 - 声称"DP 规模可以无限扩展"或"DP 能解决模型放不下单卡的问题"。 ::: {.callout-tip collapse="true" title="💡 点击展开参考答案与详细代码实现" icon="false"} ## 参考答案（仅 answer 分支）先完成题目再核对；核对后必须解释关键公式，并改一个规模重新计算。 ```{python} #| course_role: solution from dataclasses import dataclass @dataclass class Bucket: names: list[str] # 桶内梯度名（按反向完成顺序） total_bytes: int def bucket_gradients( grads: list[tuple[str, int]], # [(参数名, 梯度字节数)]，按反向完成顺序（最后一层在前） bucket_bytes: int = 25 * 1024 * 1024, # 默认桶大小 25 MB ) -> list[Bucket]: """ 贪心分桶：尽量把连续的梯度装进一个桶，直到装不下就开新桶。单个梯度超过桶上限时，单独成桶（不拆分）。 """ buckets: list[Bucket] = [] cur_names: list[str] = [] cur_bytes = 0 for name, size in grads: if size > bucket_bytes: # 超大梯度：先把当前桶收尾，再单独成桶 if cur_names: buckets.append(Bucket(cur_names, cur_bytes)) cur_names, cur_bytes = [], 0 buckets.append(Bucket([name], size)) elif cur_bytes + size <= bucket_bytes: # 填空：判断当前桶是否还装得下 cur_names.append(name) cur_bytes += size else: # 装不下：封桶，开新桶 buckets.append(Bucket(cur_names, cur_bytes)) cur_names, cur_bytes = ([], 0) # 填空：新桶初始状态 cur_names.append(name) cur_bytes += size if cur_names: buckets.append(Bucket(cur_names, cur_bytes)) return buckets def simulate_dp_step( layer_backward_s: list[float], # 每层 backward 时长（顺序：最后一层 → 第一层） bucket_comm_s: list[float], # 每桶 all-reduce 时长（顺序与分桶一致） overlap: bool, ) -> float: """ 模拟一次训练步的时间。 - 反向是串行的：层 i 的 backward 必须等层 i-1 完成。 - 每桶的通信可以在"该桶最后一层 backward 完成"后立即开始；多个桶的通信在同一通信流上串行排队。 - naive（overlap=False）：先全部 backward，再全部通信。 - overlapped（overlap=True）：通信在释放后立即开始，尽量与后续 backward 重叠。 """ # 每层 backward 的累计完成时刻 finish = [] acc = 0.0 for t in layer_backward_s: acc += t finish.append(acc) if not overlap: # naive：总时间 = 全部 backward + 全部通信 return sum(layer_backward_s) + sum(bucket_comm_s) # overlapped：通信任务按释放时间排队 # 释放时刻 = 该桶所属最后一层 backward 完成时刻。 # 简化模型：把层分成 len(bucket_comm_s) 段，每段一层或均分； # 桶 b 的释放时刻 = 该段最后一层的 finish。 release = [] seg = max(1, len(layer_backward_s) // len(bucket_comm_s)) for b in range(len(bucket_comm_s)): end_layer = min(len(layer_backward_s) - 1, (b + 1) * seg - 1) release.append(finish[end_layer]) # 通信流上串行执行：t = max(当前时间, 释放时刻) + 通信时长 t_comm = 0.0 for b in range(len(bucket_comm_s)): t_comm = max(t_comm, release[b]) + bucket_comm_s[b] # 总时间 = max(backward 全部完成时刻, 通信流结束时刻) return max(finish[-1], t_comm) def gbs_tokens(mbs: int, grad_acc: int, dp: int, seq: int) -> int: """全局 batch 大小（token 数）。""" return mbs * grad_acc * dp * seq if __name__ == "__main__": # 例：30 层，每层 backward 10ms；分 6 桶，每桶通信 12ms layer_t = [0.010] * 30 bucket_comm = [0.012] * 6 t_naive = simulate_dp_step(layer_t, bucket_comm, overlap=False) t_over = simulate_dp_step(layer_t, bucket_comm, overlap=True) print(f"naive = {t_naive * 1e3:.0f} ms, overlapped = {t_over * 1e3:.0f} ms") # 分桶统计 grads = [(f"p{i}", 2 * 1024 * 1024) for i in range(40)] # 每梯度 2MB buckets = bucket_gradients(grads, bucket_bytes=10 * 1024 * 1024) print("桶数:", len(buckets), "（不分组会是 40 次通信）") print("每桶字节:", [b.total_bytes // (1024 * 1024) for b in buckets]) # 全局 batch 示例 print("gbs_tokens:", gbs_tokens(mbs=2, grad_acc=4, dp=128, seq=4096)) ``` ## 补全后的代码（关键填空处） ```python def bucket_gradients(grads, bucket_bytes=25 * 1024 * 1024): ... for name, size in grads: if size > bucket_bytes: if cur_names: buckets.append(Bucket(cur_names, cur_bytes)) cur_names, cur_bytes = [], 0 buckets.append(Bucket([name], size)) elif cur_bytes + size <= bucket_bytes: # 填空：还装得下 cur_names.append(name) cur_bytes += size else: buckets.append(Bucket(cur_names, cur_bytes)) cur_names, cur_bytes = [], 0 # 填空：新桶初始状态 cur_names.append(name) cur_bytes += size ... ``` 时间模拟： ```python if not overlap: return sum(layer_backward_s) + sum(bucket_comm_s) # naive # overlapped：通信释放后立即排队（同一通信流） t_comm = 0.0 for b in range(len(bucket_comm_s)): t_comm = max(t_comm, release[b]) + bucket_comm_s[b] return max(finish[-1], t_comm) ``` gbs：`return mbs * grad_acc * dp * seq`。运行示例（30 层 × 10ms、6 桶 × 12ms）：naive = 372 ms，overlapped = 312 ms（部分通信被隐藏；全部隐藏时 = 300 ms）。40 个 2MB 梯度、10MB 桶 → 8 个桶（8 次通信而非 40 次）。 ## 三个问答题要点 ### Q1 - 正确性：优化器 step 必须用"所有副本梯度的平均"，所以 all-reduce 是 step 前的同步点。 - 可重叠性：梯度是按层产生的（反向从后往前），每层的同步只依赖该层梯度就绪，不依赖前面层的 backward——可以用 hook 逐参数触发，通信与剩余 backward 并行； - 与 TP 的本质区别：TP 的 all-reduce 之后**紧接着就有必须等待该结果的计算**（如 LayerNorm/下一层 matmul），同步点嵌在计算链内部、每层多次；DP 的同步点在一个训练步中只有一批、且之后没有立即依赖它的计算（等的是整个 backward 完成后的 step）。 ### Q2 - 正确性：all-reduce（SUM）是结合律运算——`Σ_mb Σ_rank` 与先局部累加再一次性 reduce 结果相同；平均 = SUM/(dp·grad_acc) 在最后一次统一做即可； - 定量：每次都同步 → 每步 grad_acc 次 all-reduce（每次 2Ψ 通信量）；只最后一次同步 → 1 次。grad_acc=8 时通信次数/通信量差 8 倍（每次同步的通信量相同）。 ### Q3 - 效率下降：ring latency 随 N 累积（每步 2(N−1) 次消息的固定开销），网络需求（总带宽 = N·2Ψ/步）随 N 增长，超过一定规模后通信无法再被 backward 完全隐藏； - 显存不变：每个 DP 副本都持有完整参数/梯度/优化器状态，dp 增加只分摊 batch，不分摊模型； - 70B 放不下：静态账本 70B×16B ≈ 1120 GB ≫ 80 GB，mbs=1 也 OOM → 需要 ZeRO（第 5 课）、TP（第 6 课）、PP（第 9 课）等。 ## 通过要点 - 教材配方顺序：gbs → seq → 单卡最大 mbs → dp → grad_acc = gbs/(mbs·dp·seq)。 - "DP 是并行版梯度累积"：gbs = mbs × grad_acc × dp；优先 dp 后 grad_acc。 ## 官方主参考 - [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook) - [PyTorch distributed documentation](https://pytorch.org/docs/stable/distributed.html) ::: ::: {.course-pagination} [**← 上一课**<br>第 3 课：分布式通信原语](lesson03.qmd){.pagination-link .pagination-prev} [**下一课 →**<br>第 5 课：ZeRO 与 FSDP](lesson05.qmd){.pagination-link .pagination-next} :::