Lecture 7 · 并行 I：基础

CS336: Language Modeling from Scratch · Stanford · Spring 2025
📅 Apr 22 · 💻 不可执行讲义（PDF） · 讲者 Tatsunori Hashimoto

---

承上启下

上一讲（Lecture 6 · Kernel 与 Triton）：

• 单 GPU 内部的核心矛盾是「算力远离数据」：HBM、SRAM、寄存器之间搬数据的代价远高于计算本身，多数算子实际上是内存受限（memory-bound）而非算力受限。
• 解法是通过 kernel fusion、tiling、用 Triton/CUDA 手写 kernel，把中间结果尽量留在片上，减少对 HBM 的往返读写。
• Benchmark 与 profiling 是性能优化的基本工作流：先测量、找瓶颈、再优化，不凭感觉。

本讲（并行基础）：

• 单 GPU 优化做得再好也不够：一是模型太大，训练状态单卡显存装不下；二是训练太慢，单卡算力跑一次预训练要几百年。
• 于是进入多 GPU / 多节点的世界。此时新的瓶颈从「片内搬数据」变成「跨 GPU 搬数据」——通信成为一等公民。
• 本讲建立概念地图：集合通信原语、数据并行、ZeRO/FSDP、流水线并行、张量并行、序列并行，以及为什么真实的大模型训练最后总是把它们组合成 3D/混合并行。
• 下一讲 Lecture 8 · 并行 II：分布式训练实战 会把这张地图上的每条路都用最小可运行代码走一遍。

---

1. 为什么必须并行？

1.1 算力限制：单卡训练要几百年

由 Lecture 2 · PyTorch 与资源核算 的资源核算，Transformer 训练的总计算量近似为：

$$C \approx 6ND$$

其中 $N$ 是参数量，$D$ 是训练 token 数，系数 $6$ 来自每个参数每个 token 前向约 $2$ FLOPs、反向约 $4$ FLOPs。

1.2 内存限制：每参数 16 字节

训练时显存不只存参数。标准的混合精度 + AdamW 配置下，每个参数需要：

内容	朴素 AdamW 每参数字节
BF16 参数	2
BF16 梯度	2
FP32 master weights	4
Adam 一阶矩 $m$	4
Adam 二阶矩 $v$	4
合计	16

即参数量为 $\Psi$ 时，仅「参数 + 梯度 + 优化器状态」就要约 $16\Psi$ 字节——70B 参数对应 $16 \times 70{\times}10^9 \approx 1.12$ TB，而单张 H100 只有 80 GB HBM。这还没算激活（activations），后者随 batch size 和序列长度增长，长上下文时往往比参数状态更大（见第 7 节）。

1.3 新计算单位：数据中心

单 GPU 内部的存储层级（Lecture 5 · GPU）：

register → shared memory → L2 → HBM

多 GPU 时代，这个层级向外延伸：

GPU HBM → NVLink（节点内）→ NVSwitch / InfiniBand（跨节点）→ 数据中心网络

数量级感受：H100 的 HBM 带宽约 3.9 TB/s，节点内 NVLink 约 900 GB/s，跨节点网络又低一到两个数量级。层级越往外，容量越大，但带宽越低、延迟越高。

---

2. 集合通信：分布式训练的词汇表

多 GPU 程序不应该手写「GPU 0 发给 GPU 1，GPU 2 等 GPU 3」这样的点对点逻辑——既容易死锁，也无法针对硬件拓扑优化。经典做法（源自 1980 年代的 HPC/MPI 传统）是使用集合通信（collective operations）：一次性声明「所有设备共同完成什么通信模式」，由底层库（NCCL/MPI）选择最优路由。

两个基本术语：world size 指参与的设备总数 $W$；rank 指每个设备的编号 $0,1,\dots,W-1$。

操作	含义	典型用途
Broadcast	一个 rank 的数据发给所有 rank	分发初始参数
Scatter	一个 rank 把切片分给各 rank	分发数据
Gather	各 rank 数据收集到一个 rank	汇总结果
Reduce	各 rank 数据做 sum/max 等聚合到一个 rank	汇总 loss
All-gather	gather 后每个 rank 都拿到完整结果	恢复完整参数/激活
Reduce-scatter	reduce 后把结果切片分给各 rank	梯度分片聚合
All-reduce	reduce 后所有 rank 都拿到完整结果	DDP 梯度同步

