CS336 · 第 7 课

分布式训练并行

单张 GPU 会遇到两个上限：模型状态可能放不下，训练速度也可能不够。增加 GPU 之后，问题并不会自动消失——计算被拆开了，但不同设备必须交换参数、梯度或激活。本讲要建立的主线是：先看必须交换什么，再按通信频率和物理拓扑选择并行策略。

这一讲到底解决什么问题？

上一讲解决的是单张 GPU 内部的数据搬运：怎样通过 fusion、tiling 和更好的 kernel 减少 HBM traffic。模型扩展到多张 GPU 后，同一个问题跨越了设备边界：本地 memory traffic 变成了 network communication。

单卡是否能训练一个模型，不能只看参数量。梯度、优化器状态、激活、临时 buffer 和通信 workspace 都会占用显存；即使模型勉强放得下，单卡完成全部 FLOPs 也可能太慢。因此我们增加 GPU，但随之得到一个新问题：

多张 GPU 分别完成一部分计算后，需要交换哪些数据，交换多少次，又应该经过哪一类互联？

DP: batch shards | full model replicas | synchronize gradients
FSDP: batch shards | model-state shards | gather/scatter states
TP: same tokens   | tensor dimension shards | communicate layer partials
PP: microbatches  | layer-range shards | send activations/gradients

logical tensor bytes → collective algorithm → physical links → wall time

先记住一个统一视角

并行策略不是一组互不相关的名词。它们都在同一组资源之间做交换：可以复制计算、保留更多显存，或把数据放到另一张 GPU 后再通信。后文每介绍一种策略，都只问三件事：分片了什么、通信了什么、通信发生多频繁。

通信原语

假设每张 GPU 都根据自己的 mini-batch 算出一份梯度。下一步不是“让 GPU 聊天”，而是一个具体的数据一致性问题：怎样让所有 GPU 获得聚合后的同一份梯度？分布式程序用 rank 标识一个参与者，用 world size 表示参与者总数，再用 collective 描述整组 rank 之间的数据变换。

先掌握三种核心 collective

这三种操作并不是孤立 API。对带宽高效的 ring 实现，可以先让每个 rank 得到最终结果的一片，再交换这些分片：

\[\text{All-Reduce}=\text{Reduce-Scatter}+\text{All-Gather}\]

从直觉到 cost model

如果每个 rank 都把完整梯度直接发送给所有其他 rank，链路会重复搬运大量数据。Ring 的直觉是让 rank 只和邻居交换分片：第一阶段边传边聚合，第二阶段再传播已经聚合好的分片。这样每个 rank 搬运的总字节数与张量大小同阶，而不是随 GPU 数线性增加。

网络 substrate：这些 bytes 经过什么硬件？

到这里我们已经知道 collective 需要传输多少数据，但通信量只是逻辑成本。真正的 wall-clock 时间取决于这些 bytes 走的是节点内 scale-up fabric，还是经过 NIC 进入节点间 scale-out 网络。

Scale-up、scale-out 与并行策略放置

• TP / EP 偏好 scale-up：它们在每层或每个 MoE block 触发高频 collective，对延迟和带宽都敏感，通常先放入 NVLink/NVSwitch 域。
• PP 可以跨较慢网络：它主要在 stage 边界传 activation/gradient；通过足够大的 microbatch 和 1F1B schedule，可以用计算覆盖较粗粒度通信。
• DP/FSDP 常扩展到 scale-out：通信频率较低但 payload 大，需要 bucket、prefetch、reduce-scatter/all-gather 与计算重叠。

NCCL 为什么必须感知 topology

NCCL 的作用是把逻辑 collective 映射到实际硬件。它会探测节点、PCIe switch、NVLink/NVSwitch 与网络接口，再选择 ring/tree、分层 collective、channel 数和 GPU↔NIC 路径。

一阶时间模型可以写成：

\[T\approx L\alpha+\frac{B}{\beta_{\mathrm{effective}}}\]

其中 \(L\) 是需要串行完成的通信步骤，\(\alpha\) 是每步启动延迟，\(B\) 是字节数，\(\beta_{\mathrm{effective}}\) 是实际有效带宽。最慢链路、拓扑冲突、拥塞、协议效率和计算重叠都会改变它，因此相同的逻辑通信量不保证相同的运行时间。

Data Parallelism (DP)

现在把 collective 放回训练问题。最直接的扩展方式是让每张 GPU 保留同一个模型、处理不同样本；本地反向结束后，各 rank 的梯度必须重新一致，因此刚才的 All-Reduce 在这里第一次真正派上用场。

核心思想

