A->B->C->D->E

MoE
dispatch

combine

all2all       不需要激活
compute   recompute 不保存激活
gmm     需要激活，激活是输入，还需要保存输出


1. 单算子exclude
------
all2all        exclude，保存输出， 1. 保存all2all输出
compute   recompute 不保存激活， 需要重算
gmm         exclude，gmm输入如果是激活也会保存，保存输出。输入优化点。 2. 保存输入（当前未优化，需要优化），3. 保存输出。
------
整个区域checkpoint，单算子exclude
会保存三份数据，1. 不是激活值，理论上可以不保存

2. 整个区域exclude，不重算
------
all2all
compute
gmm           1. 保存输入 2. 保存输出
------
会保存两份数据，都是激活值

3. 整个区域exclude，不重算
------
all2all          exclude
compute    exclude
gmm          exclude
------
保存3份输出。需要优化，做到连续的checkpoint_exclude，只保存一份输出。


假设我不想重算B和D
B的输出也不是激活
C算子的输出不是激活值，C没有激活值
D算子也没有激活值
只有E有激活值

静态图
Cell.recompute()
B.recompute(False)
D.recompute(False)


动态图
checkpopint_wrapper(Cell)

exclude_wrapper(B)
exclude_wrapper(D)

保存B和D的输出

MindSpore recompute -> exclude 输入激活优化设计

1. 背景

HyperParallel 的 MindSpore 后端已经支持 checkpoint_exclude_wrapper。在非重入 checkpoint 的首次前向阶段，
exclude 区域正常执行，并通过内层 saved_tensors_hooks 保存反向所需 tensor；checkpoint 重算时，exclude 区域不再
执行，而是直接返回首次前向缓存的输出。

当 exclude 区域的输入来自上游 recompute 区域时，exclude 内部算子可能为 backward 保存该输入。当前 pack hook
返回 tensor.data，它仍与原输入共享 device storage，导致上游本可释放、并在 backward 前重算恢复的输入激活从
首次前向一直存活到 backward，抵消了部分 checkpoint 显存收益。

目标场景如下：

checkpoint forward:
  recompute region -> activation A -> exclude region -> output B
                                      saves A for backward

本设计只解决 recompute -> exclude 输入保存问题。连续 exclude 输出融合、save_output 参数和 replay placeholder
不属于本次范围。

2. 设计目标

• exclude 为 backward 保存其可重算输入时，不保留首次前向输入的 device storage。
• checkpoint 重算经过 exclude 边界时，用重算生成的新输入固化 saved-tensor handle。
• pack hook 热路径只进行一次 Tensor 属性读取和一次轻量判断。
• 不进行 storage、layout、version 或 view 关系分析。
• 不增加公开 wrapper、公开参数或新的声明式概念。
• 保持现有数值、梯度、重复调用和流水线 dx/dw 拆分语义。

3. 非目标

• 不优化 exclude 保存输入 view、slice、reshape 结果的场景。
• 不处理 exclude -> exclude -> exclude 输出缓存融合。
• 不自动分析 Tensor 的生产算子或完整消费者关系。
• 不改变 PyTorch 后端。
• 不支持 exclude 内部原地修改输入后再保存该输入的场景。
• 不支持 checkpoint_exclude_wrapper 嵌套使用。

如果 saved-tensor pack 收到的不是带标记的原始输入对象，例如输入的 view，优化不会命中，pack hook 将回退到现有
tensor.data 行为。该回退只影响显存收益，不影响正确性。

4. 核心思路

exclude wrapper 在执行 body 前，通过 MindSpore Tensor 的 user-data 接口为输入激活临时增加属性。属性值不是布尔
标志，而是该输入对应的特殊 handle：

input tensor
  user_data[hyper_parallel_recompute_input_handle] -> RecomputedInputHandle

exclude 内层 pack hook 收到 tensor 后：

handle = tensor._get_user_data(RECOMPUTE_INPUT_HANDLE_KEY)
if isinstance(handle, RecomputedInputHandle):
    return handle
return tensor.data

首次前向中，autograd 保存的是不包含 device tensor 的轻量 handle，因此不会由该 saved-tensor slot 持有输入激活。

checkpoint 重算再次调用同一个 exclude wrapper 时，wrapper 不执行 body，但会按首次前向记录的输入路径找到当前
重算输入，并调用：

handle.materialize(recomputed_input.data)

backward 触发 unpack 时，unpack hook 从 handle 中取得已经固化的重算 tensor。

5. 内部数据结构

5.1 RecomputedInputHandle

RecomputedInputHandle
  tensor: Optional[Tensor]
  used: bool

  mark_used()
  materialize(tensor)
  unpack() -> Tensor

