pytorch-fsdp2

技能：在训练脚本中正确使用 PyTorch FSDP2（fully_shard）

此技能将指导编码 Agent 如何在训练循环中添加 PyTorch FSDP2，并正确处理初始化、分片、混合精度／卸载配置和检查点。

PyTorch 中的 FSDP2 主要通过 torch.distributed.fsdp.fully_shard 以及它原地添加到模块上的 FSDPModule 方法提供。参见：references/pytorch_fully_shard_api.md、references/pytorch_fsdp2_tutorial.md。

---

何时使用此技能

以下情况适合使用 FSDP2：

• 模型无法放入单个 GPU（参数 + 梯度 + 优化器状态）。
• 你希望采用 eager 模式的分片方案，即基于 DTensor 的逐参数分片（比 FSDP1 更便于检查，分片状态字典也更简单）。
• 你之后可能会使用 DeviceMesh 将 DP 与张量并行组合起来。

以下情况应避免使用或谨慎使用：

• 你需要在不同 PyTorch 版本间严格向后兼容的检查点（DCP 对此发出警告）。
• 你只能使用不支持 FSDP2 技术栈的旧版 PyTorch。

替代方案（FSDP2 并非最佳选择时）

• DistributedDataParallel（DDP）：如果你需要经典的分布式数据并行训练，请使用标准的数据并行封装器。
• FullyShardedDataParallel（FSDP1）：如果你需要在数据并行工作进程间分片参数，请使用原始 FSDP 封装器。

参考：references/pytorch_ddp_notes.md、references/pytorch_fsdp1_api.md。

---

Agent 必须遵守的约定

1. 使用 torchrun 启动，并为每个进程设置 CUDA 设备（通常通过 LOCAL_RANK）。
2. 自底向上应用 fully_shard()，也就是先对子模块（例如 Transformer 块）分片，再对根模块分片。
3. 调用 model(input)，而不是 model.forward(input)，以确保 FSDP2 钩子能够运行（除非你显式调用 unshard() 或注册 forward 方法）。
4. 分片后再创建优化器，并确保优化器基于 DTensor 参数构建（即在 fully_shard 之后）。
5. 使用 Distributed Checkpoint（DCP）或分布式状态字典辅助工具创建检查点，不要直接使用 torch.save(model.state_dict())，除非你明确选择汇集完整张量。

（这些规则均在官方 API 文档／教程中有直接说明；参见参考资料。）

---

分步操作流程

0）版本与环境检查

• 优先使用近期稳定版 PyTorch，其中的文档已更新 FSDP2 和 DCP 相关内容。
• 使用 torchrun --nproc_per_node <gpus_per_node> ...，并确保可以访问 RANK、WORLD_SIZE、LOCAL_RANK。

参考：references/pytorch_fsdp2_tutorial.md（启动命令和设置）、references/pytorch_fully_shard_api.md（用户约定）。

---

1）初始化分布式并设置设备

最简且正确的模式：

• dist.init_process_group(backend="nccl")
• torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
• 也可以创建 DeviceMesh，用于描述数据并行组。

参考：references/pytorch_device_mesh_tutorial.md（DeviceMesh 的用途及其管理进程组的方式）。

---

2）在 meta 设备上构建模型（超大模型推荐）

对于大型模型，先在 meta 上初始化并应用分片，再将权重实体化到 GPU：

• with torch.device("meta"): model = ...
• 对子模块应用 fully_shard(...)，然后对 model 应用 fully_shard(model)
• model.to_empty(device="cuda")
• model.reset_parameters()（或使用你的初始化流程）

参考：references/pytorch_fsdp2_tutorial.md（迁移指南明确展示了此流程）。

---

3) 自底向上应用 fully_shard()（包装策略 =“按需应用”）

不要只对最顶层模块调用 fully_shard。

Transformer 类模型的推荐分片模式：

• 遍历模块，if isinstance(m, TransformerBlock): fully_shard(m, ...)
• 然后调用 fully_shard(model, ...)

原因：

• fully_shard 会为集体通信效率创建“参数组”，并排除已通过先前调用分组的参数。自底向上应用可实现更好的重叠并降低峰值内存。

参考：references/pytorch_fully_shard_api.md（自底向上的要求及原因）。

---

4) 配置 reshard_after_forward，权衡内存与性能

默认行为：

• None 表示非根模块为 True，根模块为 False（良好的默认设置）。

经验法则：

• 如果受内存限制：保留默认值，或对多个块强制设为 True。
• 如果受吞吐量限制且内存充足：可考虑延长参数保持未分片状态的时间（根模块通常设为 False）。
• 进阶：使用 int，在前向传播后重新分片到更小的网格（例如节点内网格），前提是该网格大小是有意义的约数。

参考：references/pytorch_fully_shard_api.md（完整语义）。

---

5) 混合精度与卸载（可选但常用）

FSDP2 使用：

• mp_policy=MixedPrecisionPolicy(param_dtype=..., reduce_dtype=..., output_dtype=..., cast_forward_inputs=...)
• 如果需要 CPU 卸载，则使用 offload_policy=CPUOffloadPolicy()

经验法则：

• 在 H100/A100 级别的 GPU 上，先尝试使用 BF16 参数和归约（前提是模型在数值上保持稳定）。
• 让 reduce_dtype 符合你的梯度归约预期。
• 如果使用 CPU 卸载，请预留 PCIe/NVLink 流量和运行时开销。

参考：references/pytorch_fully_shard_api.md（MixedPrecisionPolicy / OffloadPolicy 类）。