每个 GPU 持有完整模型，处理不同的数据片。梯度通过 all-reduce 同步。

# Data Parallelism 流程
for each GPU in parallel:
    # 1. 获取不同的数据片
    x_local = data[rank * batch_size : (rank+1) * batch_size]

    # 2. 前向传播（本地计算）
    loss = model(x_local)

    # 3. 反向传播（计算本地梯度）
    loss.backward()

    # 4. 同步梯度（通信！）
    all_reduce(gradients)  # 所有 GPU 得到平均梯度

    # 5. 更新参数（所有 GPU 同步更新）
    optimizer.step()

通信开销

若梯度 payload 为 \(S\) bytes，\(N\)-rank ring All-Reduce 的每-rank 通信量约为 \(2(N-1)S/N\)。这说明 DDP 的主要通信量与参数数量同阶；实际时间仍需用有效 collective 带宽和实测 overlap 判断，不能直接除以某条链路的峰值带宽。

# DDP 示意：把一个大 batch 拆成多个 micro-batch
optimizer.zero_grad()
for micro_batch in micro_batches[:-1]:
    with model.no_sync():
        (model(micro_batch) / len(micro_batches)).backward()

# 最后一次 backward 不使用 no_sync()，触发 gradient bucket all-reduce
(model(micro_batches[-1]) / len(micro_batches)).backward()
optimizer.step()

DDP vs FSDP

DDP 解决了吞吐扩展，却没有解决模型状态冗余：每个 rank 仍保存完整参数、梯度和优化器状态。若这些状态本身放不下，就需要进一步分片，这正是下一讲 FSDP / ZeRO 的主题。

Tensor Parallelism (TP)

DP 让不同 GPU 处理不同样本，但每张 GPU 仍要独立完成整层计算。如果单层权重或中间激活太大，就必须把一次矩阵计算本身拆开；这会把通信从“每步一次梯度同步”推进到“几乎每个 Transformer block 都要交换 activation”。

核心思想

把模型的宽度切分到多个 GPU——每个 GPU 只持有参数的一部分。

以 \(Y=XW\) 为例，按输出维切分 \(W\) 后，每个 rank 只计算一部分输出；按输入维切分时，每个 rank 得到部分和，随后需要 collective 合并。Megatron-LM 将两种切法配对，让部分中间结果尽量留在本地，只在 block 的关键边界通信。

关键代价：TP 的 activation collective 通常每层发生并位于关键路径，因此倾向部署在 NVLink/NVSwitch 等高带宽、低延迟的 scale-up domain。跨节点 TP 并非不可能，但是否划算取决于拓扑、消息大小和实现。

FFN / SwiGLU:
# W_gate 与 W_up 按输出维切分（column parallel）
# 每个 rank 保留自己的中间分片，先在本地做 Swish(gate) * up
# W_down 按输入维切分（row parallel），最后合并各 rank 的部分和

Attention 的切分:
# W_q, W_k, W_v 按 head 切分（每个 GPU 负责部分 head）
# W_o 按输入维切分，最后合并各 rank 的部分和
#
# 基础 Megatron TP 常用 All-Reduce 合并；配合 sequence parallel 时，
# 相邻区域可改用 Reduce-Scatter / All-Gather 组合。

Transformer block 的 TP 通信:
1. Attention:
#    - Q, K, V 计算: 无需通信（head 独立）
#    - W_o: All-Reduce

2. FFN:
#    - W_gate, W_up 和逐元素门控: 本地分片计算
#    - W_down 后: All-Reduce（或 sequence-parallel 变体中的 Reduce-Scatter）

总计: 每个 block 2 次 All-Reduce

假设 d_model = 4096, num_heads = 32
# 每 GPU: 8 个 head, d_model_per_gpu = 1024

# W_q: (4096, 4096) → 切成 4 份: 每 GPU (4096, 1024)
# W_k: (4096, 4096) → 切成 4 份: 每 GPU (4096, 1024)
# W_v: (4096, 4096) → 切成 4 份: 每 GPU (4096, 1024)
# W_o: (4096, 4096) → 切成 4 份: 每 GPU (1024, 4096)

Pipeline Parallelism (PP)

TP 在每层切宽度，因此通信频繁。另一种办法是让每张 GPU 持有一段连续层，只在 stage 边界交接 activation；通信更粗粒度，但不同 stage 可能互相等待，形成 pipeline bubble。

核心思想

把模型的深度切分到多个 GPU——每个 GPU 持有连续的几层。

Pipeline Parallel: 4 GPU, 16 层
GPU 0: layers 0-3
GPU 1: layers 4-7
GPU 2: layers 8-11
GPU 3: layers 12-15

