FSDP全分片数据并行实战指南

可视化演示

展示FSDP模型分片：多GPU(G0-G2)间分片Param/Grad/Opt三部分，All-Gather(AG)聚合→Reduce-Scatter(RS)分散，Ckpt检查点。

核心原理

年夜饭的备料环节，老传统“各备全套”最费冰箱：十位厨师各囤十份一样的猪肉和酱油，食材堆到冷库门口。FSDP 的做法像年夜饭的分工备料——整头猪按部位切好，A 师傅保管猪腿、B 师傅保管里脊、C 师傅掌勺调料，谁做菜需要某个部位，就喊一声让对方递过来。这套逻辑对应到深度学习里，就是“全分片”：把模型这套大食材拆开，分给参与训练的每一张 GPU 保管。

在 PyTorch 里，DDP（数据并行）让每张卡都保存一份完整模型副本，十张卡就是十份参数，显存被参数复制白白烧掉。FSDP 则把模型状态——参数、梯度、优化器状态——全部按卡切分。每张卡开始时只持有整个模型的 N 分之一（N 是卡数），需要算某一层时，再用 all-gather 操作把分片临时拼回全量；算完梯度后，又用 reduce-scatter 把梯度切分汇总回各卡。这就是“全分片”的含义：模型侧的三类状态统统均摊。

显存收益可以用一个粗略的账本算清楚。假设模型参数占 P 字节，梯度占 G 字节，优化器状态占 O 字节，激活值占 A 字节，N 张卡全分片后，单卡模型侧占用约为：

$$\displaystyle M_{\text{model}} = \frac{P + G + O}{N} + A$$

1. P（参数占用）：模型权重本身占的字节数，等于参数量乘以每个参数的字节数。

$$\displaystyle P = \text{num\_params} \times \text{bytes\_per\_param}$$

2. G（梯度占用）：反向传播算出的梯度，尺寸与参数一样大，可近似认为 G ≈ P。

$$\displaystyle G \approx P$$

3. O（优化器状态）：Adam 等优化器额外维护的一阶、二阶动量，通常约为参数的 2 倍。

$$\displaystyle O \approx 2P$$

4. 除以 N（分片摊薄）：全分片把参数、梯度、优化器状态三样均分到 N 张卡上。

$$\displaystyle \frac{P + G + O}{N}$$

5. 加 A（激活值）：前向传播的中间张量要在反向时复用，无法被分片，所以单独相加。

$$\displaystyle M_{\text{model}} = \frac{P + G + O}{N} + A$$

当然，天下没有免费的午餐。分片省下了冰箱空间，却多出了“借料”的跑腿成本：每次前向和反向都要把参数拼回全量，通信量比 DDP 更高。所以 FSDP 的定位很清晰——它不是让训练更快，而是让放不下的模型放得下。理解了“省显存、费通信”这笔账，后面所有坑的答案就藏在这八个字里。

模型包装

千人宴的后厨不能只有一个大通铺，否则所有人挤在一起抢料、找不到东西。聪明的做法是把后厨隔成若干小档口，每个档口独立保管自己的食材、独立记账。FSDP 里的“包装”正是这个隔断动作：你决定把模型的哪些层圈进同一个“档口”（FSDP 块），块内的参数一起分片、一起拼回。

PyTorch 中，最直接的包装入口是 FullyShardedDataParallel(model, ...)，它把整个模型当作一个巨大的档口。但整包会让拼参时的峰值显存等于整个模型的参数总量，分片效果大打折扣。所以工程上更常用“自动包装”：通过 auto_wrap_policy 告诉 FSDP 按什么规则自动给每一层切档口。比如 size_based_auto_wrap_policy(min_num_params=阈值) 表示“参数量超过阈值的模块自动独立成块”；Transformer 架构则常用 transformer_auto_wrap_policy 按层类型切块。你也可以手动逐个包装子模块，精确控制每一刀切在哪里。

