[RFC] Torch non-reentrant checkpoint 主动重计算与 dx/dw 复用设计

背景

MindSpore non-reentrant checkpoint 已具备重计算 handle 收集、backward 前主动重计算、稳定 session 缓存以及跨 dx/dw 复用重计算结果的能力。Torch 原生 use_reentrant=False checkpoint 只在 saved tensor 首次 unpack 时按 GraphTask 懒触发重计算，不公开 checkpoint frame，也不支持由上层调度器主动触发和跨 GraphTask 复用。

在 dx/dw 分离场景中，输入梯度和权重梯度由不同 autograd GraphTask 计算。如果直接使用 Torch 原生实现，两次 backward 会分别执行 checkpoint replay，增加计算开销。NPU autograd 还可能在独立 worker 线程执行 saved-tensor unpack，主线程设置的 ContextVar 不会自动传播，不能单纯依赖动态上下文识别 session。

本 RFC 用于说明 Torch 后端补齐上述底层原语的设计。当前只交付 checkpoint 内核与 session 生命周期能力，PP stage 和调度器接入后续单独设计。

---

1. 基本信息

项目	内容
作者	@DavidFFFan
相关模块	core/activation_checkpoint、platform/torch/activation_checkpoint
相关 issue / PR	PR #1102
适用后端	PT；MS 作为已有能力参考

2. 背景

类型	需要说明的内容
功能补全	Torch 缺少主动重计算、稳定 session、跨 dx/dw GraphTask 复用和显式清理能力
能力增强	在 Torch 2.9 use_reentrant=False 语义上扩展调度基础原语，不修改或 patch 框架源码
用户需求	上层调度器可提前完成 checkpoint replay，dx/dw 分离计算时只重计算一次

本 RFC 要解决的问题：Torch checkpoint 无法主动重计算，也无法让 dx/dw 两个 GraphTask 复用同一份重计算结果。
完成后的成功标准：prefire、dx 和 dw 使用同一 session 时 checkpoint function 只执行一次 replay，梯度与无 checkpoint 基线一致，session 可可靠清理。

3. 目标和非目标

3.1 目标

1. Torch eager 模式提供与原生 non-reentrant checkpoint 一致的 forward/replay 基础语义。
2. 提供 handle collector、主动重计算、稳定 session、retain_on_unpack 和幂等清理原语。
3. 支持 dx/dw 独立 GraphTask 复用一次 prefired replay，避免重复计算。
4. NPU autograd worker 不继承主线程 ContextVar 时仍能识别 frame 对应的 session。
5. saved-tensor unpack 热路径不执行全局锁和 session ContextVar 查询。
6. 支持 per-call early_stop、RNG、device/autocast 恢复、确定性检查、kwargs 和 context_fn 组合。
7. 普通 GraphTask 路径保持 Torch 原生的多次 backward 和 nested checkpoint 行为。

3.2 非目标

1. 本期不接入 PP stage 或 dx/dw 调度器；PR #1102 只提供底层原语。
2. 本期不支持 scheduled recomputation 中动态创建 nested checkpoint，遇到时明确报错。
3. 本期不支持 reentrant checkpoint 和 checkpoint_sequential。
4. 本期不重写 torch.compile 内部 checkpoint，compile 模式回退 Torch 原生实现。
5. Torch checkpoint_exclude_wrapper 作为独立能力后续适配；当前 invocation state/session 方案与其兼容。
6. 不在 CI 中增加 Torch 2.6/2.7/2.9 完整版本矩阵，由发布环境做版本验证。

4. 相关实现参考

来源	做法	限制	对本 RFC 的影响
Torch 2.9 use_reentrant=False	forward pack 保存 holder，首次 unpack 时按 GraphTask replay	frame 为私有实现，不能主动触发或跨 GraphTask 复用	Hyper eager 实现保留算法语义并增加调度扩展
MindSpore 2.10 non-reentrant recompute	提供 handle、prefire、session retain/clear	生命周期依赖 MindSpore autodiff 执行模型	Torch 对齐功能语义，但按 Torch autograd worker 特性实现
Torch GraphTask id	普通 backward 用 GraphTask id 隔离 replay cache	私有 C++ 接口	无显式 session 时保持原生 fallback 语义

5. 对外接口

5.1 接口定义

with platform.recompute_handle_collector_ctx() as handles:
    output = checkpoint(function, *args, early_stop=True, **kwargs)

session_id = ("micro_batch", micro_batch_id)
try:
    with platform.recompute_session_ctx(session_id, retain_on_unpack=True):
        for handle in handles:
            platform.recompute_handle(handle, session_id)

    with platform.recompute_session_ctx(session_id, retain_on_unpack=True):
        dx = compute_input_grad(output)

    with platform.recompute_session_ctx(session_id, retain_on_unpack=False):
        dw = compute_weight_grad(output)
finally:
    platform.clear_recompute_session(session_id)

