7 min read#distributed-training#fsdp#data-parallelism#zero#memory-optimization#pytorch

Fully Sharded Data Parallel

Apr 10, 2026

FSDP(Fully Sharded Data Parallel, 全分片数据并行) 是 PyTorch 对 ZeRO-3 的原生实现。与标准 DDP 不同，FSDP 将模型参数、梯度和优化器状态全部分片(shard) 存储到各个 rank 上，把单卡显存需求从 $O(\Psi)$ 降到 $O(\Psi/N)$ （ $\Psi$ 为参数量， $N$ 为 GPU 数）。代价是通信量增加约 50%。

动机：DDP 的显存瓶颈

标准 DDP 在每个 GPU 上保留完整的模型副本——参数、梯度、优化器状态全部复制 $N$ 份。无论集群有多少 GPU，单卡必须放得下整个模型的所有状态。

训练显存构成

以混合精度训练(Mixed Precision Training) + Adam 优化器为例，一个 $\Psi$ 参数的模型在每个 GPU 上需要存储：

组成	精度	每参数字节数	总量
模型参数（前向/反向用）	fp16	2	$2\Psi$
梯度	fp16	2	$2\Psi$
主参数副本（优化器更新用）	fp32	4	$4\Psi$
Adam 一阶矩(momentum)	fp32	4	$4\Psi$
Adam 二阶矩(variance)	fp32	4	$4\Psi$
模型状态合计		16	$16\Psi$

一个 7B 参数模型 → $16 \times 7 \times 10^9 = 112$ GB，已超过单张 A100 (80GB) 的容量。这就是 DDP 的显存墙(memory wall)。

上表的 $16\Psi$ 是模型状态(model states) 的开销，不包含激活值(activations)——激活值取决于 batch size 和序列长度，是另一块独立的显存开销，由激活检查点(Activation Checkpointing) 等技术单独优化。

每个 rank 存什么？

这是理解 FSDP 的关键。下表对比 DDP 与 FSDP 中每个 rank 的存储状态：

状态	DDP（每个 rank）	FSDP（每个 rank）	需要通信吗？
模型参数	完整 $2\Psi$	$2\Psi / N$ （只存自己的分片）	✅ 计算前 All-Gather 聚合
梯度	完整 $2\Psi$	$2\Psi / N$	✅ 反向后 Reduce-Scatter 分发
优化器状态	完整 $12\Psi$	$12\Psi / N$	❌ 各 rank 独立更新自己的分片
激活值	完整（自己的数据子集）	完整（自己的数据子集）	❌ 每个 rank 数据不同，无需共享

关键观察：

参数、梯度、优化器状态都切成 $N$ 份，每个 rank 只存 $1/N$ → 单卡显存 $16\Psi / N$
激活值不被 FSDP 切分——每个 rank 处理不同的数据子集，因此各自的激活值本来就不同
优化器状态从不参与通信——每个 rank 只管更新自己那 $1/N$ 的参数，用自己那 $1/N$ 的梯度和优化器状态

通信发生在什么时候？通信的是什么？

FSDP 的核心机制是”用完即弃(discard after use)“：参数在计算前临时通过 All-Gather 聚合为完整形态，计算完成后立即释放非本地的部分。具体的 4 GPU 分片与 All-Gather 示例见 FSDP TransformerBlock 分片示例。

Forward 阶段

对每个 FSDP 单元（通常对应一个 Transformer 层）：

All-Gather 参数：从所有 rank 收集该层的完整参数
前向计算：使用完整参数做本地前向传播（各 rank 数据不同 → 激活值不同）
释放非本地参数：丢弃从其他 rank 收集来的参数分片，只保留自己的 $1/N$

Backward 阶段

从最后一层开始，逐层处理：

All-Gather 参数：再次聚合该层完整参数（前向后已释放，需要重新收集）
反向计算：使用完整参数计算本地梯度
释放非本地参数
Reduce-Scatter 梯度：所有 rank 的梯度求和后切分——每个 rank 只保留自己负责的 $1/N$ 梯度

