FSDP1 vs FSDP2
System · PyTorch · commit 2ee7a63

FSDP1 vs FSDP2:同一套 1/N,两种参数表示

两者都把参数、梯度、优化器状态切到每张卡只留 1/N,前向前 all-gather,反向后 reduce-scatter。 真正的分水岭是切的单位:FSDP1 把一个 module 的参数压成一条 FlatParameter 再切,FSDP2 让每个参数自己成为 DTensor(Shard(0))。 这个选择决定了 padding、state dict、优化器看到什么,以及能不能和 TP/EP 直接拼。

PyTorch Distributed · FSDP FSDP 文档 ↗ FSDP2 API ↗ FSDP2 RFC ↗
TL;DR · 30 秒抓住核心

常驻显存的账,两者一样。FULL_SHARD 下每卡参数、梯度、优化器状态都约为 1/N;7B bf16 + AdamW、N=8 时是 10.5 GB,而完整副本是 84 GB。

FSDP1 的代价在「打包」。它把 module 内所有参数 flatten 成一条 1D 的 FlatParameter:use_orig_params=True 时每个原始参数前补 padding 到 16 B 边界(bf16 对齐单位是 8 个元素,最多补 7 个),同一路径下整条再补到 N 的倍数。切出来的 chunk 边界可能落在某个参数中间。

FSDP2 换成分片单位 = 单个参数。每个参数单独 chunk 成 DTensor,优化器直接拿到 DTensor;sharded state dict 与训练表示一致,不需要为了保存而 all-gather。但「按参数分片」不等于「按参数通信」——一个 fully_shard group 仍然一次 all-gather、一次 reduce-scatter。

1/N
FULL_SHARD 下每卡常驻参数
梯度与优化器状态同为 1/N;N=8 即 12.5%
16 B
FSDP1 地址对齐粒度
bf16 每参数最多补 7 个元素
1 个参数
FSDP2 的分片单位
通信仍按 module group 合并
01

先把「每卡一份」变成「每卡 1/N 份」——难点不在公式,在表示

读完这节你应该知道:为什么 1/N 只是起点,真正影响工程的是这 1/N 在内存里怎么摆。

DDP 的账本很直观:每张卡都放一份完整的参数、梯度、优化器状态。用 P 表示参数量,bp、bg 表示参数和梯度的字节数,bopt 表示每个参数的优化器状态字节数,每卡常驻就是:

MDDP = P × (bp + bg + bopt)
完整副本;不含激活、临时缓冲与碎片

按 7B 参数、bf16 参数与梯度、AdamW 的 fp32 一阶/二阶动量(8 B/param)算,完整副本是 84 GB。这已经超过一张 80 GB 卡的可分配显存,而且还没算激活。FSDP 的思路是把这三类状态都切成 N 份,每卡只留一份:

MFSDP = P × (bp + bg + bopt) / N
FULL_SHARD 的常驻账;7B、N=8、bf16 时为 10.5 GB

这个 1/N 的结论对 FSDP1 和 FSDP2 都成立。真正的问题是:这 1/N 在内存里长什么样?是一条被 flatten 的连续 buffer,还是许多个独立的参数分片?这个表示方式会连带影响三件事:

下面的 02 先看前人怎么做,03 再把两种表示拆开。

02

从 DDP 到 ZeRO 再到 FSDP:分片对象越来越多,表示问题越来越尖

这一节建立对照:DDP 解决通信去中心化,ZeRO/FSDP 解决显存,但「怎么表示一份分片」是 FSDP 才真正暴露的问题。

DDP

只分数据,不分状态

每张卡维护完整模型,反向后用 ring all-reduce 同步梯度。实现简单、通信与计算重叠好,但显存随参数量线性增长。

卡在哪:模型 + 优化器状态超过单卡时,DDP 只能换更大的卡,不能靠加卡解决。

ZeRO / FSDP

把状态切开,计算前再拼回

ZeRO-1 切优化器状态,ZeRO-2 再切梯度,ZeRO-3 连参数一起切。FSDP 的 FULL_SHARD 对应最后一档:前向前 all-gather 参数,反向后 reduce-scatter 梯度。