生命周期：

1. 首次前向进入 exclude 时创建，此时不持有 Tensor。
2. exclude pack hook 调用 mark_used()，再将 handle 作为 opaque packed object 返回给 autograd。
3. checkpoint 重算到达 exclude 时，handle 绑定重算输入的 tensor.data。
4. backward unpack 返回该 tensor。
5. autograd 释放 packed object 后，handle 和重算 tensor 一起释放。

如果 unpack() 发生在 materialize() 前，应抛出明确的 RuntimeError，不能回退到首次前向 tensor，否则会掩盖
重算顺序或控制流错误。

handle 允许再次 materialize()，以兼容 retain graph 或 PP dx/dw 场景下的重复 replay；后一次 replay 可以覆盖前一次
绑定值，但不能在第一次 unpack 后主动清空。

5.2 InputBinding

InputBinding
  path: input path in args/kwargs
  handle: RecomputedInputHandle

path 用于在重算阶段从本次调用的 args、kwargs 中找到对应 Tensor。支持直接参数以及 tuple、list、dict 中的
Tensor 叶子。

同一 Tensor 对象在输入结构中出现多次时，按对象 identity 去重，只创建一个 handle，并使用第一次出现的路径进行
重算绑定。

5.3 CheckpointExcludeCallRecord

现有 cache 条目由单一 output 扩展为调用记录：

CheckpointExcludeCallRecord
  output: Any
  input_bindings: list[InputBinding]

_ExcludeCache 仍按 wrapper identity 保存 FIFO 队列，以匹配同一 wrapper 在一次 checkpoint 中的多次调用。
本次设计不改变 output 的缓存方式。

6. 首次前向流程

6.1 输入筛选

exclude wrapper 进入 body 前遍历 args 和 kwargs：

1. 只处理 MindSpore Tensor。
2. 使用 isinstance(tensor, mindspore.Parameter) 跳过 Parameter。
3. 对其余输入激活创建 handle 和 binding。

checkpoint 原始输入、普通 recompute 算子输出、上游 exclude 的缓存输出都可以使用同一机制：只要重算阶段再次调用
当前 exclude wrapper，并能在同一路径取得真实 Tensor，handle 就可以被固化。对于上游 exclude 输出，该机制不会
释放上游 wrapper 自己缓存的 output，但仍保持正确性；连续 exclude 输出缓存优化留给独立方案处理。

6.2 临时属性

对每个待优化输入：

1. 通过 _get_user_data 保存该 Tensor 上同名 user-data 的原值；不存在时返回 None。
2. 通过 _set_user_data 将 RecomputedInputHandle 写入 Tensor user-data。
3. 在 saved_tensors_hooks 中执行 exclude body。
4. 在 finally 中恢复原属性或删除临时属性。

该 user-data 只在 exclude body 执行期间存在，避免下游普通算子或后续 training step 误识别。handle 本身是当前调用
独占对象，不需要额外通过全局表查找。

MindSpore 2.10 的验证结果表明，saved-tensor pack 会收到新的 Python Tensor wrapper，普通 setattr 写入的 Python
属性不会保留；Tensor _set_user_data 写入的数据附着在底层 Tensor 上，可以被 pack wrapper 通过 _get_user_data
读取。因此实现必须使用 Tensor user-data，不能使用普通 Python 对象属性。

6.3 pack 行为

pack hook 不遍历输入，不比较 storage，也不识别 view：

pack(tensor)
  -> tensor 上存在有效 handle：标记 handle.used，返回 handle
  -> 否则：返回 tensor.data

如果一个 exclude 算子既保存输入又保存内部产生的 tensor，只有 exact input object 命中属性；内部 tensor 继续按现有
逻辑保存。没有被 pack 实际命中的输入 handle 保持 used=False，重算时不需要固化，也不需要为其输出增加 replay
boundary。

7. checkpoint 重算流程

重算阶段调用 exclude wrapper 时：

1. 从 invocation-local cache 中弹出匹配的 CheckpointExcludeCallRecord。
2. 遍历 input_bindings。
3. 按 path 从当前 args、kwargs 解析真实 Tensor。
4. 校验路径存在且结果仍为 Tensor。
5. 使用 recomputed_input.data 固化对应 handle。
6. 返回记录中的原始 output，不执行 exclude body。

如果输入结构、wrapper 调用次数或调用顺序在首次前向和重算之间发生变化，应在 path 解析或 cache pop 处抛出明确
错误。

7.1 首次前向如何记录输入路径

记录的不是首次前向 Tensor，而是“输入路径 + 空 handle”。下面是实现骨架：

