pytorch-fsdp2

用途

该技能教编码智能体如何在训练循环中正确加入 PyTorch FSDP2。FSDP2 主要通过 torch.distributed.fsdp.fully_shard 以及它为模块就地添加的 FSDPModule 方法暴露能力。技能附带多份参考资料，覆盖 FSDP2 教程、fully_shard API、DDP 与 FSDP1 说明、DeviceMesh、张量并行、Distributed Checkpoint（DCP）概览与用法，以及示例脚本。

何时使用

模型（参数 + 梯度 + 优化器状态）无法放入单张 GPU。
希望采用 eager 模式、基于 DTensor 的逐参数分片，相比 FSDP1 更易检查、分片 state dict 更简单。
后续可能用 DeviceMesh 将数据并行与张量并行组合。

需要谨慎或避免的情况：需要跨 PyTorch 版本严格向后兼容的检查点（DCP 对此有警告）；被迫使用不含 FSDP2 栈的旧版 PyTorch。替代方案包括 DistributedDataParallel（DDP）和 FSDP1。

主要能力与约束

技能规定智能体须遵循的约定：用 torchrun 启动并按进程设置 CUDA 设备（通常依据 LOCAL_RANK）；自下而上应用 fully_shard()（先分片子模块如 Transformer 块，再分片根模块）；调用 model(input) 而非 model.forward(input)，以触发 FSDP2 钩子（除非显式 unshard() 或注册 forward）；在分片之后创建优化器，使其基于 DTensor 参数；检查点使用 DCP 或分布式 state dict 辅助函数，而非直接 torch.save(model.state_dict())，除非有意聚合为完整张量。

使用步骤

环境检查：使用较新的稳定版 PyTorch，用 torchrun --nproc_per_node 启动，确保 RANK、WORLD_SIZE、LOCAL_RANK 可见。
初始化分布式并设置设备：dist.init_process_group(backend="nccl")，torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))，可选创建 DeviceMesh 描述数据并行组。
大模型可在 meta 设备上构建，应用分片后 model.to_empty(device="cuda") 并 reset_parameters()。
自下而上应用 fully_shard()，形成参数组以提升集合通信效率并降低峰值显存。
配置 reshard_after_forward：默认非根模块为 True、根模块为 False；显存受限可多设 True，吞吐优先可延长不分片时间；也可用整数在 forward 后重分片到更小 mesh。
可选混合精度与 offload：MixedPrecisionPolicy(param_dtype、reduce_dtype、output_dtype、cast_forward_inputs) 与 CPUOffloadPolicy()；在 H100/A100 类 GPU 上可从 BF16 起步，使用 CPU offload 需考虑 PCIe/NVLink 流量与运行时开销。
优化器、梯度裁剪与累积：分片后创建优化器；梯度累积用 set_requires_gradient_sync 而非 FSDP1 的 no_sync()；梯度裁剪按教程中 DTensor 方式处理。
检查点：首选 DCP（多 rank 并行保存/加载、支持加载时重分片、产生多个文件、就地操作）；或使用 get_model_state_dict / set_model_state_dict 与 StateDictOptions(full_state_dict、cpu_offload、broadcast_from_rank0 等)，优化器用 get_optimizer_state_dict / set_optimizer_state_dict。避免用普通 torch.save 保存 DTensor state dict。

适用场景与工作流

技能提供两类清单：将 FSDP2 改造进现有训练脚本（初始化进程组、设设备、建模型、自下而上分片、分片后建优化器、用 model(inputs) 触发钩子、加入 DCP 保存/加载）；以及最小化加入 DCP 保存/加载（用 Stateful 包装或 get_state_dict 组装状态，各 rank 调用 dcp.save 到共享路径，dcp.load 并用 set_state_dict 恢复，验证重分片假设）。

调试与常见问题

调试清单包括：各 rank 是否在不同 GPU、是否误调 forward()、fully_shard 是否自下而上、优化器创建时机、检查点路径是否一致，以及 DCP 的版本兼容警告。常见修复：改用 model(inputs) 触发钩子；分片后再建优化器；先分片子模块；对更多模块设 reshard_after_forward=True 以降低显存峰值；用 set_requires_gradient_sync 处理梯度累积。

限制

资料未提供该技能的具体安装命令、收费信息或性能数据；上述内容均为 PyTorch FSDP2 使用指导，实际效果取决于模型、硬件与 PyTorch 版本。