FSDP 配置错误会导致 dtype 冲突、NaN loss 和 checkpoint 失败，主因是未被包裹的 buffer/参数、auto_wrap_policy 不匹配、sharding 策略误用及 torch.compile 顺序错误；须手动检查包裹、禁用自动 cast、定制策略、后置 compile 并严格遵循初始化顺序。
FSDP 不是给模型套个 FullyShardedDataParallel 就能跑通的——直接 wrap 会爆 RuntimeError: expected scalar type Float but found Half，训练中途 loss 突然变 NaN，或 checkpoint 加载失败，几乎都是配置错位导致的。
为什么 FullyShardedDataParallel 一启动就报 dtype 错误
根本原因不是混合精度设错了，而是模型里存在 FSDP 没管到的 buffer 或参数。比如手动调用 self.register_buffer("cache", ...)（尤其 persistent=False 时），或者自定义 nn.Parameter 写在 __init__ 里但没被 auto_wrap_policy 覆盖。
这类 tensor 仍保持原始 dtype（如 torch.float32），而 FSDP 默认在 forward 前把参数 cast 成 fp16 或 bf16，计算时类型不匹配就崩。
- 检查所有子模块是否都被
fsdp_wrap包裹：打印model.named_modules()，确认目标层（如LlamaDecoderLayer）确实被FSDP实例接管 - 禁用 FSDP 自动 cast：传入
mixed_precision=None，改由你在forward()入口统一加torch.amp.autocast(dtype=torch.bfloat16) - 确保
GradScaler(init_scale=...)的初始值和模型刚加载时的 loss 量级匹配（例如 LLaMA-3 405B 初始 loss ≈ 10–12，init_scale设成 65536 更稳）
怎么选 ShardingStrategy 和 BackwardPrefetch
选错策略不会报错，但显存不降反升、通信卡死、吞吐掉一半——全是静默失效。
立即学习“Python免费学习笔记（深入）”；
ShardingStrategy.FULL_SHARD 是千亿模型唯一可行选项，但它要求每层都做 all-gather；若你对 Embedding 层也启用它，而 vocab_size 是 128K，那一次 all-gather 就可能卡住整机 NCCL。
- Embedding 层必须单独排除：用
lambda module: not isinstance(module, nn.Embedding)过滤 - 小 batch（≤2）+ 深层模型（≥60L）：用
FULL_SHARD+BackwardPrefetch.BACKWARD_PRE，提前拉下一层参数掩盖通信延迟 - 大 batch（≥8）或 MoE 类模型：改用
ShardingStrategy.SHARD_GRAD_OP，只分片梯度和优化器状态，避免频繁all-gather - 单机 8 卡务必配
DeviceMesh("cuda", torch.arange(8))，别用二维 mesh（如[[0,1],[2,3],...]），否则跨 NUMA 节点通信延迟翻倍
为什么 transformer_auto_wrap_policy 在 LLaMA/Qwen 上大概率失效
这个 policy 只认硬编码的类名：TransformerEncoderLayer 和 TransformerDecoderLayer。但 LLaMA-3 用的是 LlamaDecoderLayer，Qwen2 用的是 Qwen2DecoderLayer，MoE 模型里还有 Expert 子模块——全都不在默认白名单里。
结果就是：只有顶层 nn.ModuleList 被 FSDP 包裹，内部所有参数仍在单卡驻留，显存毫无改善。
- 修复方法：显式传入自定义集合，例如
transformer_layer_cls={LlamaDecoderLayer, Qwen2DecoderLayer} - 更鲁棒做法：换用
size_based_auto_wrap_policy，设min_num_params=1e8（即 ≥1 亿参数的模块才分片），彻底绕过类名依赖 - 验证是否生效：训练前跑一遍
sum(p.numel() for p in fsdp_model.parameters())，应 ≈ 总参数量 / world_size；若接近总量，说明分片失败
torch.compile 和 FSDP 的顺序陷阱
PyTorch 2.0+ 强推 torch.compile(model)，但它和 FSDP 有严格时序约束：必须先完成 FSDP 包装，再 compile，且仅支持 mode="default"。
写成 compiled = torch.compile(FullyShardedDataParallel(model)) 是对的；但若写成 FullyShardedDataParallel(torch.compile(model))，或用了 mode="max-autotune"，就会触发 internal assert 或 silent fallback 到解释器模式，性能归零。
- 永远不要在 FSDP 外层再套任何 wrapper（包括
DataParallel、DDP） - checkpoint 加载必须用
FSDP.state_dict_type(...)+FSDP.load_state_dict(...)，直接torch.load后load_state_dict会报 key mismatch - 混合精度若硬件不支持
bfloat16（如 V100），FSDP 会静默降级到fp16，但GradScaler若未同步调整 scale 值，几轮后梯度就全NaN
最易被忽略的点是：FSDP 的分片行为发生在模型构造时，而非数据加载时；所有 buffer 注册、参数初始化、甚至 to(device) 都必须在 FSDP 包装之后做——顺序颠倒，前面所有配置都白搭。
