用途
该技能面向编码智能体，教它如何把 PyTorch FSDP2 正确加入训练循环，包括初始化、分片、混合精度/offload 配置与检查点。FSDP2 在 PyTorch 中主要通过 torch.distributed.fsdp.fully_shard 以及它就地添加到模块上的 FSDPModule 方法暴露。
何时使用
- 模型无法放入单张 GPU（参数 + 梯度 + 优化器状态）。
- 希望采用 eager 模式、基于 DTensor 的逐参数分片，相比 FSDP1 更易检查、分片状态字典更简单。
- 之后可能用 DeviceMesh 把数据并行与张量并行组合起来。
需要谨慎或避免的情况：
- 需要跨 PyTorch 版本严格向后兼容的检查点（DCP 对此有警告）。
- 被迫使用没有 FSDP2 栈的旧版 PyTorch。
替代方案：DistributedDataParallel（DDP）用于经典数据并行；FullyShardedDataParallel（FSDP1）用于原始的参数分片包装。
智能体必须遵守的约定
- 用 torchrun 启动，并按进程设置 CUDA 设备（通常通过 LOCAL_RANK）。
- 自底向上应用 fully_shard()，即先分片子模块（如 Transformer block），再分片根模块。
- 调用 model(input) 而非 model.forward(input)，以便 FSDP2 钩子运行（除非显式 unshard() 或注册 forward 方法）。
- 在分片之后创建优化器，确保它建立在 DTensor 参数（post-fully_shard）之上。
- 使用分布式检查点（DCP）或分布式状态字典辅助函数保存，不要用朴素的 torch.save(model.state_dict())，除非有意收集为完整张量。
使用方式（步骤）
0）版本与环境检查：优先使用近期稳定版 PyTorch；用 torchrun --nproc_per_node <gpus_per_node> 启动，确保 RANK、WORLD_SIZE、LOCAL_RANK 可见。
1）初始化分布式并设置设备：dist.init_process_group(backend="nccl")；torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))；可选创建 DeviceMesh 描述数据并行组。
2）在 meta 设备上构建模型（推荐用于超大模型）：with torch.device("meta"): model = ...；对子模块应用 fully_shard(...)，再 fully_shard(model)；model.to_empty(device="cuda")；model.reset_parameters()（或自定义初始化）。
3）自底向上应用 fully_shard()：不要只对最顶层模块调用；对 transformer 类模型，遍历模块，若 isinstance(m, TransformerBlock) 则 fully_shard(m, ...)，最后 fully_shard(model, ...)。这样可形成参数组以提升集合通信效率，并排除已被先前调用分组的参数，带来更好的重叠与更低的峰值显存。
4）配置 reshard_after_forward 以权衡显存与性能：默认 None 表示非根模块为 True、根模块为 False；显存受限时保持默认或对更多 block 强制 True；吞吐受限且显存允许时可让参数更久保持未分片（根模块常为 False）；高级用法是传入整数，在前向后重分片到更小的 mesh。
5）混合精度与 offload（可选但常见）：使用 mp_policy=MixedPrecisionPolicy(param_dtype=..., reduce_dtype=..., output_dtype=..., cast_forward_inputs=...)；如需 CPU offload 使用 offload_policy=CPUOffloadPolicy()。经验做法：在 H100/A100 级 GPU 上从 BF16 参数/归约开始（若模型数值稳定）；reduce_dtype 与梯度归约预期保持一致；使用 CPU offload 时需为 PCIe/NVLink 流量与运行时开销做预算。
6）优化器、梯度裁剪、累积：在分片后创建优化器以持有 DTensor 参数；如需梯度累积/no_sync，使用 FSDP2 的 set_requires_gradient_sync 而非 FSDP1 的 no_sync()；梯度裁剪采用 FSDP2 教程中针对 DTensor 的做法。
7）检查点：优先 DCP 或分布式状态字典辅助函数。DCP 从多个 rank 并行保存/加载并支持加载时重分片，会生成多个文件（通常每 rank 至少一个）并就地操作。也可用 get_model_state_dict / set_model_state_dict 配合 StateDictOptions(full_state_dict=True, cpu_offload=True, broadcast_from_rank0=True, ...)，优化器用 get_optimizer_state_dict / set_optimizer_state_dict。避免用普通 torch.save 保存 DTensor 状态字典，除非有意用 DTensor.full_tensor() 转换并小心管理显存。
适用场景
- 把 FSDP2 改造进已有训练脚本（Workflow A）：torchrun 启动并初始化进程组；按 LOCAL_RANK 设置 CUDA 设备，需要多维并行时创建 DeviceMesh；构建模型（必要时用 meta），自底向上 fully_shard 后再 fully_shard(model)；分片后创建优化器；用 model(inputs) 触发钩子，用 set_requires_gradient_sync 做累积；通过 torch.distributed.checkpoint 辅助函数加入 DCP 保存/加载。
- 加入 DCP 保存/加载（Workflow B）：用 Stateful 包装状态或通过 get_state_dict 组装；所有 rank 调用 dcp.save(...) 到共享路径；调用 dcp.load(...) 并用 set_state_dict 恢复；加载到不同 mesh 时验证重分片假设。
调试与常见问题
调试清单：所有 rank 是否在不同 GPU 上（否则检查 torch.cuda.set_device(LOCAL_RANK) 与 torchrun 参数）；是否误直接调用 forward()（应使用 model(input) 或显式 unshard()/注册 forward）；fully_shard() 是否自底向上（只分片根模块会导致更差的显存/性能）；优化器创建时机是否正确（必须在分片后基于 DTensor 参数构建）；检查点路径是否一致（用 DCP 时不要与临时 torch.save 混用，除非理解转换）；注意 DCP 的 PyTorch 版本兼容性警告。
常见问题与修复：前向钩子未运行 → 调用 model(inputs)（或显式 unshard()）而非 model.forward(...)；优化器看到非 DTensor 参数 → 在所有 fully_shard 调用之后创建优化器；仅根模块被分片 → 先对子模块自底向上应用 fully_shard；前向后显存尖峰 → 对更多模块设置 reshard_after_forward=True；梯度累积失步 → 使用 set_requires_gradient_sync 而非 FSDP1 的 no_sync()。
最小参考实现结构
智能体应实现带以下标注块的脚本：init_distributed()（初始化进程组、设置设备）；build_model_meta()（meta 上建模型、应用 fully_shard、物化权重）；build_optimizer()（分片后创建优化器）；train_step()（用 model(inputs) 与 DTensor 感知模式做前向/反向/步进）；checkpoint_save/load()（DCP 或分布式状态字典辅助函数）。具体示例见 references/pytorch_examples_fsdp2.md 与官方教程参考。
限制
资料未提供该技能的具体收费、价格、支持平台、性能效果或安装命令等信息，暂未提供。