[Bug]: MindSpore 2.10 非重入 recompute 每 step 持续残留 Tensor，导致 host 内存线性增长

[Bug]: MindSpore 2.10 非重入 recompute 每 step 持续残留 Tensor，导致 host 内存线性增长

Checklist

• 已搜索现有 Issues，未发现相同问题。
• 已阅读相关代码和文档。
• 已提供可独立运行的最小复现，以及完整的 A/B 定位数据。

🐛 Describe the bug

在 MindFormers pynative 长稳训练中，通过 HyperParallel 启用 full recompute 后，host 匿名内存会随训练 step 缓慢、持续上涨，无法达到稳定平台。

进一步统计 Python GC 可见对象后发现，每个启用 recompute 的 transformer layer 每 step 会残留一个 Tensor：

• 12 层 full recompute：每 step 增加 12 个 Tensor；
• 6 层 full recompute：每 step 增加 6 个 Tensor；
• CPU 最小复现、单个 recompute block：每 step 增加 1 个 Tensor。

gc.collect() 无法回收这些对象，返回值持续为 0。CPU 最小复现不包含数据集、Muon 优化器、训练 callback 或 Ascend page cache，因此这些因素可以排除。

HyperParallel 的 checkpoint() 当前固定使用非重入重计算：

# hyper_parallel/core/activation_checkpoint/activation_checkpoint.py
return plat.checkpoint(
    function, *args, context_fn=composed_context_fn, use_reentrant=False, **kwargs
)

在当前 MindSpore 2.10.0 环境中，该路径会持续保留 saved tensor 相关对象。尝试切换到 use_reentrant=True 不能作为临时规避方案，因为会在当前场景报错：

RuntimeError: The pointer[top_cell_] is null

真实训练 A/B 结果

训练配置的模型为 12 层 DeepSeek V3，8 卡 Ascend，pynative 模式，序列长度 4096，local batch size 为 1。

12 层启用 full recompute 时，连续两个 20-step 采样点之间均增加 240 个 Tensor：

step=40   Tensor delta over 20 steps: +240
step=60   Tensor delta over 20 steps: +240
step=80   Tensor delta over 20 steps: +240
step=100  Tensor delta over 20 steps: +240

临时将 full recompute 缩小到 6 层后，增长速度同步减半：

step=40   Tensor delta over 20 steps: +120

12 层 full recompute 对应的进程匿名内存采样如下：

step=40   anon=3155.3 MiB
step=60   anon=3176.5 MiB
step=80   anon=3186.5 MiB
step=100  anon=3195.2 MiB

Tensor 增长量与启用 recompute 的层数严格对应，说明增长源位于每层每 step 都会执行一次的 recompute/saved-tensor 生命周期，而不是一次性缓存。

CPU 最小复现

以下代码在 CPU 上即可复现，不依赖 Ascend、多卡通信、模型、数据集或优化器：

"""Minimal host-memory reproducer for MindSpore non-reentrant recompute."""

import argparse
import gc
import os

import mindspore as ms
from mindspore import mint, nn

from hyper_parallel.platform.mindspore.autograd_compat import enable_mindspore_backward_compat


class Block(nn.Cell):
    """Small differentiable block used by the reproducer."""

    def construct(self, value):
        return mint.mul(value, value)


def tensor_count():
    return sum(type(obj).__name__ == "Tensor" for obj in gc.get_objects())


def rss_mib():
    with open("/proc/self/status", "r", encoding="utf-8") as stream:
        for line in stream:
            if line.startswith("VmRSS:"):
                return int(line.split()[1]) / 1024
    return 0.0


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--steps", type=int, default=100)
    parser.add_argument("--gc", action="store_true")
    parser.add_argument("--reentrant", action="store_true")
    args = parser.parse_args()

    ms.set_context(mode=ms.PYNATIVE_MODE, device_target="CPU")
    enable_mindspore_backward_compat()

    block = Block()
    value = mint.ones((1024, 1024), dtype=ms.float32)
    value.requires_grad = True

    def forward(input_value):
        output = ms.recompute(block, input_value, use_reentrant=args.reentrant)
        return output.sum()

    grad_fn = ms.value_and_grad(forward, grad_position=0)
    for step in range(1, args.steps + 1):
        loss, grad = grad_fn(value)
        del loss, grad

        if args.gc:
            collected = gc.collect()
        else:
            collected = None

        if step == 1 or step % 10 == 0:
            print(
                f"pid={os.getpid()} step={step} rss={rss_mib():.1f}MiB "
                f"tensors={tensor_count()} gc_collect={collected}",
                flush=True,
            )