---

6) 优化器、梯度裁剪、累积

• 在分片后创建优化器，使其持有 DTensor 参数。
• 如果需要梯度累积 / no_sync：
  • 使用 FSDP2 机制（set_requires_gradient_sync），而不是 FSDP1 的 no_sync()。

梯度裁剪：

• 使用 FSDP2 教程（“Gradient Clipping and Optimizer with DTensor”）中展示的方法，因为参数和梯度都是 DTensor。

参考：references/pytorch_fsdp2_tutorial.md。

---

7) 检查点：优先使用 DCP 或分布式状态字典辅助工具

推荐以下两种方法：

A) Distributed Checkpoint（DCP）——最佳默认选择

• DCP 可从多个 rank 并行保存 / 加载，并支持加载时重新分片。
• DCP 会生成多个文件（通常每个 rank 至少一个），并以“原位”方式操作。

B) 分布式状态字典辅助工具

• 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

避免：

• 除非你有意使用 DTensor.full_tensor() 进行转换并谨慎管理内存，否则不要使用普通的 torch.save 保存 DTensor 状态字典。

参考：

• references/pytorch_dcp_overview.md（DCP 的行为和注意事项）
• references/pytorch_dcp_recipe.md 和 references/pytorch_dcp_async_recipe.md（端到端用法）
• references/pytorch_fsdp2_tutorial.md（DTensor 与 DCP 状态字典流程）
• references/pytorch_examples_fsdp2.md（可运行的检查点脚本）

---

工作流检查清单（便于复制粘贴）

工作流 A：将 FSDP2 集成到现有训练脚本中

• 使用 torchrun 启动并初始化进程组。
• 根据 LOCAL_RANK 设置 CUDA 设备；如果需要多维并行，则创建 DeviceMesh。
• 构建模型（如有需要，使用 meta），从底向上应用 fully_shard，然后调用 fully_shard(model)。
• 分片后再创建优化器，以便它能获取 DTensor 参数。
• 使用 model(inputs) 运行模型，以触发钩子；累积时使用 set_requires_gradient_sync。
• 通过 torch.distributed.checkpoint 辅助函数添加 DCP 保存/加载。

参考：references/pytorch_fsdp2_tutorial.md、references/pytorch_fully_shard_api.md、references/pytorch_device_mesh_tutorial.md、references/pytorch_dcp_recipe.md。

工作流 B：添加 DCP 保存/加载（最简模式）

• 使用 Stateful 封装状态，或通过 get_state_dict 组装状态。
• 从所有 rank 调用 dcp.save(...)，保存到共享路径。
• 调用 dcp.load(...) 并使用 set_state_dict 恢复状态。
• 加载到不同 mesh 时，验证所有重新分片假设。

参考：references/pytorch_dcp_recipe.md。

调试检查清单（Agent 应首先检查的内容）

1. 所有 rank 是否都使用不同的 GPU？
   如果不是，请检查 torch.cuda.set_device(LOCAL_RANK) 和 torchrun 标志。
2. 是否意外地直接调用了 forward()？
   请使用 model(input)，或显式调用 unshard()／注册 forward。
3. 是否从底向上应用了 fully_shard()？
   如果只对根模块进行分片，预期会出现更差的内存/性能表现，也可能造成混淆。
4. 是否在正确的时机创建了优化器？
   必须在分片后基于 DTensor 参数创建。
5. 检查点路径是否一致？
   • 如果使用 DCP，除非了解转换方式，否则不要与临时的 torch.save 混用。
   • 注意 PyTorch 版本兼容性警告对 DCP 的影响。

---

常见问题和修复方法

• 前向钩子未运行 → 使用 model(inputs)（或显式调用 unshard()），不要使用 model.forward(...)。
• 优化器获取到的不是 DTensor 参数 → 在所有 fully_shard 调用后创建优化器。
• 只有根模块进行了分片 → 在根模块之前，对子模块从底向上应用 fully_shard。
• 前向传播后内存激增 → 对更多模块设置 reshard_after_forward=True。
• 梯度累积不同步 → 使用 set_requires_gradient_sync，不要使用 FSDP1 的 no_sync()。

参考：references/pytorch_fully_shard_api.md、references/pytorch_fsdp2_tutorial.md。

---

最简参考实现大纲（便于 Agent 使用）

编码 Agent 应实现一个包含以下标记代码块的脚本：

• init_distributed()：初始化进程组并设置设备
• build_model_meta()：在 meta 上构建模型，应用 fully_shard，再实例化权重
• build_optimizer()：在分片后创建优化器
• train_step()：使用 model(inputs) 和适配 DTensor 的模式进行前向传播、反向传播和优化器更新
• checkpoint_save/load()：使用 DCP 或分布式状态字典辅助函数

具体示例见 references/pytorch_examples_fsdp2.md 和官方教程参考资料。

---

参考资料

• references/pytorch_fsdp2_tutorial.md
• references/pytorch_fully_shard_api.md
• references/pytorch_ddp_notes.md
• references/pytorch_fsdp1_api.md
• references/pytorch_device_mesh_tutorial.md
• references/pytorch_tp_tutorial.md
• references/pytorch_dcp_overview.md
• references/pytorch_dcp_recipe.md
• references/pytorch_dcp_async_recipe.md
• references/pytorch_examples_fsdp2.md
• references/torchtitan_fsdp_notes.md（可选，生产环境说明）
• references/ray_train_fsdp2_example.md（可选，集成示例）