关键恒等式：

$$\text{all-reduce} \;=\; \text{reduce-scatter} \;+\; \text{all-gather}$$

即「先把归约结果切片分到各卡，再把各卡的切片拼回完整结果」等价于一次 all-reduce。

---

3. 数据并行：最简单，也最浪费内存

3.1 Naive Data Parallel

思想：复制模型、切分数据。

• 每张 GPU 放一份完整的模型副本（参数、梯度、优化器状态全都复制）；
• 全局 batch 切成 $W$ 份，每张卡吃自己那一份；
• 每张卡独立前向、反向，得到本地梯度 $g_w$；
• 用 all-reduce 求平均梯度 $\bar g = \frac{1}{W}\sum_{w=1}^{W} g_w$，使所有卡拿到同一个梯度；
• 每张卡独立执行 optimizer step——由于初始参数相同、梯度相同，各副本参数永远保持一致。

$$\bar g = \frac{1}{W}\sum_{w=1}^{W} g_w$$

其中 $g_w$ 是 rank $w$ 在本地 mini-batch 上算出的梯度。平均后的 $\bar g$ 在数学上完全等价于用整个全局 batch 做一次大 batch SGD——数据并行不改变优化轨迹，只改变计算的物理分布。

graph LR
    D0["data shard 0"] --> M0["model copy"]
    D1["data shard 1"] --> M1["model copy"]
    D2["data shard 2"] --> M2["model copy"]
    M0 --> G["all-reduce gradients"]
    M1 --> G
    M2 --> G

优点：

• 实现最容易，对模型代码几乎零侵入（PyTorch 的 DDP 就是给 backward 挂 hook 自动做 all-reduce）；
• 计算扩展性好：每卡的计算量随 $W$ 线性下降；
• 只要每卡的本地 batch 足够大，通信开销可以被计算摊薄，还能与反向传播重叠。

缺点：

• 每张卡都存完整的参数、梯度、优化器状态（$16\Psi$ 字节），内存完全不随 GPU 数下降——模型单卡装不下时数据并行独木难支；
• 全局 batch size 不能无限增大：超过临界 batch size（critical batch size）后，增大 batch 换来的「每步进步」收益递减，等价于浪费算力。梯度噪声尺度（gradient noise scale）刻画了这个上限，这也是纯数据并行扩展的统计瓶颈。

3.2 DDP 的通信量

每步需要同步一次完整梯度。设梯度总大小为 $G$ 字节（BF16 下 $G = 2\Psi$），环式 all-reduce 中每卡每步的收发量约为：

$$2\cdot\frac{W-1}{W}\cdot G \;\approx\; 2G$$

即与模型大小同阶、与 batch size 无关。实践中 DDP 还会做两件事摊薄这笔开销：分桶（bucketing）——把许多小参数的梯度攒成大桶再通信，减少启动延迟；通信计算重叠——某层梯度一算完就异步发起 all-reduce，与更早层的反向计算并行。

---

4. ZeRO：把数据并行的状态切开

ZeRO（Zero Redundancy Optimizer）的出发点：数据并行里每张卡存的 $16\Psi$ 字节状态在 $W$ 张卡上重复了 $W$ 遍，纯属冗余——能不能切片存、用时取？ZeRO 按「切多少」分三个 stage，切得越多越省内存、实现越复杂。

4.1 ZeRO Stage 1：切 optimizer state

每张卡：

• 保留完整的 BF16 参数与完整梯度；
• 只保存 $1/W$ 的优化器状态（FP32 master weights + Adam 的 $m, v$）；
• 反向后用 reduce-scatter 让每卡拿到自己负责那片参数的平均梯度，只更新对应的参数切片；
• 更新完后 all-gather 参数，让所有卡重新拥有完整的最新参数。