入参 / 配置项	类型	默认值	是否必填	含义	合法范围	错误处理
early_stop	bool	True	否	saved tensor 齐备后是否提前结束 replay	True / False	非 bool 抛 ValueError
session_id	hashable	无	是	跨 GraphTask 稳定标识	非 None 且可哈希	非法值抛 ValueError
retain_on_unpack	bool	False	否	unpack 后是否保留 session tensor	True / False	非 bool 抛 ValueError
handle	opaque	无	是	collector 返回的 checkpoint frame 句柄	仅接受 Hyper handle	非法 handle 抛 ValueError

5.2 使用示例

上层使用顺序固定为：收集 handle → prefire → dx 保留 → dw 最终消费 → finally 清理。即使最终消费者已完成，也必须执行 clear_recompute_session()，释放 partial backward 或未消费 holder 关联的数据。

5.3 接口说明

接口与 MindSpore 已有平台抽象保持一致。session_id 必须由调用方显式提供，避免匿名 id 无法传递到后续 dx/dw。handle 对用户保持不透明，上层不依赖 Torch checkpoint 私有 frame 类型。

6. 方案设计

6.1 总体流程

flowchart TD
    A["Checkpoint forward"] --> B["Collector records frame handle"]
    B --> C["Prefire with stable session"]
    C --> D["Replay and cache saved tensors on frame"]
    D --> E["dx GraphTask: retain_on_unpack=True"]
    E --> F["dw GraphTask: retain_on_unpack=False"]
    F --> G["clear_recompute_session"]
    D -. "no session" .-> H["GraphTask id fallback"]

6.2 架构参考

flowchart LR
    subgraph Core["Core API"]
        Checkpoint["checkpoint / checkpoint_wrapper"]
        State["Recompute invocation context"]
    end

    subgraph TorchBackend["Torch Backend"]
        Frame["CheckpointFrame"]
        Hooks["saved_tensors_hooks"]
        Session["Session control plane"]
        Cache["Per-session recomputed tensors"]
    end

    subgraph Consumers["Consumers"]
        Prefire["Prefire scheduler"]
        DX["dx GraphTask"]
        DW["dw GraphTask"]
    end

    Checkpoint --> Frame
    State --> Hooks
    Frame --> Hooks
    Session --> Frame
    Frame --> Cache
    Prefire --> Session
    DX --> Cache
    DW --> Cache

6.3 时序参考

sequenceDiagram
    participant U as Upper Scheduler
    participant C as Checkpoint Frame
    participant W as NPU Autograd Worker
    participant D as dx/dw Consumer

    U->>C: collect handle during forward
    U->>C: recompute_handle(handle, session)
    C->>C: bind SessionActivation to frame
    C->>C: replay and cache saved tensors
    D->>W: launch dx GraphTask
    W->>C: unpack reads frame.active_session
    C-->>W: retained recomputed tensor
    D->>W: launch dw GraphTask
    W->>C: unpack reads same session
    C-->>W: final recomputed tensor
    U->>C: clear_recompute_session(session)

6.4 关键逻辑

• _CheckpointFrame 管理 forward holder、per-session replay tensor、metadata 和 active_session。
• session 注册、frame 绑定/解绑和清理由控制面 RLock 保护。
• unpack 热路径直接读取 frame.active_session；没有 activation 时使用当前 GraphTask id。
• prefire 在实际执行 replay 的线程临时安装 session ContextVar，用于生命周期校验和 scheduled nested 检测，不依赖线程传播。
• retain_on_unpack=True 允许多个消费者复用；最终消费者设为 False，finally 继续执行幂等 clear。
• replay 使用 forward 捕获的 CPU/device RNG、autocast 和用户 recompute context。
• early_stop=True 在所有 forward holder 对应 tensor 已产生后抛内部控制流异常结束 replay。

6.5 代码改动点

模块	改动内容	是否影响已有行为
core/activation_checkpoint	统一 early_stop 和 recompute context 组合	是，Torch eager 走 Hyper 实现
platform/platform.py	补充 handle/session/clear 平台抽象	否，保持后端隔离
platform/torch/activation_checkpoint	新增 non-reentrant 内核和 scheduled recomputation	是，Torch eager checkpoint 实现切换
platform/torch/platform.py	路由 checkpoint 和调度原语	是
platform/mindspore/platform.py	对齐平台接口签名	否，不改变 MindSpore 内核行为
tests	CPU UT 与单卡 NPU ST	否

6.6 方案取舍

方案	优点	缺点	是否选择	原因
unpack 直接读取 session ContextVar	实现简单	NPU worker 不继承主线程上下文；每次 unpack 都查询	否	实际 dx/dw E2E 无法识别 session
unpack 查询全局 session-frame 映射	可跨线程	热路径需要查表或加锁	否	saved-tensor unpack 调用频繁
控制面将 activation 绑定到 frame	worker 可见；unpack 只读字段	同一 frame 不支持并发 session	是	满足有序 dx/dw 生命周期且热路径开销最低

