Chapter 05
FSDP2：和 FSDP1 的区别在哪
FSDP 是 PyTorch 自家的 ZeRO-3 实现。2024 重写后叫 FSDP2，TorchTitan 全用 FSDP2。本章讲两者差异。
5.1FSDP 是什么
Fully Sharded Data Parallel：把 model 参数、梯度、optimizer state 切到所有 DP rank 上，每张卡只持 1/N 份。forward / backward 时按需 all-gather 重组完整 weight。等价于 DeepSpeed ZeRO-3。
5.2FSDP1 的痛点
- FlatParameter：把多个 weight 拼成一个大 flat tensor，user-API 不友好；
- state_dict 怪：要专门 callback 才能拿原始 weight；
- 跟 TP 配合差：FlatParameter 跟 DTensor 不兼容；
- 不能 partial freezing：要冻一部分 layer 很麻烦；
- 编译不顺：torch.compile 兼容性差。
5.3FSDP2 的设计
| FSDP1 | FSDP2 | |
|---|---|---|
| 参数表示 | FlatParameter | DTensor |
| state_dict | callback 风格 | 原生 named_parameters() 直接用 |
| 跟 TP 配合 | 难 | 原生兼容（都基于 DTensor） |
| partial freeze | 难 | 直接 p.requires_grad = False |
| torch.compile | 限制多 | 原生支持 |
| API 入口 | FullyShardedDataParallel | fully_shard（函数） |
5.4FSDP2 API 风格
from torch.distributed.fsdp import fully_shard
# FSDP2 是函数式 API，不是包装类
for layer in model.layers:
fully_shard(layer, mesh=mesh["dp_shard"])
fully_shard(model, mesh=mesh["dp_shard"])
# 之后正常 forward/backward
out = model(x)
out.loss.backward()
optimizer.step()不再有 FullyShardedDataParallel(model) 这种包装；而是用 fully_shard() 函数原地给指定 module 标记 sharding。
5.5HSDP：Hybrid Shard Data Parallel
FSDP2 支持二维 mesh，外维 Replicate + 内维 Shard，节省跨机带宽：
mesh = init_device_mesh("cuda", (2, 8), mesh_dim_names=("dp_replicate", "dp_shard"))
# 假设 16 GPU 分布在 2 节点 × 8 GPU
# - 节点内 8 卡做 FSDP shard
# - 节点间 2 个副本之间做 DDP all-reduce
fully_shard(model, mesh=mesh, reshard_after_forward=True)对应 TorchTitan 配置：
[training]
data_parallel_replicate_degree = 2 # 节点间副本
data_parallel_shard_degree = 8 # 节点内 shard5.6TorchTitan 里的 FSDP2 用法
看 torchtitan/distributed/parallelize_llama.py::apply_fsdp：
def apply_fsdp(model, dp_mesh, ...):
fsdp_config = {"mesh": dp_mesh, "mp_policy": MixedPrecisionPolicy(
param_dtype=torch.bfloat16,
reduce_dtype=torch.float32,
)}
# 给每个 TransformerBlock 单独 fsdp
for layer in model.layers.values():
fully_shard(layer, **fsdp_config)
fully_shard(model, **fsdp_config)5.7性能对比
Llama-3 8B、8×A100、bf16、seqlen=8K：
| 方案 | MFU | tok/s/GPU | 备注 |
|---|---|---|---|
| DDP (太大跑不下) | — | — | OOM |
| FSDP1 | ~42% | ~3,400 | baseline |
| FSDP2 | ~48% | ~4,000 | compile 友好 |
| FSDP2 + torch.compile | ~54% | ~4,500 | + compile |
| FSDP2 + TP=2 | ~52% | ~4,300 | 大模型显存压力大时 |
5.8常见踩坑
| 现象 | 处理 |
|---|---|
| OOM 第 0 step | fsdp_use_orig_params 在 FSDP2 默认 true；问题在 chunk 太大 |
| save ckpt 后加载错 | 必须用 dcp（详见 ch10） |
| load HF state_dict 失败 | name 不一致；FSDP2 的 named_parameters() 不带 fsdp 前缀 |
| 慢 1.3× | 检查 reshard_after_forward；显存够时关掉 |
5.9这章你需要带走的
- FSDP2 = 基于 DTensor 重写的 FSDP；
- API 风格从类包装变成
fully_shard()函数； - HSDP（Hybrid Shard）= 跨节点 Replicate + 节点内 Shard，省带宽；
- TorchTitan 全用 FSDP2，跟 TP / PP / CP 都用 DTensor 无缝配合。