fsdp切分千亿模型需精准配置auto_wrap_policy和sharding_strategy：自定义层类名须显式传入，推荐size_based_auto_wrap_policy；单机多卡用1d devicemesh配full_shard+limit_all_gathers；混合精度优先fp16并校验bf16支持；保存检查点须用fsdp流式导出、eval模式、多文件拆分。
FSDP 能切分千亿模型，但“能跑”不等于“高效”，关键在 auto_wrap_policy 的粒度和 sharding_strategy 的匹配——错配会导致通信爆炸或显存不降反升。
为什么 transformer_auto_wrap_policy 在千亿模型上容易失效
默认的 transformer_auto_wrap_policy 只识别 TransformerEncoderLayer 和 TransformerDecoderLayer 这两类类名，但实际千亿模型（如 LLaMA-3 405B、Qwen2-72B）往往用自定义 Block 类（Qwen2DecoderLayer、LlamaDecoderLayer），或者嵌套过深（如 MoE 中的 Expert 子模块未被包裹）。结果就是：只有顶层 nn.ModuleList 被 FSDP 包裹，内部参数仍全量驻留单卡。
- 检查方法：训练前打印
model结构，确认目标层类名是否在transformer_layer_cls集合中 - 修复方式：显式传入自定义类集合，例如
transformer_layer_cls={LlamaDecoderLayer, Qwen2DecoderLayer} - 更稳妥做法：改用
size_based_auto_wrap_policy，按参数量阈值切分（如min_num_params=1e8），避免依赖类名
PyTorch 2.0+ 中 sharding_strategy 必须与模型规模对齐
千亿模型不是简单设成 ShardingStrategy.FULL_SHARD 就万事大吉。FSDP2（PyTorch ≥2.2）默认启用 DTensor 后，FULL_SHARD 会强制所有参数走 DeviceMesh 分片，若 mesh 维度没对齐硬件拓扑（比如 8 卡机器用了 mesh = DeviceMesh("cuda", [[0,1],[2,3],[4,5],[6,7]])），通信将退化为跨 NUMA 节点传输，延迟翻倍。
- 单机多卡推荐：用 1D mesh，
mesh = DeviceMesh("cuda", torch.arange(world_size)) - 千万级参数以下模型：用
ShardingStrategy.SHARD_GRAD_OP（ZeRO-2），减少 all-gather 开销 - 千亿模型必须用
FULL_SHARD，但需配合limit_all_gathers=True防止 forward 阶段内存峰值冲高
混合精度配置不兼容 bfloat16 时的静默降级风险
PyTorch 2.0+ 的 MixedPrecision 允许细粒度控制 param_dtype、reduce_dtype、buffer_dtype，但若硬件不支持 bfloat16（如 V100 或旧版 A100 驱动），bfSixteen 策略不会报错，而是自动回落到 fp16 计算 —— 此时梯度缩放（grad scaling）可能失效，loss 突然 nan。
- 务必在 init 时校验：
torch.cuda.is_bf16_supported()+nccl.version() >= (2, 10) - 千亿模型建议统一用
fp16：通信带宽压力小，且 AMP fallback 更稳定 - 禁用 buffer 混合精度：
buffer_dtype=torch.float32，避免 LayerNorm 等算子因 buffer 降级导致数值不稳定
保存千亿模型检查点的两个硬约束
直接调用 model.state_dict() 会触发全量参数聚合到 rank 0，内存瞬间暴涨数 TB，OOM 是必然结果。FSDP 提供了流式导出机制，但有两个常被忽略的前提：
-
offload_to_cpu=True必须配合rank0_only=True，否则每张卡都试图把本地分片 dump 到 CPU 内存 - 保存前需确保模型处于
eval()模式，否则activation_checkpointing可能残留未释放的中间变量 - 不要用
torch.save()直接序列化整个 model —— 改用save_policy = FullStateDictConfig(...)+FSDP.state_dict_type(...)上下文管理器
真正棘手的是：千亿模型的 state dict 键名极长（含完整模块路径），文件系统 metadata 容易打爆 inode。生产环境必须拆成 model.bin + optimizer.bin + rng_state.pth 多文件保存，且路径不能含中文或空格。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