if __name__ == "__main__":
    main()

运行命令：

PYTHONPATH=/path/to/hyper-parallel \
python repro_recompute_host_memory.py --steps 50 --gc

实测输出中的 Tensor 数量持续增长：

step=1   tensors=15
step=10  tensors=24
step=20  tensors=34
step=30  tensors=43
step=40  tensors=53
step=50  tensors=63

启用 --gc 后增长趋势不变，且 gc_collect=0。

Expected behavior

一次 forward/backward 完成后，非重入 recompute 为该 step 保存的 Tensor、hook 和 autograd graph 引用应被释放。经过少量 warm-up 后：

• Python GC 可见的 Tensor 数量应保持稳定；
• host 匿名内存和 RSS 应达到稳定平台，不应与训练 step 数线性相关；
• Tensor 数量不应与启用 recompute 的 layer 数量呈“每层每 step 增加一个”的关系。

建议：

1. 确认 HyperParallel 支持的 MindSpore 版本是否已经包含 saved-tensor unpack 生命周期修复；
2. 若 MindSpore 2.10.0 仍是支持版本，考虑在 HyperParallel 侧提供兼容处理、版本约束或明确告警；
3. 增加 MindSpore non-reentrant checkpoint 多 step 回归用例，验证 backward 后 Tensor/hook 可以释放。

Additional context

高度相关的 MindSpore 修复如下：

• PR: https://gitee.com/mindspore/mindspore/pulls/91580
• Commit: 1ca0512208657730336641865f6de4556481b235
• Title: fix memory leak and circular references when unpack saved output tensor

该改动在 SavedTensor::UnWrapToTensor 中对 saved tensor unpack 的结果统一执行 ShallowCopyAndDetachForTensor，并说明未 detach 时会形成以下循环引用：

data -> auto_grad_meta_data -> grad_node -> data

补丁的关键变化为：

- auto data =
-   saved_tensor_hook_ ? CommonUtils::ShallowCopyAndDetachForTensor(saved_tensor_hook_->RunUnpackHook()) : data_;
+ auto data =
+   CommonUtils::ShallowCopyAndDetachForTensor(saved_tensor_hook_ ? saved_tensor_hook_->RunUnpackHook() : data_);

这与当前观察到的 saved tensor 持续残留现象高度吻合，但目前尚未使用包含该提交的 MindSpore 安装包完成最终 A/B 验证，因此这里将其记录为最可能的上游根因，而不是已经验证的结论。

相关但不重复的 HyperParallel issue：

• https://gitcode.com/mindspore/hyper-parallel/issues/164

该 issue 是 fully_shard 参数 view 污染非重入 recompute placeholder 序列后导致 unpack 报错；本 issue 是不依赖 fully_shard、可在 CPU 最小用例稳定复现的跨 step Tensor/host 内存泄漏。

Environment info

MindSpore:       2.10.0
MindFormers:     e8383c0a11c14b4b112dc77a55b8484549e6ce94
HyperParallel:   1fca9c7e7c4ac1378c25b774b9bdb40f85db2ae6
Python:          3.11
Execution mode:  PYNATIVE_MODE
Device:          8-card Ascend for full training; CPU for minimal reproduction
Model:           12-layer DeepSeek V3
Sequence length: 4096
Local batch size: 1
Global batch size: 8

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

Start with hyper_parallel/core/activation_checkpoint/activation_checkpoint.py and run the CPU reproducer repro_recompute_host_memory.py with --gc on MindSpore 2.10.0. Compare behavior with a build containing commit 1ca0512208657730336641865f6de4556481b235, then determine whether the supported-version handling needs a compatibility change, version constraint, or warning. Done means Tensor counts and host memory stop growing across multiple steps, with regression coverage for non-reentrant recompute.