对比	DDP	ZeRO-1
参数	每卡完整	每卡完整
梯度	每卡完整	每卡完整
Optimizer state	每卡完整	切分
通信	all-reduce 梯度	reduce-scatter 梯度 + all-gather 参数
每卡内存	$16\Psi$	$4\Psi + 12\Psi/W$

由于 all-reduce 本来就等于 reduce-scatter + all-gather（第 2 节恒等式），ZeRO-1 的通信量与 DDP 相同（约 $2\times$ 模型大小）——只是把「归约梯度」和「取回参数」两步拆开，在中间插入了本地的优化器更新。

4.2 ZeRO Stage 2：再切梯度

ZeRO-2 在 Stage 1 基础上把梯度也切分：

• 反向传播时，每层梯度一算出来就立刻 reduce-scatter，归约结果只留在负责该分片的卡上；
• 任何时刻都不需要保存完整的梯度向量，每卡内存降到 $2\Psi + 14\Psi/W$；
• 梯度显存可以尽早释放，reduce-scatter 还可以与后续层的反向计算重叠。

代价是实现复杂度：不能再随意实例化完整梯度张量（例如做全局 grad-norm 裁剪要先各自算分片范数再 all-reduce 标量），backward hook 与参数分桶的逻辑也更精细。

4.3 ZeRO Stage 3 / FSDP：参数也切

ZeRO-3 把三样全切：参数、梯度、优化器状态各留 $1/W$，每卡内存降到 $16\Psi/W$——内存终于随 GPU 数线性下降。代价是参数不再常驻，每层用时要现取：

前向时（逐层）：

1. 进入第 $i$ 层前，all-gather 该层的参数分片，拼出完整参数；
2. 本地前向计算；
3. 用完立即释放非本卡的参数分片（只留 $1/W$）。

反向时（逐层，逆序）：

1. 再次 all-gather 该层参数（前向后没有保留完整版）；
2. 计算本层梯度；
3. reduce-scatter 梯度，每卡只留自己负责的梯度分片；
4. 释放完整参数。

sequenceDiagram
    participant G0 as GPU 0
    participant G1 as GPU 1
    participant G2 as GPU 2
    Note over G0,G2: 前向：进入第 i 层前
    G0->>G1: all-gather 第 i 层参数分片
    G1->>G2: all-gather 第 i 层参数分片
    G2->>G0: all-gather 第 i 层参数分片
    Note over G0,G2: 各卡拼出完整第 i 层 → 本地计算 → 释放非本卡分片
    Note over G0,G2: 反向：再次 all-gather 第 i 层参数并计算梯度
    G0->>G1: reduce-scatter 第 i 层梯度
    G1->>G2: reduce-scatter 第 i 层梯度
    G2->>G0: reduce-scatter 第 i 层梯度
    Note over G0,G2: 每卡只留自己的梯度分片，更新对应参数分片

通信量：前向 all-gather 参数（$1\times$）+ 反向 all-gather 参数（$1\times$）+ 反向 reduce-scatter 梯度（$1\times$），共约 $3\times$ 模型大小——是 DDP 的 $1.5$ 倍。工程上靠预取（prefetch）掩盖：算第 $i$ 层时就异步 all-gather 第 $i{+}1$ 层，通信藏在计算后面。

FSDP（Fully Sharded Data Parallel）是 PyTorch 对 ZeRO-3 思想的原生实现，Lecture 8 · 并行 II：分布式训练实战 有其 API 用法。

4.4 ZeRO 的边界

ZeRO 解决的是参数 / 梯度 / 优化器状态这三块内存，它不解决：

• 激活内存：激活随 batch 与序列长度增长，需要 activation checkpointing 或序列并行（第 7 节）；
• 计算扩展的统计极限：ZeRO 本质仍是数据并行，每卡本地 batch 太小时计算效率差，全局 batch 又受临界 batch size 约束；
• 通信延迟与跨节点带宽：ZeRO-3 每层都要 all-gather，跨节点低带宽下预取藏不住通信，吞吐会明显掉——这正是要引入模型并行（流水线/张量）的原因。

---

5. 流水线并行：按深度切模型

5.1 Layer-wise Parallel 的低利用率