代价:每步多出两次 all-gather(前向、反向各一次)和一次 reduce-scatter;通信量随 N 增大而趋近 3P 个元素。

把 ZeRO 的三个阶段和每卡常驻写成一张表,差别一目了然(b 为参数/梯度字节数,优化器状态按 bopt 计):

表 1 · 分片对象与每卡常驻(不含激活;ZeRO 阶段名为 DeepSpeed 术语)
方案参数梯度优化器状态每卡常驻
DDP完整完整完整P × (2b + bopt)
ZeRO-1完整完整1/N2Pb + P bopt/N
ZeRO-2完整1/N1/NPb + P(b + bopt)/N
ZeRO-3 / FSDP FULL_SHARD1/N1/N1/NP(2b + bopt)/N

这张表只说明「切哪些状态」,没有说明「切出来的东西长什么样」。FSDP1 和 FSDP2 的分歧就在后者。

设计动机

FSDP2 RFC 对 FlatParameter 的原话

FSDP2 的设计文档(PyTorch issue #114299)把旧实现的限制归结为 “mainly from the flat parameter construct”,目标是 “simplicity and composability in the new design”。它列出的四个直接收益是:fp8 与非 fp8 参数可以混在同一次 all-gather;frozen 与 non-frozen 参数可以混在同一通信组而不额外占显存;sharded state dict 与训练表示一致;以及给 torch.compile 留出重排通信组的空间。

同一份 RFC 还把新老组件一一对应:FSDPParamGroup 对应 FlatParamHandle,FSDPParam 大致对应 FlatParameter,fully_shard + FSDPState 对应 FullyShardedDataParallel。这不是一次小修,而是把「一个 flat buffer」拆成「每个参数一个对象」。

FSDP1 把这些策略暴露成 ShardingStrategy:FULL_SHARD、SHARD_GRAD_OP、NO_SHARD、HYBRID_SHARD。FSDP2 没有这个枚举,只用 reshard_after_forward 控制「前向后是否释放完整参数」(_fully_shard.py:102-120);该参数默认 None,非根模块解析为 True、根模块解析为 False(:201-202)。FSDP2 没有命名 NO_SHARD 模式:复制场景直接用 DDP,或用 2D mesh 的 shard 维 = 1 模拟复制(PyTorch issue #152710,非正式用法)。FSDP1 的 NO_SHARD 已弃用并建议 DDP(api.py:50-55)。

03

FSDP1 先打包再切,FSDP2 每个参数自己切

两种表示都能做到 1/N;差别在 padding、state dict,以及优化器看到的参数类型。

FSDP1:一条 FlatParameter,切完再分发

FSDP1 的 FlatParameter 是一个 1D 的 nn.Parameter,由 module 内所有原始参数 flatten 后拼接而成(torch/distributed/fsdp/_flat_param.py:202)。拼接不是简单 concat:

L = Σ si + Σ padalign,i + padworld  (use_orig_params=True)
padalign,i = (−Ti−1) mod a,   Ti−1 = Σj<i(sj + padalign,j)
a = 16 / bytes(dtype),   padworld = (−Tlast) mod N
shardr = ceil(L / N)
FSDP1 的 flat 长度与每卡有效 chunk;对齐/world padding 只在 use_orig_params=True 时存在(默认 False 时 L = Σsi,无 padding;torch.chunk 返回不等长 chunk 时补到 chunks[0].numel())

默认 use_orig_params=False(fully_sharded_data_parallel.py:412),优化器看到的是整条 FlatParameter;named_parameters() 也会暴露它。官方文档明确写了这一点,并指出 use_orig_params=True 是使用 torch.compile() 的必要条件。

FSDP2:每个参数一个 DTensor(Shard(0))

FSDP2 的入口是 torch.distributed.fsdp.fully_shard(module, *, mesh=None, reshard_after_forward=None, shard_placement_fn=None, mp_policy=..., offload_policy=..., ignored_params=None, dp_mesh_dims=None)(torch/distributed/fsdp/_fully_shard/_fully_shard.py:102)。torch/distributed/_composable/fsdp/fully_shard.py:3-8 只是向后兼容的 re-export;真正的 composable 实现在 _fully_shard/ 目录。

它把 module.parameters() 原地从普通 Tensor 换成 DTensor,默认每个参数按 dim-0 分片(_fsdp_param.py:307)。分片过程是:

shardi = ceil(si / N)
padi = shardi × N − si
FSDP2 的单参数分片;_chunk_with_empty 保证空参数也能拿到一个空 chunk
官方文档

分片单位的一句话对照

对 FSDP1,文档写的是 “flattens, concatenates, and chunks a group of tensors together”;对 FSDP2,写的是 “chunks each parameter on dim-0 across the data parallel workers (using torch.chunk(dim=0))”。

这两句只描述分片,不描述通信。两者不是一回事:FSDP2 仍然把 group 内所有参数的分片拼在一次 collective 里,下一张卡解释这一点。

最容易混的一点

分片粒度 ≠ 通信粒度

FSDP2 按参数分片,但不是每个参数发一次 collective。fully_shard(module) 会为这个 module 建一个 group;group 内所有参数在一次 all-gather 里拼起来,梯度在一次 reduce-scatter 里拆开(_fully_shard.py:145-152,_fsdp_collectives.py:341、:561)。分片单位是参数,通信单位仍是 group。

这也解释了为什么文档强调要 bottom-up、按层调用 fully_shard:group 的边界就是你给的 module 边界,直接决定 all-gather 的粒度。

自己算:每卡到底放多少

下面的账本把参数、梯度、优化器状态分开算。拖动参数量、世界规模、精度,或切换分片策略;所有数字都按同一组公式实时重算。

参数
—
梯度
—
优化器状态
—
常驻合计
—
每卡账本
—
通信量 / iter / 卡
—

口径:优化器状态固定按 AdamW fp32 的 8 B/param; 参数和梯度按所选精度(bf16 2 B、fp32 4 B);1 GB = 109 B;不含激活、通信缓冲和碎片。 「最坏峰值」= 峰值时刻的常驻 + 一份完整参数(FULL_SHARD 下这份副本只在 all-gather 期间存在,SHARD_GRAD_OP 下保留到反向结束;NO_SHARD 没有额外副本,峰值就是常驻),不含反向中的完整梯度缓冲;按层 wrapping 时峰值只多当前层的完整参数。

手算对账:两组参数与滑块同口径,逐位可复算

组 A:P=7B,N=8,bf16,FULL_SHARD。

  • 参数分片 = 7×2/8 = 1.75 GB
  • 梯度分片 = 7×2/8 = 1.75 GB
  • 优化器分片 = 7×8/8 = 7.00 GB
  • 常驻合计 = 1.75 + 1.75 + 7.00 = 10.50 GB
  • 最坏峰值(含一份完整参数)= 10.50 + 7×2 = 24.50 GB(整模型一个 group 时,all-gather 再放一份 14.00 GB 的完整参数;不含反向中的完整梯度缓冲)
  • 通信 = 3×7×(8−1)/8×2 = 36.75 GB(2 次 all-gather + 1 次 reduce-scatter)

组 B:P=13B,N=4,fp32,SHARD_GRAD_OP。

  • 全量副本 = 13×(2×4+8) = 208.00 GB;分片后常驻 52 GB 可放进 4 张卡;峰值 104 GB 仍需按层 wrapping
  • 参数常驻分片 = 13×4/4 = 13.00 GB;计算时展开为完整 52.00 GB
  • 梯度分片 = 13×4/4 = 13.00 GB
  • 优化器分片 = 13×8/4 = 26.00 GB
  • 常驻合计 = 13.00 + 13.00 + 26.00 = 52.00 GB
  • 最坏峰值(含一份完整参数)= 52.00 + 13×4 = 104.00 GB(完整参数保留到反向结束;不含反向中的完整梯度缓冲)
  • 通信 = 2×13×(4−1)/4×4 = 78.00 GB(1 次 all-gather + 1 次 reduce-scatter)

把滑块拉到同样参数应显示同样的数;若不一致,以脚本输出为准并在报告里记下差异。

这组账本说明什么:FULL_SHARD 和 SHARD_GRAD_OP 的常驻分片相同,区别在通信——前者反向还要再 all-gather 一次。SHARD_GRAD_OP 在前向后不释放完整参数,省掉这次 all-gather,但高峰持续时间更长。NO_SHARD 完全不省显存,官方已把它标为 deprecated,建议直接用 DDP。

04

通信量由策略决定,不是由 FSDP1/FSDP2 决定

读完这节你应该知道:FULL_SHARD、SHARD_GRAD_OP、NO_SHARD 各自的通信账,以及 FSDP2 在工程上换掉了什么。

先看 7B、N=8、bf16 的三种策略。常驻和通信都按上一节的公式算,GB 为十进制:

表 2 · 7B 参数、N=8、bf16、AdamW 8 B/param;最坏峰值按整模型一个 group,计一份完整参数副本(FULL_SHARD 在 all-gather 期,SHARD_GRAD_OP 保留到反向结束),不含反向中的完整梯度缓冲(不含激活;1 GB = 109 B)
策略每卡常驻最坏峰值(含完整参数)通信 / iter / 卡FSDP2 等价
FULL_SHARD10.50 GB24.50 GB36.75 GBreshard_after_forward=True(需显式;默认 None)
SHARD_GRAD_OP10.50 GB24.50 GB24.50 GBreshard_after_forward=False(需显式;默认 None,近似)
NO_SHARD84.00 GB84.00 GB24.50 GB无命名模式;复制场景用 DDP

FULL_SHARD 每卡每步的通信量是 3P(N−1)/N 个元素:前向 all-gather、反向 all-gather、反向 reduce-scatter。N 越大越接近 3P。SHARD_GRAD_OP 和 NO_SHARD 都是 2P(N−1)/N:前者是 1 次 all-gather + 1 次 reduce-scatter,后者是 1 次 all-reduce。

注意 SHARD_GRAD_OP 的「省」不在峰值数字上,而在高峰的持续时间:它前向后不 reshard,完整参数一直留到反向结束;FULL_SHARD 在前向后就把完整参数释放,反向需要时再 all-gather。对激活很大、层数很多的模型,这个区别会直接决定能不能放下。

FSDP1 的运行时:三条 stream 和一个 rate limiter

FSDP1 在 _runtime_utils.py 里把 unshard、post-backward、pre-unshard 放在不同 CUDA stream 上(_runtime_utils.py:240-270),用 limit_all_gathers 限制同时在飞的 all-gather 数量(:297-308)。前向的 pre-hook 调 _unshard;梯度在 _post_backward_hook(:701)调用的 _reduce_grad 里走 dist.reduce_scatter_single(:766、:848-858),_post_backward_reshard 只负责 reshard(:793-799)。

FSDP2 的运行时:显式 stream,不做 record_stream

FSDP2 把通信状态收进 FSDPParamGroup:all-gather 的 copy-in、all-gather、reduce-scatter 各有自己的 stream(_fsdp_param_group.py:93-120)。官方文档明确说它采用一套避免 torch.Tensor.record_stream 的内存管理方式,因此不需要 limit_all_gathers=True;也没有 bucket_cap_mb,group 的边界完全由你调用 fully_shard 的 module 决定。

混合精度也换了 API:FSDP1 用 MixedPrecision,FSDP2 用 MixedPrecisionPolicy(_fsdp_api.py:15)。语义是「all-gather 时把参数 cast 到 param_dtype,优化器仍按原始 dtype 更新」;reduce_dtype 控制梯度归约精度。这不是 autocast 的替代品,cast 发生在 all-gather 这个边界上。

和 TP / PP 组合:FSDP2 的分片是「可叠加」的

FSDP 本身是数据并行维度的分片;TP 是层内维度,PP 是层间维度。组合时,FSDP 的 shard group 只应覆盖同一份模型副本——对 PP 来说就是同一个 stage 的 DP 副本,不能跨 stage all-gather 不存在的参数 推断(由 FSDP 的 all-gather 语义与 PP 的分 stage 所有权推出)。

二维 mesh 的构造入口是 torch.distributed.device_mesh.init_device_mesh(torch/distributed/device_mesh.py:1640),DeviceMesh 类本身在 :188。传给 fully_shard 的 DP mesh 是 1D 或 2D:1D 是纯 FSDP,2D 是 HSDP(_fully_shard.py:185-190);和 TP/EP 组合后,内部 SPMD mesh 可以是 2–4 维(_fsdp_param.py:740-790)。

05

一次迭代的六步:shard → all-gather → forward → reshard → all-gather → reduce-scatter

这张图把两种表示放在同一张时间线上:左边是 FSDP1 的 flat chunk,右边是 FSDP2 的 per-parameter DTensor。

FSDP1 FlatParameter 与 FSDP2 per-parameter DTensor 的分片对比 上半部分:FSDP1 把四个参数 flatten 成一条带对齐 padding 的 FlatParameter,再按 N=4 切成四个 chunk;chunk 边界可能切断参数。下半部分:FSDP2 让每个参数单独按 N=4 切成 DTensor 分片,rank r 持有每个参数的第 r 块;collective 仍按 group 合并。 FSDP1 · FlatParameter 先把 module 内所有参数 flatten 成 1D,补对齐与世界规模 padding(仅 use_orig_params=True 路径),再按 N=4 切 原始参数 P1 P2 P3 P4 FlatParameter world pad 按 N=4 切 chunk 0 chunk 1 chunk 2 chunk 3 每卡持有 rank 0 rank 1 rank 2 rank 3 chunk 边界切断 P2 / P4 all-gather:4 个 chunk → 一条完整 flat reduce-scatter → 4 chunk FSDP2 · per-parameter DTensor 每个参数单独 chunk,rank r 持有每个参数的第 r 块;collective 仍按 group 合并 P1 pad P2 pad P3 pad P4 pad chunk 0 chunk 1 chunk 2 chunk 3 pad rank 视图 rank 0 rank 1 rank 2 rank 3 每个彩色块 = 一个 DTensor(Shard(0));rank r 持有每参数第 r 块,共 4 个分片。 通信仍按 group 合并:一次 all-gather 拼回 4 个完整参数,reduce-scatter 反向拆分。 all-gather / reduce-scatter:group 级,不是 per-parameter 级
图 1 · FSDP1 的 flat chunk 与 FSDP2 的 per-parameter DTensor 对比(示意图,不按真实比例)。 FSDP1 的 padding 与切分见 _flat_param.py:565-567、:759-763、:1104; FSDP2 的单参数 chunk 与 DTensor 包装见 _fsdp_param.py:362-369、:394-398。
读图:为什么 FSDP1 的 chunk 会切断参数以及为什么 FSDP2 的通信仍然按 group 走

FSDP1:切的是整条 FlatParameter,不是单个参数。chunk 的边界由 L/N 决定,和原始参数的边界无关。图中 rank 0 的 chunk 同时包含 P1 的尾部、对齐 padding 和 P2 的开头;rank 2 的 chunk 同时包含 P3、padding 和 P4 的开头。要把某个参数取出来,必须先用 flat metadata(_numels_with_padding、_shapes、_fqns)把 flat view 拆回原始形状。

FSDP2:每个参数单独 chunk,所以 rank r 的存储是「每个参数的第 r 块」这一组独立分片。它们不是一条连续 buffer;每个分片自己是一个 DTensor。但 collective 不是按参数发的:FSDPParamGroup 会把 group 内所有参数的 all-gather input 拼在一起,一次通信拿回整个 group 的完整参数(_fsdp_collectives.py:341-380)。所以「分片单位 = 参数」和「通信单位 = group」同时成立。

图中 FSDP2 每行末尾的斜线 padding 表示该参数自己的 ceil(s_i/N) 余量;它不会跨参数共享。FSDP1 的斜线 padding 则是整条 flat 的对齐/world-size 余量。两者的 padding 归属不同,这是下一节显存对比的根源。

把六步按时间排开:

  1. 初始化:FSDP1 构造 FlatParameter 并按 N 切;FSDP2 把每个参数换成 DTensor(Shard(0))。
  2. 前向前:all-gather 出完整参数。FSDP1 是一条 flat buffer;FSDP2 是 group 内所有参数分片的拼接。
  3. 前向:用完整参数计算。FSDP1 的 use_orig_params=True 会提供原始参数的 view;FSDP2 的普通参数是 all-gather buffer 的 view,DTensor 参数(TP/EP)保持自己的 placement。
  4. 前向后:如果 reshard_after_forward=True(FSDP2 需显式;默认 None 按非根 True、根 False 解析),释放完整参数,只留分片;False 则保留到反向结束。
  5. 反向前:如果前向后释放了,需要再次 all-gather;这就是 FULL_SHARD 的第二次 all-gather。
  6. 反向后:reduce-scatter 梯度,释放完整参数,优化器在分片上更新。
06

省下的是 1/N,不是 1/N²——真正的差别在状态字典和组合

这一节回答两个容易混的问题:per-parameter DTensor 到底省不省显存?FSDP2 换来的工程收益是什么?

先算一笔「分片粒度」的账

取两个很小的参数,大小分别是 3 和 5 个元素,N=4,bf16。这个例子故意让参数不能被 N 整除,好把 padding 的归属看清楚:

表 3 · 3 + 5 个元素、N=4、bf16 时每卡持有的元素数(FSDP1 的 a=8 仅适用于 use_orig_params=True 行;FSDP2 的每参数分片大小(含 padding)是 ceil(si/N);按 P1=3、P2=5 的顺序)
实现计算每卡元素数padding 归属
FSDP1 默认(use_orig_params=False)ceil((3+5)/4)2无 padding(8 能被 4 整除)
FSDP1(use_orig_params=True)P2 前对齐 pad 5,再 world pad 3:(3+5+5+3)/44P2 前对齐 + 整条 world-size 余量
FSDP2ceil(3/4)+ceil(5/4) = 1+23每个参数自己的 chunk 余量

结论很反直觉:FSDP2 的 per-parameter 分片不是无条件的显存胜利。在这个例子里,它比 FSDP1 默认多 50%(相对 2 个元素/卡),比 use_orig_params=True 少 25%(相对 4 个元素/卡)。原因很简单:FSDP1 把所有参数合成一条再切,边界可以互相借位;FSDP2 每个参数都要各自向上取整,小参数多的时候反而更容易浪费。

这也解释了为什么「FSDP2 更省显存」不能一概而论。两者的常驻比例都是 1/N;真正的差异是 padding 的归属、临时 buffer 的管理方式,以及 state dict 能不能直接写分片。对 LLM 这种参数大、形状通常能被 N 整除的模型,padding 差异很小;对大量小参数(LayerNorm、bias、router)的模型,FSDP1 的 flat 打包反而可能更省 推断(由表 3 的取整公式外推)。

状态字典:一个要拼回,一个直接写

FSDP2 的官方文档把 per-parameter sharding 的直接收益写成 “communication-free (sharded) state dicts”:训练时参数就是 DTensor,保存 sharded state dict 时按分片写即可,不需要先把完整参数 all-gather 回来。RFC 的表述更直接:“Sharded and local state dicts are equivalent.”

FSDP1 也支持 sharded state dict,但保存前同样要先 unshard:_common_unshard_pre_state_dict_hook(torch/distributed/fsdp/_state_dict_utils.py:148-167)说明 FULL 与 SHARDED 保存都会走它(composable NO_SHARD 等例外直接返回),_sharded_pre_state_dict_hook(:504-530)也会调用它;sharded load 路径还会 all-gather(_sharded_pre_load_state_dict_hook,:582-662;:658 的 dist.all_gather_single 仅在 _use_dtensor=False 时执行,_use_dtensor=True 走 :668 的 DTensor all-gather,仍有通信)。因此显存尖峰不只属于 full 保存,sharded 保存同样要临时放一份完整参数;FSDP2 的 sharded 才不需要 all-gather。FSDP2 把「训练表示」和「分片表示」统一了,checkpoint 路径更短;代价是它不直接支持 full state dict,需要 DTensor.full_tensor() 或分布式 checkpoint(DCP)。

组合:FSDP2 的 DTensor 可以和 TP/EP 叠加

FSDP1 的 FlatParameter 是普通 1D Tensor,不是 DTensor;TP/EP 的分片信息不在它里面。FSDP2 的每个参数本来就是 DTensor,_init_sharding_spec_tp 可以把 DP mesh 和 TP/EP mesh 拼成一个 SPMD mesh,再给参数一个组合 placement(_fsdp_param.py:740-790)。这是「per-parameter」相对「flat buffer」最实质的工程差异之一:分片不再是二选一,而是可以沿多个 mesh 维度叠加。

设计文档的说法

RFC 对收益的原始表述

RFC 列出的四条收益是:fp8 权重和非 fp8 参数可以混在同一次 all-gather;frozen 与 non-frozen 参数可以混在同一通信组而不额外占显存;sharded state dict 与训练表示一致;以及给 torch.compile 留出重排通信组的空间。它还称原型在 Llama-like 模型上 “achieving on-par throughput while using less memory”——这是设计文档的自述,不是论文表格里的 benchmark。

RFC 写作时给出的代码量对比是约 3k 行 vs 原 FSDP 约 14k 行;这是当时原型的状态,不应直接当成今天两个实现的 LOC 对比。

一张表把差异收拢

表 4 · FSDP1 与 FSDP2 的实现差异(源码 commit 2ee7a63 与官方文档)
维度FSDP1FSDP2
APIFullyShardedDataParallel(module, ...) 包装器fully_shard(module, ...),返回 FSDPModule;原地替换参数
参数表示一个 module 一条 FlatParameter每个参数一个 DTensor(Shard(0))
分片单位整条 flat 按 N chunk;边界可切断参数单参数按 N chunk;每参数独立 padding
策略开关ShardingStrategy:FULL_SHARD / SHARD_GRAD_OP / NO_SHARD / HYBRID_SHARDreshard_after_forward=True/False/int(默认 None);无命名 NO_SHARD,复制场景用 DDP
通信单位每个 FlatParameter group 一次 all-gather / reduce-scatter每个 FSDPParamGroup 一次;参数分片在 collective 内拼接
state dictfull / sharded / local;full 与 sharded 保存都要先 unshard,sharded load 还会 all-gathersharded 与训练表示一致;full 走 DCP / full_tensor()
混合精度MixedPrecisionMixedPrecisionPolicy;cast 发生在 all-gather 边界
TP/EP 组合flat 参数不是 DTensor;TP 的分片信息不在 FlatParameter 里参数本身是 DTensor;DP + TP/EP 可落在同一 SPMD mesh
调优旋钮auto_wrap_policy、bucket_cap_mb、limit_all_gathersgroup 边界由 fully_shard 的 module 决定;无 bucket_cap_mb;无 limit_all_gathers

最后一行值得展开:FSDP2 没有 bucket_cap_mb,不是说它不关心通信效率,而是把「桶」换成了你划分的 module group。你调 fully_shard 的粒度,就是在调 all-gather/reduce-scatter 的粒度。官方文档也明确说,只对最顶层 module 调一次 fully_shard 会变成前向、反向各一次巨大的阻塞通信,“almost never what you want”。

07

什么时候别用 FSDP2——以及 FSDP1 迁移前该核对什么

这一节列的是边界条件,不是使用教程。先确认你的场景没有撞上这些约束。

先别分片

模型能放进单卡

如果参数 + 梯度 + 优化器状态 + 激活能放进一张卡,DDP 或单卡更简单:没有 all-gather/reduce-scatter 的额外通信,也没有 sharded state dict 的复杂度。

FSDP 的收益随 N 增长,但通信也随 N 增长;小模型上通信占比更高,先算账再决定。

先算网络

互联带宽不够

FULL_SHARD 每步约 3P(N−1)/N 个元素的通信量。N 很大时,all-gather 和 reduce-scatter 的延迟会压过计算;没有 NVLink/InfiniBand 的机器上,先做通信账再决定。

FSDP2 没有 bucket_cap_mb,group 边界就是通信边界;不要把整个模型塞进一个 group 再指望它自动重叠。

FSDP2 的硬约束

FSDP1 迁移前要核对的点

和 TP / PP 的边界

最后,FSDP 不解决激活显存。如果训练仍然 OOM,先上 activation checkpointing,再考虑继续增大 N。分片只减少参数/梯度/优化器状态,不减少每层的中间激活;激活的大小取决于 batch、序列长度和模型宽度,本页的账本没有把它算进去。