【需求】PyTorch 分支支持 dx/dw 分离场景的重计算和重算预取

背景

当前 HyperParallel PyTorch 分支还不支持 dx/dw 分离场景下的重计算调度，也缺少重算预取能力。在 pipeline F&B fusion、Interleaved 1F1B、fully_shard/HSDP 等场景中，dx 和 dw 可能被拆成独立 backward 阶段或独立调度单元。如果每个阶段各自触发 checkpoint recompute，会导致同一段 activation 被重复重算，影响性能；如果希望提前调度 recompute，也缺少可收集、可手动触发的 recompute handle。

因此需要在 PyTorch 分支补齐 dx/dw 分离场景的重计算支持，并支持重算预取，使上层 scheduler 可以控制 recompute 的触发时机和缓存生命周期。

需求目标

1. 支持 dx/dw 分离场景的重计算

   • dx 和 dw 分离执行时，可以共享同一轮 recompute 结果。
   • 避免同一 checkpoint block 在 dx、dw 阶段重复重算。
   • 支持多个 backward 调用显式归属到同一个 recompute session。

2. 支持重算预取

   • forward 创建 checkpoint/recompute 单元时，上层可以收集对应 handle。
   • scheduler 可以在 backward 真正 unpack activation 前主动触发 recompute。
   • 预取产生的 recomputed activation 可以被后续 dx/dw 阶段消费。

3. 支持缓存生命周期管理

   • 支持 retain_on_unpack 或等价机制：第一次 consumer 读取 recomputed activation 后可选择保留，供后续 dw/consumer 继续使用。
   • 支持显式 clear：dx/dw 或多 consumer 完成后，清理 session 关联的 recomputed tensors、counter、is_recomputed 状态和 holder handle。
   • 默认未启用该能力时，保持 PyTorch checkpoint 当前行为不变。

PyTorch 现状说明

PyTorch 原生 torch.utils.checkpoint 已有 GraphExecGroup，可以让同一 group 下的 checkpoint region 在多次 backward 中按同一个 gid 记录重算状态；但原生实现会在每个 saved activation unpack 后立即清掉 holder handle，并要求同一 GraphExecGroup 下的多次 backward 不能访问同一个 saved activation。

因此，PyTorch 原生能力不能直接满足 dx/dw 分离共享同一份 recompute activation，也不能提供 HyperParallel scheduler 所需的重算预取 handle。

实现建议

优先在 HyperParallel PyTorch 平台层封装/扩展，形成可由 scheduler 使用的 recompute session 与 prefetch API：

• 可放在 hyper_parallel/platform/torch/activation_checkpoint/ 或 hyper_parallel/core/activation_checkpoint/。
• 参考 PyTorch 非重入 checkpoint 的 _CheckpointFrame、_checkpoint_hook、_recomputation_hook 设计。
• 参考 PyTorch GraphExecGroup 的 gid 思路，但需要补齐 dx/dw 共享所需的 retain/clear 语义。
• 提供 handle collector，使 scheduler 能够收集 checkpoint/recompute handle 并提前调用 recompute。
• 与 pipeline / fully_shard / HSDP dx-dw split 调度集成时，保证同一 session 下只重算一次，后续阶段消费缓存。

验收标准

• 新增 PyTorch 分支单测：同一 checkpoint block 在 dx/dw 两次 backward 中只触发一次 recompute。
• 覆盖第一次 backward 读取后保留缓存、第二次 backward 复用并最终清理的场景。
• 覆盖 scheduler 手动收集 recompute handle 并提前触发 recompute 的预取场景。
• 覆盖异常、未完整消费或提前退出时显式 clear 不遗留 session 状态。
• 未启用 dx/dw recompute session 或 prefetch 时，现有 checkpoint 行为保持兼容。

关联

• HyperParallel #170：Interleaved 1F1B 支持 F&B fusion 阶段 dx/dw 分离。
• MindSpore PR mindspore/mindspore#92629：非重入 recompute 支持自定义 session id、retain unpack、handle 收集，用于支撑 dx/dw 分离只重算一次和重算预取。

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 186
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/186

Start by reading the activation-checkpoint code under hyper_parallel/platform/torch/activation_checkpoint/ and hyper_parallel/core/activation_checkpoint/, then compare PyTorch’s non-reentrant checkpoint hooks and GraphExecGroup. Done means dx/dw backward stages share one recompute, prefetch handles and retain/clear lifecycle are covered by tests, and existing behavior remains compatible when the feature is disabled.