最直接的模型并行：GPU 0 放前几层，GPU 1 放后几层，激活在卡间接力。问题在于计算是串行依赖的——GPU 0 算的时候 GPU 1 干等，GPU 1 算的时候 GPU 0 干等。$p$ 张卡的利用率只有 $1/p$：内存确实省了，算力全浪费了。

5.2 Pipeline Parallelism：用 micro-batch 填满流水线

解法：把 batch 切成 $m$ 个 micro-batch，像工厂流水线一样依次喂入。GPU 0 处理完 micro-batch 0 交给 GPU 1 后，立刻开始处理 micro-batch 1——各阶段错峰同时工作。

gantt
    title 2 阶段流水线示意
    dateFormat X
    axisFormat %L
    section GPU0
    MB0 F :a0, 0, 1
    MB1 F :a1, 1, 1
    MB2 F :a2, 2, 1
    section GPU1
    等待 :b0, 0, 1
    MB0 F :b1, 1, 1
    MB1 F :b2, 2, 1
    MB2 F :b3, 3, 1

流水线注满与排空的空闲时间称为气泡（bubble）。设流水线阶段数为 $p$、micro-batch 数为 $m$，气泡开销相对于有效计算时间的比例近似为：

$$\text{bubble ratio} \;\approx\; \frac{p-1}{m}$$

（若按占总时间的比例算则是 $\frac{p-1}{m+p-1}$。）直觉：注满流水线要 $p-1$ 个 micro-batch 时隙、排空又要 $p-1$ 个，而有效工作是 $m$ 个时隙，所以 $m \gg p$ 时气泡可忽略。

micro-batch 也不是越多越好：每个 micro-batch 太小会让 GPU kernel 效率下降（矩阵太瘦），且通信次数与调度开销增多。

5.3 为什么还要用流水线？

尽管有气泡，流水线并行仍有独特价值：

• 参数按层切分，每卡只存 $1/p$ 的层，参数内存线性下降；
• 通信量极小且只发生在相邻阶段之间：每个 micro-batch 只传一次层间激活（点对点 send/recv），远小于张量并行的每层通信——这使它成为跨节点低带宽链路上的首选模型并行方式；
• 对层数很多的深模型是自然的切法。

缺点与工程复杂度：

• batch 太小时气泡吃掉利用率；
• 前向/反向的调度、各 micro-batch 激活的保存、梯度累积都要精细处理——朴素 GPipe 调度要保存全部 $m$ 个 micro-batch 的激活，内存峰值高；
• 1F1B 调度（one-forward-one-backward）让每个阶段尽早交替做前向和反向，把在途激活数从 $m$ 降到 $p$ 量级；交错流水线（interleaved schedule）给每卡分多段不连续的层进一步缩小气泡；近年的 zero-bubble 调度把反向拆成「对输入求梯度」和「对权重求梯度」两步错开填缝，理论上能把气泡填平。

---

6. 张量并行：按宽度切模型

6.1 思想：把单个矩阵乘切开

流水线按深度切（不同层放不同卡），张量并行（tensor parallelism）按宽度切（同一层的大矩阵拆到多卡）。

对线性层 $Y = XW$，有两种切法：

• 按列切（column parallel）：$W = [W_1 \mid W_2]$，则 $Y = [XW_1 \mid XW_2]$——每卡拿完整输入 $X$，算出输出的一段列，结果需要拼接（all-gather）才是完整 $Y$；
• 按行切（row parallel）：$W = \begin{bmatrix} W_1 \\ W_2 \end{bmatrix}$，输入相应按列切 $X = [X_1 \mid X_2]$，则 $Y = X_1 W_1 + X_2 W_2$——每卡算出一个部分和，需要 all-reduce 相加才是完整 $Y$。

graph LR
    X["X"] --> W0["W0"]
    X --> W1["W1"]
    W0 --> Y0["Y0"]
    W1 --> Y1["Y1"]
    Y0 --> AG["all-gather / all-reduce"]
    Y1 --> AG
    AG --> Y["Y"]

Megatron-LM 的经典组合是列切接行切。以 MLP 块 $Z = \mathrm{GeLU}(XA)\,B$ 为例：

