Appearance
FSDP / DeepSpeed ZeRO
全分片数据并行 (Fully Sharded Data Parallel, FSDP) 和 DeepSpeed ZeRO (Zero Redundancy Optimizer) 是解决大模型训练内存瓶颈的两大核心技术。它们通过将模型参数、梯度和优化器状态分片到所有数据并行 worker 上，显著减少单设备内存占用。
问题背景
在标准数据并行 (DDP) 中，每个 GPU 都复制完整的模型参数、梯度和优化器状态：
- 模型参数: $4 imes ext{params}$ bytes (FP32 主权重 + FP16 参数)
- 梯度: $2 imes ext{params}$ bytes (FP16 梯度)
- 优化器状态: $12 imes ext{params}$ bytes (Adam: 动量 + 二阶矩 + FP32 副本)
一个 10B 参数的模型在混合精度训练下需要 ~60GB 内存，远超单张 GPU 容量。
DeepSpeed ZeRO 三阶段
ZeRO 通过逐步分片三类状态来减少冗余：
ZeRO-1: 优化器状态分片
将 Adam 的动量和二阶矩分片到各数据并行 rank：
- 内存节省: ~4x (从 16x 到 4x 模型大小)
- 通信: 梯度 AllReduce 后需要 AllGather 优化器状态
ZeRO-2: + 梯度分片
在 ZeRO-1 基础上增加梯度分片：
- 内存节省: ~8x
- 通信: 梯度只在需要时收集
ZeRO-3: + 参数分片
将模型参数也分片到各 rank，每个 rank 只保存部分参数：
- 内存节省: 与数据并行度成正比（如 64 张 GPU 节省 64x）
- 通信: 前向/反向传播时需要动态收集参数
| 阶段 | 分片内容 | 内存节省 | 通信开销 |
|---|---|---|---|
| DDP | 无 | 1x | 低 |
| ZeRO-1 | 优化器状态 | ~4x | 中 |
| ZeRO-2 | +梯度 | ~8x | 中 |
| ZeRO-3 | +参数 | ~N× (DP度) | 高 |
ZeRO-Offload 与 ZeRO-Infinity
ZeRO-Offload
将优化器状态和计算卸载到 CPU 内存：
- 单张 V100 可训练 10B+ 参数模型
- CPU 计算优化器更新，GPU 专注前向/反向传播
ZeRO-Infinity
进一步将参数、梯度和优化器状态卸载到 NVMe SSD：
- 单节点可训练 万亿级 参数模型
- 通过重叠数据传输隐藏存储延迟
PyTorch FSDP
PyTorch 1.11+ 原生集成的全分片数据并行：
- API:
torch.distributed.fsdp.FullyShardedDataParallel - 特点: 更简洁的 API，与 PyTorch 生态深度集成
- 功能: 自动分片、混合精度、激活值检查点、向量级并行
- 演进: 早期由 Meta 的 FairScale 项目探索，后合并入 PyTorch 核心
FSDP vs ZeRO 对比
| 维度 | PyTorch FSDP | DeepSpeed ZeRO |
|---|---|---|
| 归属 | PyTorch 原生 | Microsoft 开源 |
| API 风格 | 简洁、Pythonic | 功能丰富、配置驱动 |
| 集成度 | 与 PyTorch 深度集成 | 支持多种框架 (PyTorch/TensorFlow) |
| 附加功能 | 基础分片 | + 3D 并行 + MoE + 压缩 + 推理引擎 |
| 适用场景 | PyTorch 用户、研究环境 | 超大规模训练、生产环境 |
重要框架与工具
- DeepSpeed: Microsoft 开源，集成 ZeRO 系列、3D 并行、MoE、推理加速
- PyTorch FSDP: PyTorch 原生，
torch.distributed.fsdp，更简洁的 API - FairScale: Meta 的 FSDP 早期实现，已合并入 PyTorch
- Colossal-AI: 统一并行框架，集成 ZeRO/Gemini/Sequence Parallelism
- Hugging Face Accelerate: 简化 FSDP/DeepSpeed 配置的高级 API
与相关概念的关系
- FSDP vs DDP: FSDP 是 DDP 的内存优化扩展，计算模式相似但参数按需收集
- ZeRO vs Model Parallelism: ZeRO 是纯数据并行优化，可与张量/流水线并行组合
- Offload vs Latency: 卸载增加通信延迟，需要与计算重叠隐藏
- Activation Checkpointing: 与 FSDP/ZeRO 正交，通过重计算换激活值内存
Sources
- Rajbhandari et al., "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models" (SC20)
- Ren et al., "ZeRO-Offload: Democratizing Billion-Scale Model Training" (USENIX ATC 2021)
- Rajbhandari et al., "ZeRO-Infinity: Breaking GPU Memory Wall for Extreme Scale Deep Learning" (2021)
- PyTorch FSDP Documentation: https://pytorch.org/docs/stable/fsdp.html
- DeepSpeed: https://github.com/microsoft/DeepSpeed