Input → GPU 0 → activations → GPU 1 → ... → GPU 3 → Output

把 batch 拆成多个 micro-batch 后，不同 stage 可以同时处理不同 micro-batch。micro-batch 越多，越容易填满流水线；但调度、activation memory 和 stage balance 也更复杂。

朴素 pipeline (4 stages, 1 batch):
Time →
GPU 0: [====L0-3====]  idle   idle   idle
GPU 1:  idle  [====L4-7====]  idle   idle
GPU 2:  idle   idle  [====L8-11===]  idle
GPU 3:  idle   idle   idle  [====L12-15==]

# 对理想化流水线，bubble time / useful compute ≈ (p-1) / m
# 若写成总时间占比，则约为 (p-1) / (m+p-1)
# p = pipeline stages, m = micro-batches
# 当 m >> p 时，bubble 可以忽略

Pipeline Parallel 的通信:
每对相邻 GPU 之间传 activation；反向再传 activation gradient。
bf16 activation bytes = batch_size × seq_len × d_model × 2

# 例如: 一个 bf16 activation tensor，B=4, S=4096, D=8192
# = 4 × 4096 × 8192 × 2 bytes = 256 MiB

# 反向还要传 activation gradients；若叠加 TP/sequence sharding，张量大小会变化。
# P2P 通信可以与其他 micro-batch 的计算 overlap，但受 schedule 和依赖约束。

Sequence Parallelism (SP)

TP 已经切分矩阵乘法，却不一定切分 LayerNorm、Dropout 等逐 token 区域的 activation。为了让这部分显存也随并行度下降，需要继续沿序列维分片。

动机

Tensor parallel 主要切分矩阵乘法的宽度维；LayerNorm、Dropout 等逐 token 区域仍可能复制 activation。Megatron 的 sequence parallel 用序列维分片这些区域，以继续降低 activation memory。

核心思想

1. LayerNorm / Dropout 是逐 token 独立的，可以自然地沿序列维并行。
2. 在 row-parallel 与 column-parallel 区域之间，用 Reduce-Scatter 和 All-Gather 改变 activation 的分片布局。
3. 相关 activation memory 因而可近似随 TP degree 降低。

这不同于为超长上下文设计的 context parallel / ring attention；后者会分片 attention 所需的序列上下文，并引入另一组通信模式。

混合并行策略

没有一种并行策略同时满足“最省显存、最少通信、没有 bubble”。真实集群通常先根据拓扑选择高频通信的局部并行度，再用其他维度扩展容量和吞吐。

3D / 4D 并行

常见分解同时使用 DP（不同数据）、TP（层内矩阵切分）、PP（层间切分），并让 SP 配合 TP 分片部分 activation：

\[N_{\mathrm{GPU}}=N_{\mathrm{DP}}\times N_{\mathrm{TP}}\times N_{\mathrm{PP}}\]

并行策略选择

训练 70B 模型，16 节点 (128 GPU):
TP = 8 (节点内)
PP = 4 (4 个 stages)
DP = 128 / (8 × 4) = 4 (数据并行)
总计: 8 × 4 × 4 = 128 GPU

# 仅参数 shard（未考虑 embedding tie、负载不均等）:
# 70B / (TP 8 × PP 4) ≈ 2.19B 参数/GPU
#
# 但不能据此宣布“放得下”：梯度、master weights、Adam moments 是否再按 DP 分片，
# 以及 activation、temporary buffers、pipeline schedule 都会改变峰值内存。

本讲知识地图：每种策略在交换什么

知识检查

如果只能记住 5 件事

1. 多 GPU 的核心新成本是 communication；先问交换什么数据，再讨论并行名称。
2. Collective 是分布式训练的数据变换原语：DDP 主要使用 All-Reduce，FSDP 主要使用 All-Gather 与 Reduce-Scatter。
3. 逻辑通信量不等于运行时间；bytes 必须结合 latency、effective bandwidth、topology 和 overlap 解读。
4. TP 的通信高频且常在关键路径，因此倾向放在高带宽、低延迟的 scale-up domain；PP 通信更粗粒度，但要管理 bubble。
5. 真实系统通常组合 DP / FSDP / TP / PP / SP：没有固定最优配置，选择取决于模型 shape、显存、batch、网络和实现。

延伸阅读

• Korthikanti et al. (2023): Reducing Activation Recomputation in Large Transformer Models — 并行策略组合
• Narayanan et al. (2021): Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
• Rajbhandari et al. (2020): ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• Stanford CS336 Spring 2025: Lecture 7 - Parallelism Basics
• Colossal-AI — 并行策略可视化