• 第一个线性层 $A$ 按列切：每卡得到 $\mathrm{GeLU}(XA_i)$——GeLU 是逐元素操作，对列分片可以独立施加，中间不需要任何通信；
• 第二个线性层 $B$ 按行切：恰好吃进上一步的列分片作为输入，各卡输出部分和；
• 整个 MLP 块前向只需最后一次 all-reduce。

Attention 的切法同理：多头天然按头切（每卡负责一部分头，QKV 投影按列切），输出投影按行切，同样整块只需一次 all-reduce。这样一个 Transformer 层的前向共 2 次 all-reduce（attention 一次 + MLP 一次），反向再 2 次，合计每层 4 次。

6.2 优缺点

优点：

• 没有流水线气泡，所有卡对同一个 micro-batch 同时工作；
• 单层的参数和计算被多卡分摊，能支撑单层大到一张卡放不下的模型（超大 hidden dim / FFN）；
• 与 NVLink 这类节点内高带宽互联是绝配。

缺点：

• 每层都要通信（且在关键路径上，难以完全用计算掩盖），对带宽和延迟都极其敏感；
• 因此实践中张量并行度通常 $\le 8$，被限制在一个节点（一台机器 8 卡 NVLink 全互联）之内；
• 通信内容是激活（大小 $\propto$ batch × seq × hidden），batch 大时这笔账不小（见 Lecture 8 · 并行 II：分布式训练实战 的数值估算）。

---

7. 激活内存与序列并行

7.1 激活不只是参数

反向传播需要保存前向的中间结果（activations）。激活内存的大小与 batch size $b$、序列长度 $s$、hidden dim $h$、层数、注意力实现方式、是否做 checkpointing 都有关。

结论：即使 ZeRO-3 把 $16\Psi$ 的参数状态切干净了，长序列大 batch 下激活仍可能爆显存。对付激活有两条路：重算（activation checkpointing / recomputation）——只存少量检查点，反向时重新前向一遍换取内存（约多花 1/3 计算）；切分——即下面的序列并行。

7.2 Tensor Parallel 的激活盲区

张量并行把大矩阵乘的输入/输出激活切了 $1/t$（$t$ 为张量并行度），但 LayerNorm、dropout、residual 这类逐点（pointwise）操作落在通信边界之外，每张卡仍持有完整的 $[b, s, h]$ 激活。对应上面公式：$34sbh$ 里约有 $10sbh$ 属于这些逐点项，不随 $t$ 缩小——这阻碍激活内存的线性扩展。

7.3 Sequence Parallel

序列并行（sequence parallelism）补上这块盲区：在逐点操作区域，激活沿序列维切开：

$$[b,\, s,\, h] \;\longrightarrow\; \text{每卡持有 } [b,\, s/t,\, h]$$

逐点操作对每个 token 独立，沿序列切完全不影响结果。与张量并行拼接时，原来的每次 all-reduce 被替换为一次 all-gather（进入张量并行区前把序列拼全）加一次 reduce-scatter（离开时按序列切回去）——由第 2 节恒等式，总通信量不变，但逐点区域的激活内存变成了 $1/t$。

收益：激活内存随卡数接近线性下降，与张量并行天然互补（Megatron 中两者总是一起开）。更长上下文的场景还有沿序列切注意力本身的 context parallelism / Ring Attention（第 8 节提及）。

---

8. 3D 并行：真实训练的组合拳

三大主轴：

轴	切分对象	解决什么	主要代价
Data Parallel	batch	扩计算	梯度 all-reduce
Tensor Parallel	width	单层参数/计算	高频层内通信
Pipeline Parallel	depth	模型深度/跨节点	气泡与调度

在三大主轴之上还能叠加：

• ZeRO / FSDP：给数据并行维度切训练状态；
• Sequence Parallel：给张量并行组切逐点激活；
• Expert Parallel：MoE 的专家分布到不同卡，token 经 all-to-all 路由（见 Lecture 4 · 混合专家模型 MoE）；
• Context Parallel / Ring Attention：超长上下文时把注意力计算本身沿序列切到多卡。

8.1 经验规则