Optimizer Step

每个 rank 用自己持有的 $1/N$ 梯度 + $1/N$ 优化器状态，更新自己的 $1/N$ 参数。无需任何通信。

通信量分析

操作	时机	每 GPU 通信量
All-Gather 参数	每层前向前	$\sim 2\Psi$ （整个前向累计）
All-Gather 参数	每层反向前	$\sim 2\Psi$ （整个反向累计）
Reduce-Scatter 梯度	每层反向后	$\sim 2\Psi$ （整个反向累计）
FSDP 总计		$\sim 6\Psi$ 字节
DDP 总计（一次 AllReduce）		$\sim 4\Psi$ 字节

FSDP 的通信量比 DDP 多约 50%。多出的 $2\Psi$ 来自前向阶段的参数 All-Gather——DDP 不需要这一步，因为它始终持有完整参数。

详细的单步训练流程见 FSDP 训练步骤逐层图解。

降低通信开销

FSDP 的总通信量 $\sim 6\Psi$ 看起来比 DDP 的 $\sim 4\Psi$ 高不少，但这是理论字节量。实际训练中，通过把通信放到与计算并行的 CUDA stream 上，大部分通信时间可以被隐藏在计算时间内。理想情况下：

$\text{每步 wall-clock} \approx \max(\text{compute 时间},\ \text{comm 时间})$

而不是两者相加。当 compute 时间 $\geq$ comm 时间（大模型通常如此），通信开销几乎为零。FSDP 的这套优化分两类：时间上重叠通信与计算，以及降低每次通信的字节量。

核心优化：Backward Prefetching

反向传播按逆序进行：layer $N \to N-1 \to \ldots \to 1$ 。朴素实现是串行的——等 layer $i$ 的 backward 计算完成 → 启动 layer $i-1$ 的 All-Gather → 等 All-Gather 完成 → 再做 layer $i-1$ 的 backward。通信和计算完全交替，总时间 $\approx$ compute $+$ comm。

BACKWARD_PRE（FSDP2 默认）：在 layer $i$ 的 backward 计算开始之前就异步发起 layer $i-1$ 的 All-Gather。时间轴变成：

compute:  [bwd-N       ][bwd-(N-1)    ][bwd-(N-2)    ]
comm:     [AG-(N-1)][RS-N][AG-(N-2)][RS-(N-1)][AG-(N-3)]
           ↑ 与 bwd-N 计算并行

当 bwd-N 还在 compute stream 上跑时，comm stream 已经在给 layer $N-1$ 做 All-Gather，同时把 layer $N$ 的梯度做 Reduce-Scatter。bwd-N 完成时，下一层的参数已就绪（或接近就绪），bwd-(N-1) 几乎零等待启动。

HuggingFace 的 FSDP1 vs FSDP2 文档明确写了：“FSDP2 uses BACKWARD_PRE by default, as only this allows communication and computation overlap”——FSDP2 干脆移除了另一个选项，因为只有 BACKWARD_PRE 能真正做到重叠。

Forward Prefetching

前向类似：计算 layer $i$ 时预取 layer $i+1$ 的参数。FSDP2 支持两种模式：

隐式预取(implicit prefetching)：执行图里自动插入预取 hook，用户无感知
显式预取(explicit prefetching)：通过 set_modules_to_forward_prefetch() / set_modules_to_backward_prefetch() 手动指定顺序

前向的重叠机会通常比反向少——前向计算往往更快，通信容易”追不上”；而反向计算包含权重梯度 + 激活梯度两部分，时间预算更充裕。

reshard_after_forward：空间换通信

fully_shard(module, reshard_after_forward=...) 是 FSDP2 里最关键的显存-通信权衡旋钮：

