两者都把参数、梯度、优化器状态切到每张卡只留 1/N,前向前 all-gather,反向后 reduce-scatter。
真正的分水岭是切的单位:FSDP1 把一个 module 的参数压成一条 FlatParameter 再切,FSDP2 让每个参数自己成为 DTensor(Shard(0))。
这个选择决定了 padding、state dict、优化器看到什么,以及能不能和 TP/EP 直接拼。
常驻显存的账,两者一样。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 只是起点,真正影响工程的是这 1/N 在内存里怎么摆。
DDP 的账本很直观:每张卡都放一份完整的参数、梯度、优化器状态。用 P 表示参数量,bp、bg 表示参数和梯度的字节数,bopt 表示每个参数的优化器状态字节数,每卡常驻就是:
按 7B 参数、bf16 参数与梯度、AdamW 的 fp32 一阶/二阶动量(8 B/param)算,完整副本是 84 GB。这已经超过一张 80 GB 卡的可分配显存,而且还没算激活。FSDP 的思路是把这三类状态都切成 N 份,每卡只留一份:
这个 1/N 的结论对 FSDP1 和 FSDP2 都成立。真正的问题是:这 1/N 在内存里长什么样?是一条被 flatten 的连续 buffer,还是许多个独立的参数分片?这个表示方式会连带影响三件事:
下面的 02 先看前人怎么做,03 再把两种表示拆开。
这一节建立对照:DDP 解决通信去中心化,ZeRO/FSDP 解决显存,但「怎么表示一份分片」是 FSDP 才真正暴露的问题。
每张卡维护完整模型,反向后用 ring all-reduce 同步梯度。实现简单、通信与计算重叠好,但显存随参数量线性增长。
卡在哪:模型 + 优化器状态超过单卡时,DDP 只能换更大的卡,不能靠加卡解决。
ZeRO-1 切优化器状态,ZeRO-2 再切梯度,ZeRO-3 连参数一起切。FSDP 的 FULL_SHARD 对应最后一档:前向前 all-gather 参数,反向后 reduce-scatter 梯度。
代价:每步多出两次 all-gather(前向、反向各一次)和一次 reduce-scatter;通信量随 N 增大而趋近 3P 个元素。
把 ZeRO 的三个阶段和每卡常驻写成一张表,差别一目了然(b 为参数/梯度字节数,优化器状态按 bopt 计):
| 方案 | 参数 | 梯度 | 优化器状态 | 每卡常驻 |
|---|---|---|---|---|
| DDP | 完整 | 完整 | 完整 | P × (2b + bopt) |
| ZeRO-1 | 完整 | 完整 | 1/N | 2Pb + P bopt/N |
| ZeRO-2 | 完整 | 1/N | 1/N | Pb + P(b + bopt)/N |
| ZeRO-3 / FSDP FULL_SHARD | 1/N | 1/N | 1/N | P(2b + bopt)/N |
这张表只说明「切哪些状态」,没有说明「切出来的东西长什么样」。FSDP1 和 FSDP2 的分歧就在后者。
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)。
两种表示都能做到 1/N;差别在 padding、state dict,以及优化器看到的参数类型。
FSDP1 的 FlatParameter 是一个 1D 的 nn.Parameter,由 module 内所有原始参数 flatten 后拼接而成(torch/distributed/fsdp/_flat_param.py:202)。拼接不是简单 concat:
use_orig_params=True 时,从第二个参数起,按需在参数前插入 padding,使其起始地址落在 16 B 边界(_flat_param.py:565-567、:710-713、:2790-2795);对齐单位 a = 16 / bytes(dtype),单个参数最多补 a−1 个元素。按前向/反向参数 dtype 算,bf16 下对齐单位是 8 个元素(16/2),fp32 下是 4 个元素(16/4)。use_orig_params=True):拼完后再补到 N 的倍数,保证每个 rank 的 chunk 等长(_flat_param.py:759-763;aligned_numel 只在 use_orig_params=True 时非 0,见 :565-566、:609-613)。默认 use_orig_params=False 时没有 world padding;torch.chunk 可能返回不等长 chunk,代码补到 chunks[0].numel(),即每卡有效 ceil(L/N)。torch.flatten(tensor).chunk(N)(_flat_param.py:1104),每个 rank 拿一个 chunk;默认 use_orig_params=False 时 chunk 可能不等长,代码按 chunks[0].numel() 补齐;chunk 边界可能落在某个原始参数中间。use_orig_params=True)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 的入口是 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)。分片过程是:
_chunk_with_empty(param_data, N, dim=0),第 i 个参数的分片大小是 ceil(si/N)(_fsdp_param.py:362-369)。nn.Parameter(to_sharded_dtensor(...))(_fsdp_param.py:394-398)。所以优化器直接拿到 DTensor,而不是 flat view。padded_sharded_param_size;padding 是按参数算的,不会把一个参数的边界和另一个参数混在一条 chunk 里。_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 的粒度。
下面的账本把参数、梯度、优化器状态分开算。拖动参数量、世界规模、精度,或切换分片策略;所有数字都按同一组公式实时重算。
口径:优化器状态固定按 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。
组 B:P=13B,N=4,fp32,SHARD_GRAD_OP。
把滑块拉到同样参数应显示同样的数;若不一致,以脚本输出为准并在报告里记下差异。
这组账本说明什么:FULL_SHARD 和 SHARD_GRAD_OP 的常驻分片相同,区别在通信——前者反向还要再 all-gather 一次。SHARD_GRAD_OP 在前向后不释放完整参数,省掉这次 all-gather,但高峰持续时间更长。NO_SHARD 完全不省显存,官方已把它标为 deprecated,建议直接用 DDP。
读完这节你应该知道:FULL_SHARD、SHARD_GRAD_OP、NO_SHARD 各自的通信账,以及 FSDP2 在工程上换掉了什么。
先看 7B、N=8、bf16 的三种策略。常驻和通信都按上一节的公式算,GB 为十进制:
| 策略 | 每卡常驻 | 最坏峰值(含完整参数) | 通信 / iter / 卡 | FSDP2 等价 |
|---|---|---|---|---|
| FULL_SHARD | 10.50 GB | 24.50 GB | 36.75 GB | reshard_after_forward=True(需显式;默认 None) |
| SHARD_GRAD_OP | 10.50 GB | 24.50 GB | 24.50 GB | reshard_after_forward=False(需显式;默认 None,近似) |
| NO_SHARD | 84.00 GB | 84.00 GB | 24.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 在 _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 把通信状态收进 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 这个边界上。
FSDP 本身是数据并行维度的分片;TP 是层内维度,PP 是层间维度。组合时,FSDP 的 shard group 只应覆盖同一份模型副本——对 PP 来说就是同一个 stage 的 DP 副本,不能跨 stage all-gather 不存在的参数 推断(由 FSDP 的 all-gather 语义与 PP 的分 stage 所有权推出)。
_init_sharding_spec_tp 明确处理「DP mesh + TP/EP mesh」的组合,支持 1D TP/EP 或 2D EP+TP,mesh ndim 2–4(_fsdp_param.py:740-790)。也就是说 TP 的分片和 FSDP 的分片可以落在同一个 DTensor 的 placement 里,而不是先把 TP 参数拆回普通 Tensor 再 flatten。(Replicate(), Shard(0)):第 0 维跨节点复制,第 1 维在节点内分片。FSDP1 对应的策略是 HYBRID_SHARD,通过 device_mesh 构造。fully_shard。不要把一个 stage 的参数和另一个 stage 放进同一个 group。二维 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)。
这张图把两种表示放在同一张时间线上:左边是 FSDP1 的 flat chunk,右边是 FSDP2 的 per-parameter DTensor。
_flat_param.py:565-567、:759-763、:1104;
FSDP2 的单参数 chunk 与 DTensor 包装见 _fsdp_param.py:362-369、:394-398。
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 归属不同,这是下一节显存对比的根源。
把六步按时间排开:
FlatParameter 并按 N 切;FSDP2 把每个参数换成 DTensor(Shard(0))。use_orig_params=True 会提供原始参数的 view;FSDP2 的普通参数是 all-gather buffer 的 view,DTensor 参数(TP/EP)保持自己的 placement。reshard_after_forward=True(FSDP2 需显式;默认 None 按非根 True、根 False 解析),释放完整参数,只留分片;False 则保留到反向结束。这一节回答两个容易混的问题:per-parameter DTensor 到底省不省显存?FSDP2 换来的工程收益是什么?
取两个很小的参数,大小分别是 3 和 5 个元素,N=4,bf16。这个例子故意让参数不能被 N 整除,好把 padding 的归属看清楚:
| 实现 | 计算 | 每卡元素数 | 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)/4 | 4 | P2 前对齐 + 整条 world-size 余量 |
| FSDP2 | ceil(3/4)+ceil(5/4) = 1+2 | 3 | 每个参数自己的 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)。
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 列出的四条收益是: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 对比。
| 维度 | FSDP1 | FSDP2 |
|---|---|---|
| API | FullyShardedDataParallel(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_SHARD | reshard_after_forward=True/False/int(默认 None);无命名 NO_SHARD,复制场景用 DDP |
| 通信单位 | 每个 FlatParameter group 一次 all-gather / reduce-scatter | 每个 FSDPParamGroup 一次;参数分片在 collective 内拼接 |
| state dict | full / sharded / local;full 与 sharded 保存都要先 unshard,sharded load 还会 all-gather | sharded 与训练表示一致;full 走 DCP / full_tensor() |
| 混合精度 | MixedPrecision | MixedPrecisionPolicy;cast 发生在 all-gather 边界 |
| TP/EP 组合 | flat 参数不是 DTensor;TP 的分片信息不在 FlatParameter 里 | 参数本身是 DTensor;DP + TP/EP 可落在同一 SPMD mesh |
| 调优旋钮 | auto_wrap_policy、bucket_cap_mb、limit_all_gathers | group 边界由 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”。
这一节列的是边界条件,不是使用教程。先确认你的场景没有撞上这些约束。
如果参数 + 梯度 + 优化器状态 + 激活能放进一张卡,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 再指望它自动重叠。
fully_shard 后是 DTensor,优化器必须用它构造;先建优化器再调 fully_shard 的旧习惯要改。fully_shard 的签名里没有 sharding_strategy(_fully_shard.py:102-120)。需要完全不分片时用 DDP;FSDP1 的 NO_SHARD 也已被标记 deprecated,官方建议 DDP(api.py:50-55)。DTensor.full_tensor() 或 DCP;如果下游工具只吃普通 full state dict,迁移时要预留一层转换。fully_shard 应该按层调用,再调根模块;只调根会得到前向、反向各一次巨大的阻塞 all-gather(_fully_shard.py:145-152 与官方文档)。_fsdp_param.py:353-359 明确 raise NotImplementedError;默认 dim-0 分片可以处理不均匀大小。_fully_shard/_fsdp_state.py:236-248,错误信息明确提示 “another FSDP group”)。可以用 fully_shard([module_a, module_b]) 把两个模块放进同一个 group。use_orig_params=False 是默认值:named_parameters() 暴露的是 FlatParameter,不是原始参数;需要逐参数优化器设置、冻结/可训练混用、或 torch.compile 时,要显式设 True(fully_sharded_data_parallel.py:412;官方文档称 True 是 torch.compile() 的必要条件)。use_orig_params=False 时每个实例的参数必须整体冻结或整体可训练。迁移到 FSDP2 前先确认这些约束是否已经绕开。NO_SHARD 已 deprecated:如果现在的代码只是拿 FSDP 当 DDP 用,直接换 DDP 更合适(api.py:50-55)。N 会同步增加 all-gather/reduce-scatter 的通信量。最后,FSDP 不解决激活显存。如果训练仍然 OOM,先上 activation checkpointing,再考虑继续增大 N。分片只减少参数/梯度/优化器状态,不减少每层的中间激活;激活的大小取决于 batch、序列长度和模型宽度,本页的账本没有把它算进去。