该方案的主要代价是 HyperParallel 需要维护一份 Torch eager non-reentrant checkpoint 内核，并在 Torch 升级时持续对照原生实现。

7. 组件依赖

依赖组件	强依赖 / 弱依赖	当前状态	未 ready 时本期能力
FSDP	不涉及	不涉及	不受影响
TP	不涉及	不涉及	不受影响
PP	弱依赖	stage 尚未接入	底层 handle/session 原语可独立交付和测试
checkpoint	强依赖	Torch eager 内核由本 RFC 补齐	无 checkpoint 时功能不生效
optimizer	不涉及	不涉及	不受影响
PT / MS 后端	PT 强依赖，MS 参考	MS 已有类似能力	Torch 独立实现，MS 行为不变

完整能力需要：后续 PP stage 将 micro-batch session、prefire、dx、dw 和 clear 接入调度生命周期。
本期最小可交付能力：单点 API 能收集 handle、主动重计算，并让独立 dx/dw GraphTask 复用一次 replay。

8. 约束与兼容性

类型	内容
不支持项	reentrant、scheduled nested、同一 frame 并发 session、完整 compile/context_fn 组合
性能收益	dx/dw 分离由两次 replay 降为一次；本期不设端到端吞吐指标
显存收益	保持 activation checkpoint 基础收益；session retain 会延长 prefired tensor 生命周期
性能劣化	控制面注册/绑定有锁；unpack 热路径仅字段读取，无锁和 ContextVar 查询
PT / MS 差异	对外 handle/session 语义对齐；Torch 使用 GraphTask、saved tensor hook 和 frame-bound activation
和已有行为不一致	Torch eager 不继承原生全局 set_checkpoint_early_stop，使用 Hyper per-call early_stop；compile 回退原生

约束：调用方必须遵循“prefire 完成 → dx/dw 有序消费 → finally clear”，不得在同一 frame 上并发激活不同 session。

9. 验证设计

9.1 测试范围

测试维度	覆盖范围
后端	MindSpore、Torch
执行模式	eager 模式
调度方式	独立调度（不依赖 PP stage）
梯度计算场景	dx/dw 分离：输入梯度与权重梯度分别计算，并复用同一份主动重计算结果

9.2 用例分层

用例级别	数量	覆盖内容	通过标准
UT	新增生命周期用例，相关回归 93 项	参数、原生语义、multi-frame、重复 session、early-stop、partial、异常、SAC、清理	全部通过，数值与无 checkpoint 基线一致
Level0 单卡 NPU	3 个 scheduled 核心用例	真实 worker dx/dw、NPU dropout RNG、NPU bfloat16 autocast	function 只 replay 一次，梯度/RNG/dtype 一致
Level1	0	PP stage 和端到端调度	后续 stage PR 增加

9.3 交互验证

组合	是否验证	通过标准
scheduled recompute + dx/dw	是	prefire、dx、dw 共用一次 replay，dx/dw 与基线一致
scheduled recompute + SAC	是，CPU UT	SAC cache 只在 prefire 消费一次
scheduled recompute + RNG	是，NPU ST	dropout mask/梯度一致，prefire 不推进调用方 RNG state
scheduled recompute + autocast	是，NPU ST	prefire 在 autocast 外运行仍恢复 bfloat16 配置
scheduled recompute + PP stage	否	后续调度接入验证
PT / MS 对齐	接口和能力语义对照	平台 API 一致，后端实现差异符合本 RFC

9.4 性能 / 显存验证

场景	基线	开启本特性	指标	通过标准
dx/dw 分离	两个 GraphTask 各 replay 一次	prefire replay 一次，dx/dw 复用	checkpoint function 调用次数	forward + replay 共 2 次
unpack 热路径	ContextVar/全局映射方案	frame 字段读取	Python 控制路径	不查询 ContextVar，不获取全局锁

本期以功能和调用次数验收，不新增性能门禁；stage 接入后补充端到端 step time 和峰值显存数据。

10. 实现计划

PR	内容	依赖	验证
#1102	Torch eager non-reentrant 内核、handle/session 原语、dx/dw E2E	无	CPU UT + 单卡 NPU Level0
后续 PR	Torch checkpoint_exclude_wrapper 适配	#1102 invocation/recompute state	UT + NPU ST
后续 PR	PP stage 生命周期接入 prefire/dx/dw/clear	#1102	PP 调度 ST + 性能/显存
后续 PR	scheduled nested checkpoint 依赖建模（如有明确需求）	stage 设计	专项 UT/ST

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

Read the RFC's listed modules: core/activation_checkpoint, platform/platform.py, platform/torch/activation_checkpoint, and the platform files, then review PR #1102 and the stated CPU UT and single-card NPU Level0 validation. The described completion criteria are handle collection, prefire, dx/dw replay reuse, reliable session cleanup, and matching gradients and RNG behavior; PP integration is explicitly deferred.