取值	前向后行为	显存	通信
`True`（默认）	立即释放非本地参数	低	反向前还要再 All-Gather 一次
`False`	保留完整参数到反向结束	高	省掉反向前的 All-Gather

reshard_after_forward=False 本质上把 FSDP 的 FULL_SHARD 退化为 SHARD_GRAD_OP（ZeRO-2） 的行为——前向后参数仍完整地躺在显存里，反向时直接用，总通信量从 $\sim 6\Psi$ 回到 $\sim 4\Psi$ ，与 DDP 持平。代价是前向后到反向开始之间要多占一份完整参数的显存。模型能放得下时，这是最直接的通信优化。

多机训练：HYBRID_SHARD 的必要性

多机集群的致命瓶颈是节点间带宽远低于节点内：

链路	典型带宽
节点内 NVLink / NVSwitch	300~900 GB/s
节点间 InfiniBand / RoCE	25~50 GB/s

相差一个数量级。如果每层都跨节点 All-Gather，训练会被慢速网络拖垮。HYBRID_SHARD 策略：节点内做 FULL_SHARD（高频 All-Gather / Reduce-Scatter 走 NVLink），节点间做 Replicate（只在反向后做一次梯度 AllReduce，同 DDP）。高频通信被限制在节点内高速链路，跨节点只保留最少同步。

但 HYBRID_SHARD 有一个硬约束：整个模型的状态 ( $16\Psi$ ) 必须能放进一个节点的总显存——因为节点内要做 FULL_SHARD，节点内的所有 GPU 必须能共同存下一份完整模型状态。以 8 × A100 (80GB) = 640 GB 的节点为例，能装下的上限大约是 $640 / 16 \approx 40\text{B}$ 参数（还要预留激活值、临时 buffer、碎片）。所以：

模型 $\leq$ 节点容量（如 7B/13B/34B on 8×A100）：HYBRID_SHARD 是最佳选择
模型 > 节点容量（如 70B+ on 8×A100）：只能跨节点 FULL_SHARD（所有高频通信都过慢网络），或引入 TP / PP 做模型并行。生产环境的 70B+ 训练主流是 TP（节点内）+ FSDP（跨节点）+ 可选 PP 的 2D/3D 并行组合。

HYBRID_SHARD 的详细配置见下方”分片策略”一节。

降低通信字节量

时间重叠之外，还可以直接降低每次通信的字节量：

混合精度通信：bf16 / fp16 参数和梯度的通信量相比 fp32 直接减半。FSDP2 通过 MixedPrecisionPolicy 独立控制前向和 reduce 的精度：
```
from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard

mp = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,   # All-Gather 用 bf16
    reduce_dtype=torch.float32,   # Reduce-Scatter 用 fp32（数值稳定）
)
for layer in model.layers:
    fully_shard(layer, mp_policy=mp)
```
这里的典型做法是 param_dtype=bf16（前向/反向计算和参数通信都用 bf16），reduce_dtype=fp32（梯度 reduce-scatter 用 fp32 避免数值累积误差）。
fp8 All-Gather（FSDP2 新能力）：借助 DTensor 的 tensor subclass 机制，把 All-Gather 的 dtype 降到 fp8，再次减半通信量。主要服务于 Float8 Linear 训练场景。
Tensor fusion：把一个 FSDP 单元内多个小参数的通信合并成一次大的 NCCL 调用，减少 kernel launch 和网络 setup 的固定开销。FSDP1 的 FlatParameter 天然实现了这件事；FSDP2 虽然每个参数是独立 DTensor，但 fully_shard 仍然在单元层面批量调度通信。

实践建议

bf16 通信总是开：最便宜的优化，直接通信量减半
多机训练一定开 HYBRID_SHARD：把慢速网络上的通信降到最低
reshard_after_forward=True 是安全默认：显存最省；只有确信通信是瓶颈且显存有富余时才改 False
BACKWARD_PRE 不用管：FSDP2 已经是默认，而且另一个选项被移除了
用 PyTorch Profiler 看 stream 重叠情况：看 compute stream 是否被 comm stream 卡住；如果 comm 时间吃不完 compute 时间，可以考虑更激进的优化（小 batch 场景尤其要关注）