包装粒度直接决定显存峰值。设所有参数均摊后每卡占 P/N，而当前这一个档口的完整参数大小为 P_block，则前向拼参瞬间的单卡峰值约为：

$$\displaystyle M_{\text{peak}} = \frac{P}{N} + P_{\text{block}} + A$$

1. P/N（常驻分片）：所有参数均分后，每卡长期占用的那部分。

$$\displaystyle \frac{P}{N} = \text{各卡常驻参数}$$

2. P_block（块内全量）：单个 FSDP 包装块在拼参时需要的完整参数。

$$\displaystyle P_{\text{block}} = \max(\text{各包装块参数量})$$

3. 加 A（激活值）：激活值显存仍然存在，且与包装粒度关系不大。

$$\displaystyle M_{\text{peak}} = \frac{P}{N} + P_{\text{block}} + A$$

档口划分是一门平衡术：切得太粗，拼参峰值高，照样 OOM；切得太细，每次前向都要多次跨卡借料，通信时间直线上升。实践中的经验是：先按模型结构切（Transformer 的每个 block 一层），再用阈值策略微调，让每个档口大小适中。记住，包装决定的不是“能不能省”，而是“省多少、跑多快”。

分片策略

年夜饭的分工不止一种：最激进的方案是连菜带锅全摊开，每位厨师只保管自己那几筐料；温和一点的方案是大家各自备齐所有主菜食材，只把价格昂贵的调料集中保管；还有一种最省事的方案，干脆各备各的全套，互不借用。FSDP 的三种分片策略，恰好对应这三种厨房分工。

• FULL_SHARD：参数、梯度、优化器状态全部按卡切分，最省显存，但每次前向反向都要 all-gather 拼回参数，通信最重。

• SHARD_GRAD_OP：参数每卡保留全套（不切），只有梯度分片。省掉了梯度这块显存，通信量比 FULL_SHARD 低。

• NO_SHARD：什么都不切，等价于 DDP，每卡一份完整模型，仅梯度做 all-reduce。

三种策略的通信代价可以量化对比。按每步训练计算，FSDP 全分片需要“前向拼参数一趟、反向拼参数一趟、梯度切分一趟”，共三趟模型大小的通信；而 DDP 只需要梯度 all-reduce 一趟（收和发合计两趟）：

$$\displaystyle C_{\text{FSDP}} = 3P, \qquad C_{\text{DDP}} = 2P$$

1. 前向 all-gather（C₁）：前向计算前，把分片的参数从各卡拼回全量，通信量为 P。

$$\displaystyle C_1 = P$$

2. 反向 all-gather（C₂）：反向算梯度时还要再用一次参数，所以再拼一趟，通信量为 P。

$$\displaystyle C_2 = P$$

3. reduce-scatter（C₃）：梯度算完后，按分片规则切分并汇总回各卡，通信量为 P。

$$\displaystyle C_3 = P$$

4. DDP 的 all-reduce（C_DDP）：梯度一次性全部规约，收与发合计约 2P。

$$\displaystyle C_{\text{DDP}} = 2P$$

策略选择的决策树其实很朴素：显存快爆了，优先 FULL_SHARD；显存刚好够，想要更快收敛，试试 SHARD_GRAD_OP；模型不大、卡又多，NO_SHARD（DDP）反而省心。厨房里没有万能分工，只有“按食材多少、传菜通道宽窄”来定的方案——显存压力大就多摊一些，带宽紧张就少借几次料。

检查点管理

年夜饭从下午一直炖到深夜，中途最怕断电、锅翻、临时改菜单。聪明的厨师会每隔一小时记一次账：食材还剩多少、哪道菜炖到几成熟、火候怎么调的。这套“中途记账”在 FSDP 里就是检查点（checkpoint）：把模型训练到一半的状态存下来，下次断电了从最近一笔账接着做，而不是从头再来。

