变更摘要

此 PR 为 activation checkpoint 的 checkpoint_exclude_wrapper 新增 save_output 参数，用于控制连续 SAVE 区域是否缓存中间输出。当相邻 SAVE 区域设置 save_output=False 时，通过 deferred input handle（_RecomputedInputHandle）在 saved-tensor pack 中保存输入引用而非完整张量数据，并在重放阶段使用零元素占位符替代实际输出，从而节省显存。同时引入 invocation ID 标记机制避免同一次 checkpoint 调用产生的 SAVE 输出被重复挂载 recompute input handle，并通过 _RecomputeBoundary 可微边界节点保证 RECOMPUTE→SAVE 边界的正确触发。输出结构遍历（_finalize_save_outputs）与 SAVE 来源标记合并在一次递归中完成，减少 Python Host 侧开销。

主要改动

• checkpoint_exclude_wrapper 新增 save_output 参数：在 activation_checkpoint.py、platform.py 及 MindSpore 实现中，为接口增加 save_output: bool = True 关键字参数，默认值保持向前兼容；当设为 False 时不再缓存该区域的真实输出。

• 引入 _RecomputedInputHandle 延迟输入句柄机制：新增 _RecomputedInputHandle 类，在 forward 阶段通过 _pack_saved_tensor 将输入张量替换为句柄存入 saved-tensor pack，重放时由 _materialize_recompute_inputs 按路径绑定重算产生的实际张量，避免保留冗余输入存储。

• _ExcludeCache 重构为 _ExcludeCacheEntry 结构：将原来仅缓存输出的 _CheckpointExcludeCache 改为存储 _ExcludeCacheEntry（包含 output、input_bindings、output_tensor_count），支持 elided 输出场景下重放时构造零元素占位结果。

• _finalize_save_outputs 合并边界与来源标记遍历：新增递归函数对输出中的 Tensor 叶子统一应用 _RecomputeBoundary 边界节点和 _SAVE_OUTPUT_SOURCE_KEY 来源标记，支持 tuple、list、dict 及 namedtuple 等嵌套结构，且在同次遍历中统计 tensor_leaf_count。

• 新增 _RecomputeBoundary 可微边界节点：继承 _Function 的自定义 autograd 节点，通过零元素 requires_grad Tensor 作为 trigger 保存在 ctx 中，保证占位输出不参与求导时仍能触发外层 checkpoint hook。

• invocation ID 标记避免重复挂载 handle：_mark_recompute_inputs 在标记输入时跳过 _SAVE_OUTPUT_SOURCE_KEY 等于当前 invocation_id 的张量，使同次 checkpoint 调用的 SAVE 输出直接传给下游区域而不额外挂载 recompute input handle。

• 新增 ST 与 UT 测试覆盖：新增 checkpoint_exclude_matmul.py（含 20 层 RMSNorm+MatMul+SiLU 显存对比、相邻 SAVE 链测试、Host 性能对比等场景）及 test_checkpoint_exclude_save_chain.py；UT 中新增对 handle 机制、边界节点、输入标记回滚、多输出 elided 重放、非法参数校验等 10 余个测试用例。

代码审查

经过全面审查，此 diff 质量很高，仅发现一处次要问题。

审查总结

按优先级统计

总体风险判断

低风险。此变更为新功能 (save_output=False 优化连续 SAVE 区域的输出显存)，实现质量良好：异常路径有正确的回滚机制（finally 恢复 user data、try/except 回滚 handle），缓存和 replay 路径对称，placeholder 机制保证了重放阶段 autograd 图结构的完整性。测试覆盖充分，包含单元测试（Tensor user data handle 生命周期、边界节点、输入标记回滚、嵌套路径解析、弱引用释放验证、多输出 replay）和系统测试（显存优化量对比、Host 性能回退检测、SAVE→SAVE 链端到端验证）。代码中唯一可改进之处是 _mark_recompute_inputs 中 except BaseException 过于宽泛，属于健壮性方面的可选改进。

各文件审查确认

💬 仅评论

检测到Pull Request代码更新，即将终止当前正在执行的门禁任务（若存在），请基于当前版本的代码重新触发。

CLA Signature Pass

DavidFFFan, thanks for your pull request. All authors of the commits have signed the CLA. 👍