FSDP 高级配置与性能调优
FSDP（FullyShardedDataParallel）是 PyTorch 原生的全分片数据并行方案。本文介绍 FSDP 的高级配置选项和生产环境性能调优策略。
分片策略
python
from torch.distributed.fsdp import ShardingStrategy
# FULL_SHARD: 完全分片（ZeRO-3），最省显存
model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)
# SHARD_GRAD_OP: 仅分片梯度和优化器状态（ZeRO-2）
model = FSDP(model, sharding_strategy=ShardingStrategy.SHARD_GRAD_OP)
# NO_SHARD: 不分片（标准 DDP）
model = FSDP(model, sharding_strategy=ShardingStrategy.NO_SHARD)包装策略
FSDP 的包装策略决定了分片粒度，直接影响性能：
python
from torch.distributed.fsdp import MixedPrecision
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
# 基于 Transformer 层的自动包装
auto_wrap_policy = transformer_auto_wrap_policy(
transformer_layer_cls={LlamaDecoderLayer}
)
# BF16 混合精度
mp_policy = MixedPrecision(
param_dtype=torch.bfloat16,
reduce_dtype=torch.bfloat16,
buffer_dtype=torch.bfloat16,
)
model = FSDP(
model,
auto_wrap_policy=auto_wrap_policy,
mixed_precision=mp_policy,
sharding_strategy=ShardingStrategy.FULL_SHARD,
device_id=torch.cuda.current_device(),
)包装粒度影响
| 包装粒度 | 通信次数 | 每次通信量 | 总通信量 | 推荐 |
|---|---|---|---|---|
| 整个模型 | 2 | 全部参数 | 2×参数 | 否 |
| Transformer 层 | 2×层数 | 层参数 | 2×参数 | 推荐 |
| 单个算子 | 2×算子数 | 算子参数 | 2×参数 | 过细 |
最优包装
基于 Transformer 层的包装是最优粒度：通信次数适中，每次通信量可被 NCCL 优化，与计算重叠效果好。
激活检查点
FSDP 支持激活检查点，用计算换显存：
python
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
apply_activation_checkpointing,
checkpoint_wrapper,
)
# 对每层应用激活检查点
def check_fn(submodule):
return isinstance(submodule, LlamaDecoderLayer)
apply_activation_checkpointing(model, checkpoint_wrapper_fn=checkpoint_wrapper, check_fn=check_fn)预取优化
python
from torch.distributed.fsdp import CPUOffload
model = FSDP(
model,
# 参数预取：在前向传播计算当前层时预取下一层参数
forward_prefetch=True,
# CPU Offload
cpu_offload=CPUOffload(offload_params=True),
)性能对比
| 配置 | 显存 (7B) | 吞吐 (tok/s) | 相对速度 |
|---|---|---|---|
| DDP (8GPU) | 28GB | 5200 | 100% |
| FSDP ZeRO-2 | 18GB | 4800 | 92% |
| FSDP ZeRO-3 | 4GB | 4000 | 77% |
| FSDP ZeRO-3 + ActCkpt | 3GB | 3200 | 62% |
显存与速度权衡
激进的分片和检查点策略节省显存但降低速度。建议优先使用 ZeRO-2（最佳平衡），仅在显存不足时升级到 ZeRO-3。