1. 先让模型放进显存：按需组合 ZeRO/FSDP、张量并行、流水线并行、activation checkpointing——放不下什么都别谈。
2. 高频通信放高带宽链路：张量并行（每层多次通信）尽量收在 NVLink 全互联的单节点内。
3. 跨节点用低频通信的并行：流水线并行（每 micro-batch 一次点对点）或数据并行（每步一次、可与计算重叠）更耐低带宽。
4. batch 够大时数据并行最好用：通信可摊薄可重叠；batch 太小时 DDP 通信占比飙升，且受临界 batch size 的统计上限约束。
5. 激活内存单独治理：只切参数不够，序列并行与重算是长序列训练的标配。

8.2 真实模型案例

模型	并行策略特点
DeepSeek-V3	16 路流水线（DualPipe 调度）+ 64 路专家并行 + ZeRO-1 数据并行，不用张量并行
Llama 3 405B	4D 并行：TP=8（节点内）× PP=16 × 数据并行（FSDP 式），长上下文阶段再加 context parallel
Gemma 2	TPU 上 ZeRO-3 式状态分片 + 数据/模型并行（编译器自动分片）
Yi 系列	ZeRO-1 + 张量/流水线并行的 Megatron 式组合

---

总结

mindmap
  root((并行 I))
    Collective
      all-reduce
      all-gather
      reduce-scatter
    Data Parallel
      batch sharding
      gradient sync
      memory not sharded
    ZeRO
      stage 1 optimizer
      stage 2 gradients
      stage 3 params
      FSDP
    Pipeline
      depth sharding
      micro-batches
      bubble
    Tensor
      width sharding
      high bandwidth
      Megatron
    Activation
      checkpointing
      sequence parallel
    Hybrid
      3D parallel
      topology aware

关键要点：

1. 并行训练的目标是让内存和计算随卡数线性扩展——但通信总在阻止理想线性，所有策略都是围绕「省通信」设计的。
2. 集合通信是分布式训练的基本语言——all-reduce、all-gather、reduce-scatter 的数据流向必须烂熟，恒等式 all-reduce = reduce-scatter + all-gather 是 ZeRO 系列的基石。
3. ZeRO 不是新的模型并行，而是更聪明的数据并行——按 stage 依次切优化器状态（$4\Psi{+}12\Psi/W$）、梯度（$2\Psi{+}14\Psi/W$）、参数（$16\Psi/W$）；ZeRO-1/2 通信量不变，ZeRO-3 多付 $1.5\times$ 通信换全量切分。
4. 流水线并行省内存但有气泡——气泡比例约 $(p-1)/m$，micro-batch 数量是填气泡的关键；1F1B/交错/zero-bubble 调度是工程主战场。
5. 张量并行无气泡但通信频繁——「列切 → 逐元素激活 → 行切」让每个子块只通信一次；适合 NVLink 节点内，并行度一般 $\le 8$。
6. 激活内存同样重要——注意力分数项随 $s^2$ 增长；序列并行切逐点激活且通信量不变，与重算一起构成长序列训练的标配。
7. 真实系统是组合策略——数据并行 × 张量并行 × 流水线并行，再叠加 ZeRO、sequence/expert/context parallel，按硬件拓扑把「通信频率」与「链路带宽」对齐。

下一讲预告（Lecture 8 · 并行 II：分布式训练实战）：

本讲建立概念地图，下一讲用 torch.distributed 和 NCCL 写最小可运行代码：初始化进程组、跑 collective benchmark，并手写数据并行、张量并行、流水线并行的简化版本。

---

复习自测

---

参考资料

• 💻 2025 Lecture 7 - Parallelism basics.pdf（官方讲义）
• 📄 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• 📄 Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• 📄 GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
• 📖 PyTorch FSDP Tutorial
• 📄 Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
• 📄 Reducing Activation Recomputation in Large Transformer Models（序列并行与激活内存公式，Korthikanti et al. 2022）
• 📄 An Empirical Model of Large-Batch Training（梯度噪声尺度与临界 batch size，McCandlish et al. 2018）
• 📖 HuggingFace Ultra-Scale Playbook（各并行策略的交互式图解）