def _collect_input_bindings(args, kwargs):
    bindings = []
    previous_marks = []
    seen_tensor_ids = set()

    for path, tensor in _iter_tensor_inputs(args, kwargs):
        if _is_parameter(tensor) or id(tensor) in seen_tensor_ids:
            continue

        handle = _RecomputedInputHandle()
        previous = tensor._get_user_data(_RECOMPUTE_INPUT_HANDLE_KEY)
        tensor._set_user_data(_RECOMPUTE_INPUT_HANDLE_KEY, handle)

        bindings.append(_InputBinding(path=path, handle=handle))
        previous_marks.append((tensor, previous))
        seen_tensor_ids.add(id(tensor))

    return bindings, previous_marks

路径只保存 Python 容器位置，例如：

call:
  wrapper(x, [mask, hidden], bias=bias)

paths:
  args[0]       -> x
  args[1][0]    -> mask
  args[1][1]    -> hidden
  kwargs[bias]  -> bias

内部将路径表示为自描述的不可变 token tuple：

(("arg", 1), ("index", 1))                      # args[1][1]
(("kwarg", "bias"),)                            # kwargs["bias"]
(("arg", 0), ("key", "items"), ("index", 1))  # args[0]["items"][1]

_InputBinding 不保存 Tensor 引用，因此不会因为调用记录自身而延长首次前向输入的生命周期。临时保存 Tensor 的
previous_marks 只用于 body 退出时恢复属性，不能放进 cache record。

pack hook 命中属性时，将对应 handle 标记为 used=True：

def _pack_saved_tensor(tensor):
    handle = tensor._get_user_data(_RECOMPUTE_INPUT_HANDLE_KEY)
    if isinstance(handle, _RecomputedInputHandle):
        handle.mark_used()
        return handle
    return tensor.data

exclude body 执行完成后，把 output 和 bindings 一起加入现有 invocation-local FIFO cache：

record = _CheckpointExcludeCallRecord(
    output=output,
    input_bindings=bindings,
)
cache.save(id(self), record)
return _apply_replay_boundary(output, bindings)

这里 cache 保存的是 exclude body 的原始 output；首次前向和 replay 都根据 bindings 中是否存在 used=True 的 handle
应用相同的 boundary，保证 MindSpore 外层 placeholder 的 pack 顺序一致。

7.2 重算时如何找到对应输入

checkpoint replay 会使用重算数据流重新调用同一个 wrapper，因此此时的 args、kwargs 已经包含上游重新生成的
Tensor。wrapper 先按 wrapper identity 从 FIFO cache 取出本次调用对应的 record，再解析路径：

def _resolve_input(args, kwargs, path):
    root_kind, root_key = path[0]
    value = args[root_key] if root_kind == "arg" else kwargs[root_key]

    for _, token_value in path[1:]:
        value = value[token_value]
    return value

replay 分支的主体如下：

record = cache.pop(id(self))

for binding in record.input_bindings:
    if not binding.handle.used:
        continue

    recomputed_input = _resolve_input(args, kwargs, binding.path)
    if not isinstance(recomputed_input, ms.Tensor):
        raise RuntimeError(
            "Checkpoint replay did not reproduce a tensor input required "
            "by a checkpoint-excluded region"
        )
    binding.handle.materialize(recomputed_input.data)

return _apply_replay_boundary(record.output, record.input_bindings)

这里的 FIFO 很重要。例如同一个 wrapper 在一次 checkpoint 中调用两次：

forward:  wrapper(A1) -> record1
          wrapper(A2) -> record2

replay:   wrapper(A1') -> pop record1 -> handle1.materialize(A1'.data)
          wrapper(A2') -> pop record2 -> handle2.materialize(A2'.data)

不同 checkpoint invocation 使用不同的 _ExcludeCache resource，因此不会把不同 step 或嵌套 checkpoint
的 handle 交叉固化。

7.3 materialize(recomputed_input.data) 做了什么

materialize 本身只做一次引用赋值：

class _RecomputedInputHandle:
    def __init__(self):
        self._tensor = None
        self.used = False

    def mark_used(self):
        self.used = True

    def materialize(self, tensor):
        self._tensor = tensor

    def get_recomputed_tensor(self):
        if self._tensor is None:
            raise RuntimeError(
                "Checkpoint-excluded input was requested before recomputation"
            )
        return self._tensor

调用方传入的是 recomputed_input.data：

• 不进行 clone 或 device copy，不产生同尺寸的新显存分配。
• 与重算输入共享 storage，使 backward 读取的是重算得到的数值。
• 使用 detached data，避免 handle 持有重算阶段临时 autograd graph。
• handle 作为原始 autograd saved-tensor 的 packed payload 被持有，因此该重算 storage 会存活到对应 unpack。