FSDP vs DDP 对比

维度	DDP	FSDP
单卡显存	$16\Psi$ （完整复制）	$16\Psi / N$ （全部分片）
通信量	$\sim 4\Psi$ （一次 AllReduce）	$\sim 6\Psi$ （多次 All-Gather + Reduce-Scatter）
通信频率	每步一次（反向后）	每层多次（前向 + 反向）
可训练模型上限	受单卡显存限制	受集群总显存限制
计算-通信重叠	反向时 AllReduce 与计算 overlap	All-Gather / Reduce-Scatter 逐层 overlap
数值等价性	—	与 DDP 数值等价

实际案例：7B 参数 + 8 × A100 (80GB)

DDP： $112$ GB/卡 → OOM
FSDP： $112 / 8 = 14$ GB/卡 → 轻松放下

分片策略

PyTorch FSDP 提供多种分片粒度，对应不同的显存-通信权衡。FSDP1 通过 sharding_strategy 枚举配置，FSDP2 改用 DeviceMesh + Placement 表达——两者覆盖的是同一组概念。

FSDP1：sharding_strategy 枚举

策略	对应 ZeRO Stage	分片内容	通信开销	适用场景
`FULL_SHARD`	Stage 3	参数 + 梯度 + 优化器	最高（ $\sim 6\Psi$ ）	模型超出单卡容量
`SHARD_GRAD_OP`	Stage 2	梯度 + 优化器（参数保留完整）	中等（ $\sim 4\Psi$ ）	模型勉强放下，需多省点显存
`HYBRID_SHARD`	节点内 Stage 3	节点内全分片，节点间复制	低跨节点通信	多机集群，利用节点内高带宽
`NO_SHARD`	—（退化为 DDP）	不分片	最低	小模型基线

SHARD_GRAD_OP（ZeRO Stage 2）的一个重要特点：前向阶段不需要 All-Gather 参数（因为参数完整保留），只在反向后做 Reduce-Scatter 梯度，因此通信量与 DDP 持平，同时节省了梯度和优化器的存储。FSDP2 中这个模式通过 reshard_after_forward=False 表达（见上方「降低通信开销 → reshard_after_forward」一节）。

FSDP2：DeviceMesh + Placement

FSDP2 废弃了 sharding_strategy 枚举，改用给 fully_shard(module, mesh=...) 传递不同形状和命名的 DeviceMesh。Shard / HYBRID_SHARD / Replicate 这三种分片模式对应不同的 mesh 配置：

模式	Mesh 形状	DTensor placement	等价 FSDP1 枚举
Shard（全分片）	1D `(N,)`	`(Shard(0),)`	`FULL_SHARD`
HYBRID_SHARD（HSDP）	2D `(nodes, gpus_per_node)`，命名 `("replicate", "shard")`	`(Replicate(), Shard(0))`	`HYBRID_SHARD`
Replicate（退化为 DDP）	2D `(N, 1)`（shard 维度 size=1）	`(Replicate(), Shard(0))`（shard 维度只有 1 个 rank，实际不分片）	`NO_SHARD`

Shard：所有 GPU 做 FSDP

from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard

# 8 GPU 全都参与 FSDP，1D mesh
mesh = init_device_mesh("cuda", (8,), mesh_dim_names=("dp",))

for layer in model.layers:
    fully_shard(layer, mesh=mesh)
fully_shard(model, mesh=mesh)

# 每个参数 → DTensor(placements=(Shard(0),))
# 每个 GPU 持有参数的 1/8

HYBRID_SHARD：节点内分片 + 节点间复制

