FSDP 完全分片数据并行

定义

FSDP(Fully Sharded Data Parallel,完全分片数据并行)是 PyTorch 官方提供的分布式训练方案,将模型参数、梯度和优化器状态全部分片存储到数据并行组的各 GPU 上,思想与 DeepSpeed ZeRO Stage 3 一致。

原理

训练时每个 GPU 只常驻自己负责的参数分片;前向或反向到某一层时,通过 all-gather 临时聚合该层完整参数,用完即释放;反向传播后用 reduce-scatter 将梯度归约并分片。显存占用随并行卡数近似线性下降,代价是增加通信量,可通过通信与计算重叠缓解。

使用

FSDP 自 PyTorch 1.11 起正式提供,支持混合精度、激活检查点、CPU 卸载,并可与张量并行组合;新一代 FSDP2 改进了接口与 DTensor 集成。Hugging Face Accelerate、TRL 等库均支持一键启用。

对比

相较 DeepSpeed,FSDP 的优势是 PyTorch 原生、无额外依赖、与生态耦合更好;功能丰富度上 DeepSpeed 提供的卸载与推理配套更多,两者是微调与训练中最常见的两种选择。