基本原理（Fully Sharded Data Parallel，完全分片数据并行）
- 一种参数分片技术，为超大规模模型的分布式训练设计
- 参数分片：模型的参数被分片存储到不同的 GPU，每个 GPU 仅存储自己负责的参数分片（Sharded Parameters）
- 数据分发：与 DDP 类似
- 前向传播：FSDP 将需要的参数分片加载到显存，在完成计算后，卸载这些参数以节省显存
- 后向传播：梯度计算完成后，梯度也被分片，并通过 AllReduce 操作同步到所有 GPU 优化器分片：优化器状态（如动量）也被分片存储并在必要时同步
参数 ShardingStrategy 的不同取值决定了模型的划分方式
- FULL_SHARD：将模型参数、梯度和优化器状态都切分到不同的 GPU 上，类似 ZeRO-3
- SHARD_GRAD_OP：将梯度、优化器状态切分到不同的 GPU 上，每个 GPU 仍各自保留一份完整的模型参数，类似 ZeRO-2
- NO_SHARD：不切分任何参数，类似 ZeRO-0

配置文件示例

fsdp_config:
  world_size: 8
  local_rank: 0
  shard_optimizer_state: true
  mixed_precision: true
  fp16: true
  activation_checkpointing: true
  device: cuda
  offload_params: true
  offload_optimizer_state: false
  use_reentrant: true
  model_parallel_size: 2
  checkpoint_interval: 1000
  checkpoint_dir: ./checkpoints

参数说明

参数名称	可选值	示例	含义
`world_size`	正整数	8	分布式训练中的总进程数，通常等于设备数
`local_rank`	正整数	0	当前进程的本地进程号，用于标识分布式训练中的某个节点
`shard_optimizer_state`	`true` 或 `false`	true	是否对优化器状态进行切分以节省内存
`mixed_precision`	`true` 或 `false`	true	是否启用混合精度训练，以提高训练速度并减少内存占用
`fp16`	`true` 或 `false`	true	是否启用 16 位浮动点数精度 (FP16)，在一些硬件上可以显著提升训练速度
`activation_checkpointing`	`true` 或 `false`	true	是否启用激活检查点，以减少显存使用
`device`	`cuda` 或 `cpu`	cuda	训练使用的设备类型，通常为 `cuda`（GPU）或 `cpu`
`offload_params`	`true` 或 `false`	true	是否将模型参数卸载到 CPU 内存而不是 GPU 显存
`offload_optimizer_state`	`true` 或 `false`	false	是否将优化器状态卸载到 CPU 内存而不是 GPU 显存
`use_reentrant`	`true` 或 `false`	true	是否使用重入式 (reentrant) 机制，这对于某些特定的训练优化是必需的
`model_parallel_size`	正整数	2	模型并行的大小，指定了在多个设备之间划分模型的数量
`checkpoint_interval`	正整数	1000	每隔多少步进行一次检查点保存
`checkpoint_dir`	字符串（路径）	./checkpoints	存储检查点文件的目录路径

ZeRO 和 FSDP
- 更细粒度的分片方式：FSDP 支持按层（layer-wise）乃至按参数组（parameter group-wise）进行参数分片，分片粒度更精细
- 按需聚合参数，计算后即时释放，采用按需激活的参数加载机制：正向传播时，仅对当前计算所需的层执行 all-gather 聚合参数，该层计算完成后便立即释放对应参数；反向传播过程也遵循相同逻辑；这与 ZeRO Stage 3 每次都 all-gather 完整 FP16 权重的方式形成显著区别
- 梯度通信更早、更细粒度触发：在反向传播阶段，FSDP 无需等待所有层的梯度全部计算完成，而是在单个层梯度计算完毕后，立即执行 reduce-scatter，完成分片梯度的同步与累积，通信时机更靠前、粒度更细
- 与激活重计算的常用组合：激活重计算（Activation Checkpointing，依托 torch.utils.checkpoint 实现）并非 FSDP 独有特性，但二者通常搭配使用；激活重计算通过在反向时重新计算激活值、而非在前向时缓存激活值，进一步降低显存占用

FSDP-1 和 FSDP-2

FSDP1 的分片逻辑
- 将模块内所有参数先扁平拼接为一个巨大的 FlatParameter 一维张量，再在不同 GPU 间整体分片
- 所有参数被合并为一个同质化对象，丢失了原始参数的独立元数据（如 dtype、requires_grad），只能靠复杂的偏移量逻辑管理
- 流程繁琐：需要先拼接、再分片，后续还原时还要反向推断每个进程的参数片段，实现复杂度高
- 难以适配部分参数冻结（如 LoRA）等场景，必须额外封装修改
FSDP2 的分片逻辑
- 引入 DTensor（分布式张量），每个参数 / 层独立表示为一个 DTensor，直接按指定维度在 GPU 间分片
- 保留了每个原始张量的完整元数据，包括分片方式、设备信息、参数属性等，无需额外维护偏移量
- 流程直观高效：直接对单个参数分片，避免了 FSDP1 中拼接、分片、还原的复杂逻辑

主要对比

改进点	核心说明
1. 简化实现	消除了 `FlatParameter` 的复杂拼接逻辑，直接使用 `DTensor` 管理参数，避免了内部偏移量的维护工作，代码更简洁易维护
2. 原生支持部分参数冻结	`DTensor` 的设计天然兼容 LoRA 等微调方法，开箱即用，无需额外封装或修改
3. 优化 Checkpoint 管理	- 默认使用 `SHARDED_STATE_DICT` 存储分片检查点，每个 Rank 仅保存自身负责的参数与元数据，无需提前 `all-gather`，存储效率大幅提升。 - 支持异步 Checkpoint 保存：先将参数复制到 CPU，主训练线程继续运行，读写进程异步写入磁盘，减少 I/O 阻塞
4. 稳定的内存管理	不再依赖 `torch.Tensor.record_stream`（易导致非确定性行为），改用底层 CUDA 多流同步，实现稳定可预测的内存行为，减少训练中断与性能波动
5. 通信与计算重叠优化	- 默认启用 FSDP1 中需手动配置的 `BACKWARD_PRE` 预取选项。 - 新增隐式预取（implicit prefetching），更好地重叠 `all-gather` 通信与计算，隐藏通信延迟，提升 GPU 利用率
6. 更好的兼容性	对 FP8 等混合精度类型、2D/3D 并行策略（如数据并行 + 张量并行 + 流水线并行）的支持更完善，适配更复杂的大规模训练场景

FSDP

On this page

页面信息