fsdp

PyTorch FSDP：把参数、梯度和优化器状态分片到多卡，训练放不下的模型

补充说明

FSDP（Fully Sharded Data Parallel）把模型参数、梯度和优化器状态按数据并行组切分到各 GPU，前向/反向时按需 all-gather，计算完再释放。它适合单卡放不下、但又不想上张量并行的模型；和 DeepSpeed ZeRO-3 的目标类似，区别在实现、配置方式和生态集成。

FSDP 有两种常见形态：经典 FSDP wrapper 和 PyTorch 2.x 的 FSDP2（fully_shard）。FSDP2 使用 DTensor，支持更灵活的分片和 torch.compile，新项目优先评估 FSDP2。使用 FSDP 需要 torchrun 启动，不能直接用 python train.py 跑多卡。

常用示例

# 8 卡启动
torchrun --standalone --nproc_per_node=8 train_fsdp.py

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy

model = MyModel().cuda()
mp = MixedPrecision(param_dtype=torch.bfloat16,
                    reduce_dtype=torch.bfloat16,
                    buffer_dtype=torch.bfloat16)

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=mp,
    device_id=torch.cuda.current_device(),
    use_orig_params=True,
)

# 按层 wrapping，避免一次性 all-gather 整个模型
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torch.distributed.fsdp import FSDP

model = FSDP(
    model,
    auto_wrap_policy=transformer_auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
)

FSDP2 的写法：

from torch.distributed.fsdp import fully_shard

for layer in model.layers:
    fully_shard(layer)
fully_shard(model)

易错点与排错

• 用 python 直接启动：FSDP 需要进程组和 LOCAL_RANK，多卡必须用 torchrun/torch.distributed.launch。
• 没有 wrap policy：整个模型作为一个 FSDP 单元，all-gather 峰值可能仍然放不下；按 Transformer block 分层 wrapping。
• 参数初始化不一致：各 rank 必须用相同种子/逻辑初始化，或在 FSDP 后从 rank 0 广播；否则分片参数不一致。
• no_sync 使用不当：梯度累积时用 no_sync 减少通信，但最后一个 micro-batch 必须同步；否则梯度缺失。
• checkpoint 保存/加载错误：FSDP 的 state_dict 是分片的，要用 FullStateDictConfig/LocalStateDictConfig 或分布式 checkpoint；直接 torch.save(model.state_dict()) 可能只保存部分。
• BatchNorm 与 FSDP：BN 的统计量需要同步，FSDP 不自动处理；小 batch 或多卡场景考虑 SyncBatchNorm 或 LayerNorm 架构。
• 混合精度 dtype 不匹配：param_dtype、reduce_dtype、buffer_dtype 要一致规划；某些算子对 BF16 不支持。
• FSDP 与 torch.compile 顺序错误：先 compile 还是先 fully_shard 随版本不同；按 PyTorch 版本文档调整，并用小模型验证。

相关命令

• torchrun — PyTorch 分布式启动器。
• deepspeed — ZeRO 系列显存优化。
• accelerate — Hugging Face 分布式训练封装。
• megatron-lm — 张量/流水线并行。
• activation-checkpointing — 进一步降低激活显存。
• distributed-training-troubleshooting — 挂起、OOM 与 NCCL 排障。

参考链接

• PyTorch FSDP 文档
• PyTorch FSDP 教程
• FSDP2 设计

读完自测

合上页面，先回答，再展开对照。

1. FULL_SHARD 和 SHARD_GRAD_OP 的主要区别是什么？
2. 为什么 FSDP 要按 Transformer block 做 auto wrap？
3. FSDP 的 checkpoint 为什么不能直接 torch.save(model.state_dict())？