# 2 节点 × 每节点 8 GPU
# 外层 "replicate"：节点间复制
# 内层 "shard"：节点内分片
mesh = init_device_mesh(
    "cuda",
    mesh_shape=(2, 8),
    mesh_dim_names=("replicate", "shard"),
)

for layer in model.layers:
    fully_shard(layer, mesh=mesh)
fully_shard(model, mesh=mesh)

# 每个参数 → DTensor(placements=(Replicate(), Shard(0)))
# 节点内 8 GPU 分片（1/8 each），两个节点各自持有一份完整的"分片集合"

运行时通信行为：

高频 All-Gather / Reduce-Scatter 只在 shard 维度（节点内） 执行 → 走 NVLink
Reduce-Scatter 完成后，梯度还要在 replicate 维度（节点间）做一次 AllReduce → 走 InfiniBand，但整个 step 只一次
节点内 rank- $i$ 与另一个节点的 rank- $i$ 始终保持参数同步

这就是前面「降低通信开销」节讨论的 HSDP 机制——高频通信限制在节点内高速链路，跨节点只保留最少同步。

Replicate：退化为 DDP

FSDP2 没有独立的 NO_SHARD 枚举。想要纯 DDP 行为最自然的做法是直接用 torch.nn.parallel.DistributedDataParallel。但如果希望在 FSDP2 统一的 API 下表达（例如在一套混合并行代码里临时关掉 FSDP 做 ablation），可以用一个退化的 2D mesh：

# 8 GPU 全部放在 replicate 维度，shard 维度 size=1
mesh = init_device_mesh(
    "cuda",
    mesh_shape=(8, 1),
    mesh_dim_names=("replicate", "shard"),
)

fully_shard(model, mesh=mesh)
# placement 仍是 (Replicate(), Shard(0))
# 但 shard 维度只有 1 个 GPU → 实际不分片
# 行为等价于 DDP：反向时跨 replicate 维度做一次梯度 AllReduce

选型决策

场景	推荐配置
单节点训练	Shard（1D mesh）
多节点 + 模型能放进单节点	HYBRID_SHARD（2D mesh，`("replicate", "shard")`）
多节点 + 模型放不进单节点	Shard 跨所有节点，或 TP + FSDP 的 2D 组合
小模型基线对比	直接用 DDP

DeviceMesh 的真正威力在于可组合性：2D mesh 一个维度命名 "tp"、一个命名 "dp"，然后 mesh["tp"] 传给 parallelize_module、mesh["dp"] 传给 fully_shard——这种 2D/3D 并行组合是 FSDP1 枚举式 API 做不到的。生产环境 70B+ 训练的主流 TP + FSDP 组合就是这样搭起来的。

FSDP1 vs FSDP2

PyTorch 有两代 FSDP 实现。FSDP2 是当前推荐的默认方案：

维度	FSDP1	FSDP2
分片单位	FlatParameter（模块级：把一个模块的所有参数压平成一个大 1D 张量再切分）	DTensor（参数级：每个参数在 dim-0 上独立切分）
通信粒度	粗——计算某一个参数也要 All-Gather 整个模块的 FlatParameter	细——每个参数独立表示为 DTensor，不与同模块其他参数捆绑
冻结参数	差（FlatParameter 要求所有参数 `requires_grad` 一致，无法混合冻结/可训练参数）	好（每个 DTensor 独立持有 `requires_grad` → LoRA/PEFT 开箱即用）
Checkpoint	需 All-Gather 通信才能保存完整参数	直接保存分片 DTensor，无通信
与 TP/PP 组合	困难	原生支持嵌套组合
API	`FullyShardedDataParallel(module)` 包装器	`fully_shard(module)` 就地转换

FSDP2 是当前推荐方案。关于 DTensor 的详细介绍（结构、Placement 类型、DeviceMesh、与 TP 的组合）见其专页。FSDP1 FlatParameter 在 LoRA 场景下的具体问题和 workaround 见 raw/note/fsdp-partial-parameter-freezing.md。