固化完成后，invocation cache 可以按现有生命周期清理 call record；真正需要 backward 的 handle 仍由 autograd packed
payload 持有。

8. replay boundary

8.1 为什么需要 boundary

MindSpore 非重入 recompute 由外层 saved-tensor placeholder 的 unpack 触发。exclude 使用内层
saved_tensors_hooks，会覆盖外层 checkpoint hook。如果一个 checkpoint 的 backward 最先读取的是本设计的特殊
handle，而此前没有其他外层 placeholder 触发 replay，handle 可能在 checkpoint 重算前被 unpack。

因此不能仅假设“backward 一定已经重算”。需要为使用延迟输入 handle 的 exclude 输出增加一个轻量 replay boundary，
保证 backward 进入 exclude 原始反向节点前，checkpoint replay 已经执行到该 exclude。

8.2 boundary 行为

在 exclude body 和内层 saved-tensor hook 退出后，对输出 Tensor 叶子应用内部 identity autograd function：

forward:
  save one dependency through outer checkpoint hook
  return value-equivalent output

backward:
  access saved dependency, triggering checkpoint replay if needed
  pass gradient through unchanged

首次前向中的 boundary save 由外层 checkpoint hook 打包为 native placeholder，不保留额外 output storage。checkpoint
重算时，exclude wrapper 先固化输入 handles，再对缓存 output 重放同样的 boundary，使 MindSpore 外层 recompute hook
按相同顺序填充 placeholder。

为防止某个多输出分支未参与 loss，所有需要梯度的输出 Tensor 叶子都应经过 boundary；非 Tensor 和不参与梯度的输出
保持不变。

boundary 只在当前 exclude 调用实际创建了至少一个 RecomputedInputHandle 时启用。没有可优化输入的 exclude 保持
现有路径，避免无意义的 autograd 节点和 pack 调用。

8.3 与 remat 的关系

该 boundary 只保证 replay 时序，不处理 storage、view 或 layout。它对应 torch_remat 中“确保 replay 到达保存区域”
的必要机制，但不会引入其 view 重建和复杂 saved-input 分类逻辑。

9. backward unpack 流程

