CS336 · 第 8 课
上一讲说明了如何把计算分到多张 GPU；这一讲处理更现实的容量约束：即使计算可以并行，每个 data-parallel rank 仍可能保存一整份参数、梯度和优化器状态。ZeRO / FSDP 的目标是消除这些冗余，但代价是按需通信。
Spring 2026 官方来源：Lecture 8 PDF · 课程日程
本页是中文学习导读；讲次主题、教师与材料版本以 Stanford 官方页面和讲义为准。
DDP 通过复制模型来扩展吞吐，因此同步后每个 rank 拥有相同参数、梯度与 optimizer state。对大模型而言，问题不是集群总显存不够，而是这些相同状态被重复存了很多份。
DDP 状态冗余 → ZeRO 分阶段分片 optimizer / gradient / parameter → FSDP 按 unit All-Gather 参数、Reduce-Scatter 梯度 → mixed precision 与 checkpointing 继续压缩峰值 → 与 TP / PP 组合成多维并行。
上一课知道 DDP 每卡有完整副本。本课逐项问：parameter、gradient、Adam \(m/v\) 在 forward、backward、update 的哪个时刻必须完整？若某对象只在局部时刻需要，就可能其余时间只保存 shard。
parameter gradient optimizer m/v
DDP copy copy copy
ZeRO-1 copy copy shard
ZeRO-2 copy shard shard
ZeRO-3 / FSDP FULL_SHARD
shard shard shard
steady-state shard ≠ peak memory during all-gather
state sharding ≠ activation shardingZeRO-3 steady state 可让每 rank 常驻约 3 个 parameter elements。但某 FSDP unit forward 前，要 All-Gather 该 unit 的完整参数进行 matmul，计算后 reshard；backward 又可能 gather，随后 Reduce-Scatter gradients。由此可见“除以 4”描述常驻份额，不自动描述峰值。
Stage 3 容量最好，却增加 parameter materialization、通信次数与管理开销；小模型或慢网络上 DDP 可能更快。Activation checkpointing 也不是 sharding：它丢弃部分 forward activations，backward 时重算，用额外 FLOPs 换显存。
在标准 Data Parallelism 中，每个 GPU 都存了完整副本：
每个 GPU 存储（70B 模型, fp32 + AdamW）:
参数: 4 × 70B = 280 GB
梯度: 4 × 70B = 280 GB
优化器 m: 4 × 70B = 280 GB
优化器 v: 4 × 70B = 280 GB
─────────────────────────────
总计: 1120 GB ← 远超单卡 80 GB！
# 问题: 8 个 GPU 各存一份 → 8 × 1120 GB = 8960 GB
# 其中 7/8 是冗余的！核心洞察：在 DDP 中，每个 GPU 最终得到的是相同的梯度（All-Reduce 后），然后用相同的梯度更新参数。这意味着每个 GPU 的优化器状态在每一步都是相同的——完全冗余！
ZeRO（Zero Redundancy Optimizer, Rajbhandari et al. 2020）通过分片消除冗余：
| Stage | 分片什么 | 每 rank 持久模型状态 | 主要通信 |
|---|---|---|---|
| Stage 1 | 优化器状态 | 参数/梯度复制，optimizer states 为 1/N | Reduce-Scatter 梯度 + All-Gather 更新后参数；带宽量可与 DDP 同阶 |
| Stage 2 | + 梯度 | 参数复制，梯度与 optimizer states 为 1/N | 分层 Reduce-Scatter + 参数 All-Gather |
| Stage 3 | + 参数 | 参数、梯度、optimizer states 近似 1/N | 按需参数 All-Gather + 梯度 Reduce-Scatter |
分片越深入，持久显存越低，但计算前越需要重新取得当前参数。Stage 3 的关键代价是：每个 FSDP unit 通常在前向、反向分别 All-Gather 参数，并在反向 Reduce-Scatter 梯度；能否隐藏这些通信取决于 unit 粒度、计算量、网络和 prefetch。
标准 DDP:
# 每个 GPU 存完整的 m 和 v（各 4P bytes）
# ZeRO Stage 1:
# 把 m 和 v 均分到 N 个 GPU
# GPU i 只存 m[i*N/N : (i+1)*N/N] 和 v[...]
# 每个 GPU 内存:
# 参数: 4P（完整）
# 梯度: 4P（完整）
# 优化器: 8P/N（分片）
# 总计: 8P + 8P/N bytes
# 8 GPU: 8P + P = 9P（vs 标准 16P）→ 节省 44%ZeRO Stage 2:
# 梯度也分片——每个 GPU 只存自己负责的那部分参数的梯度
# 关键: All-Reduce → Reduce-Scatter
# Reduce-Scatter 让每个 GPU 只得到自己那部分的聚合梯度
# 每个 GPU 内存:
# 参数: 4P（完整）
# 梯度: 4P/N（分片）
# 优化器: 8P/N（分片）
# 总计: 4P + 12P/N bytes
# 8 GPU: 4P + 1.5P = 5.5P → 节省 66%ZeRO Stage 3:
# 参数也分片——每个 GPU 只存 1/N 的参数
# 前向/反向计算时:
# 1. All-Gather 获取完整参数
# 2. 计算
# 3. 丢弃非本地的参数分片
# 每个 GPU 内存:
# 参数: 4P/N（分片）
# 梯度: 4P/N（分片）
# 优化器: 8P/N（分片）
# 总计: 16P/N bytes
# 8 GPU: 2P → 节省 87.5%！70B 模型, 8 GPU, fp32 + AdamW
标准 DDP: 16P = 1120 GB/GPU ← 放不下
ZeRO Stage 1: 9P = 630 GB/GPU ← 还是放不下
ZeRO Stage 2: 5.5P = 385 GB/GPU ← 还是放不下
ZeRO Stage 3: 2P = 140 GB/GPU ← 还是放不下（但接近了！）
# 加上 bf16 混合精度:
# 参数 bf16: 2P/N
# 主参数 fp32: 4P/N
# 梯度 bf16: 2P/N
# 优化器 fp32: 8P/N
# 总计: 16P/N bytes
# 8 GPU: 2P = 140 GB → 还是超
# 16 GPU: 持久模型状态约 70 GB/GPU
# 这仍未计入完整参数 All-Gather buffer、activation、allocator 碎片和通信 workspace，
# 因此不能仅凭 70 < 80 GB 断言一定放得下。FSDP (Fully Sharded Data Parallel) 是 PyTorch 原生的 fully sharded data parallelism。其 FULL_SHARD 会分片 parameters、gradients 与 optimizer states，因此状态分片行为与 ZeRO Stage 3 的核心思想高度类似；但两者不是同一个框架实现，也不应把所有 FSDP 配置都等同于 ZeRO-3。
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
# 包装模型
model = FSDP(
model,
sharding_strategy=ShardingStrategy.FULL_SHARD, # ZeRO Stage 3
cpu_offload=CPUOffload(offload_params=False),
auto_wrap_policy=transformer_auto_wrap_policy,
)这解释了 FSDP 的核心权衡：它用短暂的完整参数 buffer 和更多 collective，换取大部分时间只保存 \(1/N\) 的模型状态。
令某个 FSDP unit 的参数/梯度 payload 为 S bytes、world_size=N：
# 前向 All-Gather: (N-1)/N × S per rank
# 反向参数 All-Gather: (N-1)/N × S per rank
# 反向梯度 Reduce-Scatter: (N-1)/N × S per rank
# 总计: 3(N-1)/N × S，N 大时约 3S
# 同一 payload 的 ring DDP All-Reduce 约为 2(N-1)/N × S
# FSDP 通信更多，但内存更少
# 适合: 模型太大，DDP 放不下时| 策略 | 描述 | 内存 | 通信 |
|---|---|---|---|
| FULL_SHARD | 参数、梯度、optimizer state 全分片；行为与 ZeRO-3 高度类似 | 持久模型状态最低 | 参数 gather 最频繁 |
| SHARD_GRAD_OP | 反向后分片梯度与 optimizer states；前向参数保持未分片 | 介于 DDP 与 FULL_SHARD | 依 PyTorch 版本和 reshard 行为而定 |
| NO_SHARD | 不分片；新代码通常直接使用 DDP | 模型状态复制 | 梯度 All-Reduce |
| HYBRID_SHARD | 组内 FULL_SHARD，组间复制 | 按 shard group 缩小 | 组内 gather/scatter + 组间同步 |
HYBRID_SHARD（拓扑相关选项）：
混合精度训练 (bf16 前向 + fp32 主参数):
参数:
bf16 参数: 2P（用于前向/反向计算）
fp32 主参数: 4P（用于优化器更新）
梯度:
fp32 梯度: 4P（用于优化器）
优化器:
fp32 m: 4P
fp32 v: 4P
总计: 18P bytes（vs fp32 的 16P，稍多但激活值省一半）
FSDP + 混合精度:
参数、梯度和 optimizer states 的持久 shard 可近似为上述预算的 1/N
8 GPU, 70B 模型:
每 GPU: 18P / 8 = 2.25P = ~157 GB → 还是超
16 GPU: 18P / 16 = ~79 GB 的持久状态下界
实际峰值还包括 All-Gather 后的完整 FSDP unit、activations、临时 buffer 与碎片；
79 GB 持久状态不可能安全装入 80 GB GPU，需要更细 sharding、offload、
更低状态精度或更多设备，并常配合 activation checkpointing。
用计算换内存——不保存中间激活，反向时重算：
from torch.utils.checkpoint import checkpoint
# 不 checkpoint: 保存较多前向激活，通常随层数 L 线性增长
# checkpoint: 只保存选定边界，反向时重算区间内部
# 内存/计算权衡取决于分段策略，不能统一写成必然 O(√L) 或 O(1)
# PyTorch 中 activation checkpointing 与 FSDP wrapping 是两个步骤；
# 可使用 torch.utils.checkpoint.checkpoint，或 distributed checkpoint_wrapper
# 按 Transformer block 应用，而不是传入不存在的 FSDP 构造参数。Ring All-Reduce 的两个阶段:
Phase 1: Reduce-Scatter
每个 GPU 把数据分成 N 份
在 ring 中传递并累加
经过 N-1 步，每个 GPU 有 1/N 的完整结果
Phase 2: All-Gather
每个 GPU 把自己的 1/N 结果在 ring 中传递
经过 N-1 步，每个 GPU 有完整结果
总通信量: 2 × (N-1)/N × data_size ≈ 2 × data_size
固定 payload 时，每-rank 字节量是 O(1) 并趋近 2 × data_size；
但需要 2(N-1) 个 ring steps，因此 latency 项仍随 N 增长。
# NCCL 会按 collective、消息大小、拓扑和协议选择/组合算法:
# - Tree 往往有更少的通信 steps，适合 latency 敏感的小消息
# - Ring 往往更容易逼近链路带宽，适合 bandwidth 敏感的大消息
# - 节点内可利用 NVLink/NVSwitch，跨节点再结合网络与分层算法
# - 节点间: 利用 InfiniBand/RoCE
# PyTorch 接口:
import torch.distributed as dist
# All-Reduce
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
# Reduce-Scatter
dist.reduce_scatter_tensor(output, input, op=dist.ReduceOp.SUM)
# All-Gather
dist.all_gather_into_tensor(output_tensor, input_tensor)# 简化的 NCCL benchmark 示意；正式测量优先使用 nccl-tests 并多次迭代
tensor = torch.randn(100 * 1024**2, device='cuda') # 400 MiB (float32)
for _ in range(5): # warmup
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(20):
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
torch.cuda.synchronize()
duration = (time.perf_counter() - start) / 20
size_bytes = tensor.element_size() * tensor.numel()
algorithm_bandwidth = (
2 * (world_size - 1) / world_size * size_bytes / duration
)
# 这是按 ring 字节模型换算的 algorithm bandwidth，不等于单条链路规格。
# 结果必须连同 GPU/NIC、拓扑、world_size、dtype 和 NCCL 版本一起报告。Spring 2026 官方 Lecture 8 不把 ZeRO/FSDP 当作终点，而是把它与模型不同轴的切分组合起来：
| 轴 | 切分对象 | 主要通信 | 常见放置 |
|---|---|---|---|
| FSDP / ZeRO | 参数、梯度、optimizer state | all-gather + reduce-scatter | 可跨节点扩展 |
| Tensor Parallel | 单层矩阵的宽度维 | 每层 activation all-reduce / reduce-scatter / all-gather | 优先节点内高速互联 |
| Sequence Parallel | LayerNorm/Dropout 等非 TP activation 的序列维 | 与 TP 边界配合的 reduce-scatter / all-gather | 通常与 TP 同组 |
| Pipeline Parallel | 网络深度/层 | stage 间 P2P activation 与 gradient | 可跨相对较慢网络 |
| Expert Parallel | MoE experts 与 token | token dispatch / combine 的 all-to-all | 优先低延迟、高带宽域 |
2D parallelism 常指 FSDP × TP 网格；更完整的 3D/4D 方案再加入 PP、EP/SP。总设备数可写成各独立并行轴 degree 的乘积，但前提是分组定义、batch/sequence 可整除，并且每一轴的通信资源没有互相争用。
# 若两条通信轴使用独立资源并能与计算重叠：
step_time ≈ max(compute, fsdp_collective, tp_collective)
# 若共享同一 NIC / NVLink channel 且不能重叠：
step_time ≈ compute + fsdp_collective + tp_collective实际框架会用 bucket、prefetch、异步 collective、microbatch 与 schedule 尝试隐藏通信；因此并行方案必须用 profiler 观察时间线，而不能只比较理论 communication volume。
# 硬件: 16 节点 × 8 H100 = 128 GPU
# 并行策略:
TP = 8 (节点内)
PP = 4 (4 个 pipeline stages)
DP = 128 / (8 × 4) = 4
# 内存（每 GPU）:
参数: 70B / (8 × 4) = 2.19B 参数
bf16: 4.4 GB
fp32 主参数: 8.8 GB
梯度: 8.8 GB
优化器: 17.6 GB
# activation、通信 buffer 和峰值参数 gather:
# 必须用实际 shape、FSDP unit 边界与 profiler 测量
# 吞吐量:
MFU ≈ 0.5
每 GPU: 989 × 0.5 = 495 TFLOP/s
128 GPU: 63.4 PFLOP/s
训练 15T tokens: 6 × 70B × 15T / 63.4P
≈ 9.94e7 秒 ≈ 1150 天
# 这个结果说明 128 张 H100 不足以在数天内完成如此大的 token 预算；
# 原先“99 小时”低估了约 280 倍。真实规划还要加入非训练 FLOPs 与停机。理想: N 个 GPU → N× 加速
实际: 通信开销导致亚线性加速
弱 scaling 固定每 GPU 工作量，强 scaling 固定总工作量。
效率没有跨模型/集群通用的 80-95% 或 60-80% 区间，必须报告
GPU 数、网络、global batch、并行分解、序列长度和测量口径。
影响 scaling 的因素:
1. 模型大小（越大，计算占比越高，scaling 越好）
2. Batch size（越大，通信占比越低）
3. 网络带宽（NVLink vs InfiniBand）
4. 并行策略选择
问题 1：理想均匀分片下，ZeRO Stage 3 的持久模型状态约为标准复制的多少？
问题 2：FSDP 前向传播需要什么通信？
问题 3：FSDP 的 HYBRID_SHARD 策略是指？
问题 4：固定 payload、忽略 latency 时，Ring All-Reduce 每 rank 的带宽通信量随 N 如何增长？
问题 5：Gradient checkpointing 的核心思想是什么？
ZeRO-3 在 4 ranks 上把参数常驻份额约降到 1/4，某 FSDP unit 计算时该 rank 是否永远只持有 1/4 unit 参数？
答案：否；计算前通常 All-Gather 完整 unit，因此峰值包含 materialized unit。
若答“永远 1/4”，你把 steady-state shard 账本当成了时间轴上的 peak 账本。
系统已经给出可承受的模型、token 与计算预算；下一课反过来问：固定 compute 时，参数量 \(N\) 与数据量 \(D\) 应如何分配才可能获得更低 loss。