但 FSDP 的账本有个特殊之处：每张卡只记得自己那一片食材的状态。直接 torch.save(model.state_dict(), "model.pt")，保存的只是本卡的分片，不是完整模型。正确做法是让 rank0 卡把所有分片聚合起来，拼成一份完整的“总账本”再落盘。保存时机的选择也有讲究：每隔固定步数存一次、或按 epoch 存一次；保存的频率越高，断电损失越小，但磁盘写入和聚合通信的开销越大。

聚合保存的本质，就是把 N 个分片状态按原顺序拼回完整状态。记第 i 张卡的状态为 S_i，则：

$$\displaystyle S_{\text{full}} = \sum_{i=1}^{N} S_i$$

1. S_i（分片状态）：第 i 张卡保存的状态，包含它保管的那部分参数、梯度和优化器状态。

$$\displaystyle S_i = \text{卡 } i \text{ 的分片状态}$$

2. 求和符号 Σ（拼接）：从第 1 张卡加到第 N 张卡，按分片顺序依次拼接。

$$\displaystyle \sum_{i=1}^{N}$$

3. S_full（完整状态）：拼接结果等价于单卡全量模型的 state_dict，可直接用于单卡推理或后续微调。

$$\displaystyle S_{\text{full}} = \sum_{i=1}^{N} S_i$$

版本兼容和恢复技巧同样重要：PyTorch 升级后，FSDP 的分片存储格式可能变化，加载旧检查点偶尔会报 unexpected key；跨版本训练时，建议只保存和加载模型权重，优化器状态则尽量同版本。恢复训练时，务必用与保存时完全相同的包装策略和分片策略重建模型，再 load_state_dict，否则等于拿着别人的钥匙开自己的锁。

性能调优

后厨出菜慢的瓶颈往往不在炒菜本事，而在传菜路上：每炒好一盘就叫一个人跑一趟，十分钟的菜全耗在走廊里。聪明的厨师会先把炒好的菜攒在传菜口，凑够一车再一起送。FSDP 里的“攒菜”就是梯度累积：本来每算一个 batch 就要同步一次梯度，现在连续算 K 个小 batch，梯度先攒在本地，攒够了再统一同步一次。

梯度累积直接降低了通信频率。设单卡单步的 batch 为 B，累积 K 步后再做一次优化器更新，等效 batch 变为：

$$\displaystyle B_{\text{eff}} = B \times K$$

1. B（单步 batch）：每张卡每个小步实际喂给模型的样本数。

$$\displaystyle B = \text{单卡单步 batch}$$

2. K（累积步数）：连续计算 K 个小 batch 的梯度，累加但不更新参数。

$$\displaystyle K = \text{累积步数}$$

3. 乘积（等效规模）：累积 K 步后做一次更新，等效于一次喂入 K 倍样本。

$$\displaystyle B_{\text{eff}} = B \times K$$

除了梯度累积，还有几招能进一步压通信：一是通信与计算重叠，FSDP 在算当前层时，后台提前向其他卡请求下一层的参数，边炒边传；二是混合精度，用 BF16 传梯度，通信量直接减半；三是调整超参——打包粒度、分片策略、梯度累积步数之间互相牵制，往往要组合调优。调优的顺序建议是：先看显存够不够（不够就切更细或换 FULL_SHARD），再看通信占比高不高（高就累积梯度、开重叠、上混合精度），最后才动学习率这些收敛参数。

报错排查

再熟练的后厨也会遇到事故：食材不够了、传菜超时了、调料过期了。FSDP 训练同样有一批“高频事故”，先分清是哪一类，再对症下药。