downstream backward
  -> replay boundary backward accesses native placeholder
       -> MindSpore triggers checkpoint replay
       -> upstream recompute regenerates input A'
       -> exclude replay materializes RecomputedInputHandle(A')
       -> exclude returns cached output and replays boundary
  -> original exclude backward starts
       -> unpack(RecomputedInputHandle)
       -> returns A'.data

这样，首次前向的 A 可以在 forward 后释放；A' 只从 replay 开始存活到 exclude backward 消费完成。

10. 正确性限制

10.1 exact Tensor object

优化只命中带属性的原输入对象。如果算子保存的是输入 view 或另一个别名对象，该对象没有属性，pack 会保存真实
tensor。这是预期的安全回退。

10.2 输入原地修改

如果 exclude body 原地修改输入，并为 backward 保存修改后的同一 Tensor，重算阶段只会得到进入 exclude 前的输入，
而 exclude body 又不会重跑，因此无法恢复修改后的值。

本设计将“exclude 不原地修改需要保存的输入”作为使用约束。首版不增加 version/storage 检查，避免引入额外 host
开销；需通过文档和负向测试明确该限制。

10.3 参数输入

Parameter 不参与该优化。参数本身由模型和参数管理机制持有，延迟保存不能释放其 storage，反而可能干扰参数更新、
FSDP unshard/reshard 或 version 语义。

10.4 多次 replay

PP 可能预触发 recompute，并在 dx、dw 阶段复用 recompute session。handle 必须支持覆盖固化值，cache 与 MindSpore
recompute session 的 retain 行为必须保持一致，不能在第一次 replay 或 unpack 后提前释放调用记录。

11. host 开销

pack 热路径新增成本：

• 一次 _get_user_data。
• 一次 sentinel 或 handle 类型判断。

以下工作只发生在 exclude wrapper 边界，而不是每次 pack：

• 遍历 args/kwargs Tensor 叶子。
• Tensor identity 去重。
• 创建 handle 和输入路径。
• 设置并恢复 Tensor user-data。
• 输出 replay boundary 包装。

不执行 storage weak-ref、layout、stride、offset、version 或 view recipe 计算。

12. 影响范围

主要实现文件：

• hyper_parallel/platform/mindspore/activation_checkpoint/checkpoint_exclude_wrapper.py

可能需要的内部状态调整：

• hyper_parallel/core/activation_checkpoint/recompute_state.py 不需要修改；调用记录继续使用现有 invocation resource。

测试文件：

• tests/ut/platform/mindspore/activation_checkpoint/test_checkpoint_exclude_wrapper.py
• tests/mindspore/st/activation_checkpoint/checkpoint_exclude_gmm.py

本设计不修改公开 API 签名，不要求改动 Platform、MindSporePlatform 或 PyTorch 后端。

13. 测试方案

13.1 handle 与 pack 单元测试

• 无属性 Tensor：pack 继续返回 tensor.data。
• 带 handle 属性的 exact input：pack 返回 handle。
• handle 不持有首次前向 Tensor。
• materialize 前 unpack：抛出 RuntimeError。
• materialize 后 unpack：返回重算 Tensor。
• 同一 Tensor 在多个参数位置出现：只创建一个 handle，属性能正确恢复。
• tuple、list、dict 和 kwargs 中的 Tensor 能按路径正确固化。
• Parameter 不添加标记。

13.2 PyNative 集成测试

• recompute op -> exclude -> recompute op：exclude body 只执行一次，value、dx、dw 与基线一致。
• exclude 位于 checkpoint 尾部，确保没有其他普通 placeholder 时 replay boundary 仍先触发重算。
• 自定义算子明确保存 exact input，确认 pack 返回特殊 handle。
• 保存输入 view：不命中 handle，回退现有保存路径且梯度正确。
• 连续两个 exclude：上游缓存输出能在 replay 中固化下游 handle，默认语义不变。
• 同一 wrapper 多次调用及连续多个 training step：cache FIFO 和属性清理正确。
• 前向异常：临时属性在 finally 中恢复。

13.3 PP 回归测试

• 保留现有 pre-fired recompute、dx/dw split 测试。
• 分别覆盖 retain_on_unpack=True 和 False。
• 验证 handle 在多阶段 replay 中不会提前清空或绑定到错误 invocation。

13.4 显存测试

在 Ascend NPU 上使用 RMSNorm -> GMM -> square -> mean 端到端网络。完整 block 使用 checkpoint，GMM
分别采用普通重算和 checkpoint_exclude_wrapper：

1. 普通重算、优化后的 exclude 和关闭输入 rematerialization 的 legacy exclude 分别在独立进程运行，排除 allocator
   执行顺序对峰值的影响。
2. 对比 loss、输入梯度和所有参数梯度的统计量及固定位置采样值。
3. 普通重算中 block 和 GMM 均执行两次；exclude 模式中 block 执行两次、GMM 只执行一次。
4. 输入规模为 8192 x 2048 x bf16，一个 RMSNorm 输出 activation 为 32 MiB。
5. 先执行一次同形状 warmup，清理梯度和临时结果并重置显存统计；第二次执行在 forward 完成、backward 开始前
   采集当前 allocated memory，同时记录完整 forward/backward 的峰值用于分析。
6. 完整 backward 峰值会受到缓存输出、重算输入、GMM workspace 和梯度生命周期共同影响。实测 pure recompute、
   optimized exclude、legacy exclude 的 forward 增量分别为 0/64/64 MiB，完整峰值分别为
   256/288/256 MiB。pack 记录确认 optimized 路径命中了 GMM 保存的 RMSNorm 输入，但 MindSpore GMM 输出的
   autograd graph 仍保留了等量大对象，因此该算子场景没有观察到净 forward 显存下降。
7. ST 据此验证可解释的显存上界：exclude forward 常驻不超过两份大 activation 加 allocator 余量，不比 legacy
   forward 明显更高，完整峰值相对 pure recompute 不超过约 1.5 份 activation。若后续 MindSpore GMM 的 graph
   retention 改善，该断言允许 optimized 显存进一步下降。

14. 实施顺序

1. 实现 RecomputedInputHandle、InputBinding 和调用记录。
2. 实现输入遍历、identity 去重、属性设置与 finally 恢复。
3. 修改 pack/unpack hook，保持未命中路径不变。
4. 在 replay 分支按输入路径固化 handles。
5. 实现只在有延迟 handle 时启用的 replay boundary。
6. 完成单元、PyNative、PP 和 NPU 显存测试。

15. 结论

本方案只优化 recompute -> exclude 边界：exclude 的 exact input 通过 Tensor 属性关联特殊 handle，首次前向不保存其
device storage，checkpoint replay 到达 exclude 时再固化真实 tensor。pack 热路径不做 view/storage/layout 判断；
复杂输入安全回退到现有保存行为。连续 exclude 输出融合作为独立问题，不进入本次设计和实现。