第一类是显存爆掉（OOM）。它通常不是玄学，而是账没算清：包装粒度太粗导致前向拼参峰值过高，或 batch 太大导致激活值爆表。排查时先按公式估算 M_peak = P/N + P_block + A，确认是哪一项超了；然后对症：把档口切细、减小 batch、开混合精度、或改用 FULL_SHARD。第二类是通信超时。NCCL 的集合通信如果等待太久，会抛超时异常。这往往是机器间带宽不足、某张卡掉队拖慢整体、或 timeout 参数设得太小。解决思路是：检查网络拓扑、用 torch.distributed.barrier() 定位掉队卡、调大 init_process_group(timeout=...)。第三类是版本与加载问题：PyTorch 升级后分片格式变化、保存和加载时包装配置不一致，都会报 unexpected key 或加载失败。

通信超时的判定逻辑可以写成一条不等式：

$$\displaystyle T_{\text{comm}} > T_{\text{timeout}} \Rightarrow \text{Failure}$$

1. T_comm（通信耗时）：一次集合通信从发起到完成的总用时，受带宽、卡数、数据量影响。

$$\displaystyle T_{\text{comm}} = \text{单次集合通信耗时}$$

2. T_timeout（超时阈值）：初始化进程组时通过 timeout 参数设置的上限。

$$\displaystyle T_{\text{timeout}} = \text{timeout 设置值}$$

3. 大于号右移（触发失败）：等待时间一旦越过阈值，进程组主动报错，避免无限挂起。

$$\displaystyle T_{\text{comm}} > T_{\text{timeout}} \Rightarrow \text{Failure}$$

排查的通用心法是“先环境、后代码”：先看卡数、带宽、PyTorch 和 torch.distributed 版本，再看包装与分片配置，最后才怀疑模型本身。厨房里出了事，先查是不是停水停电，别急着怪菜谱——大多数 FSDP 报错，都是配置和环境问题，不是算法问题。

AI 应用场景

Python 演示

import osimport torchimport torch.nn as nnimport torch.distributed as distfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDPfrom torch.distributed.fsdp import ShardingStrategyfrom torch.distributed.fsdp.wrap import size_based_auto_wrap_policy  def init_dist():    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))    dist.init_process_group(backend="nccl")  class TinyModel(nn.Module):    def __init__(self):        super().__init__()        self.fc1 = nn.Linear(512, 512)        self.fc2 = nn.Linear(512, 512)        self.fc3 = nn.Linear(512, 512)     def forward(self, x):        x = torch.relu(self.fc1(x))        x = torch.relu(self.fc2(x))        return self.fc3(x)  def main():    init_dist()    rank = dist.get_rank()    model = TinyModel()    wrap_policy = size_based_auto_wrap_policy(min_num_params=2048)    model = FSDP(        model,        auto_wrap_policy=wrap_policy,        sharding_strategy=ShardingStrategy.FULL_SHARD,    )    opt = torch.optim.Adam(model.parameters(), lr=2e-4)    accum = 4    micro_bs = 64    data = torch.randn(micro_bs, 512)     for epoch in range(3):        opt.zero_grad()        for step in range(accum):            out = model(data.to(rank))            loss = out.sum() / accum            loss.backward()        opt.step()        if rank == 0:            print("epoch", epoch, "loss", loss.item())     if rank == 0:        torch.save(model.state_dict(), "fsdp_model.pt")    dist.barrier()    dist.destroy_process_group()  if __name__ == "__main__":    main()

1. torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))每个进程绑定自己那张卡，避免多进程抢卡。
2. dist.init_process_group(backend="nccl")建立多卡通信组，NCCL 负责 GPU 间的集合通信。
3. size_based_auto_wrap_policy(min_num_params=2048)参数量超过阈值的模块自动切成独立 FSDP 块。
4. FSDP(model, auto_wrap_policy=..., sharding_strategy=FULL_SHARD)包装模型并启用全分片。
5. loss = out.sum() / accum把损失按累积步数缩放，等效于扩大 batch 又不炸显存。
6. data.to(rank)直接把 rank 当作设备号，把数据送到本进程对应的卡上。
7. torch.save(model.state_dict(), "fsdp_model.pt")在 rank0 聚合各卡分片后保存完整检查点。

常见误区

面试题