[Bug]: 关闭 recompute 后 MindSpore HSDP pynative 训练 host 内存仍持续上涨

[Bug]: 关闭 recompute 后 MindSpore HSDP pynative 训练 host 内存仍持续上涨

Checklist

• recompute 已完全关闭，配置中没有 recompute 段。
• 已完成 8 卡真实训练 A/B、Memray native 栈分析和 CPU/NPU 最小复现。
• 训练结束后已确认没有残留进程或 NPU 占用。

问题描述

MindFormers + Hyper-Parallel 的 8 卡 Ascend pynative 训练在关闭 recompute 后，host 匿名内存仍随 step 缓慢持续上涨。

关闭 recompute 后，之前观察到的 +12 Tensor/step 已消失，GC 可见的 Tensor 数量稳定为 2067；但无 profiler、从外部读取 rank 0 的 /proc/<pid>/smaps_rollup 时，仍有以下增长：

step=100  anon=3086.5 MiB
step=200  anon=3089.5 MiB
step=300  anon=3092.0 MiB

即 200 step 增长约 5.5 MiB，约为 28 KiB/step。因此，该问题不是 recompute Tensor 残留的延续，而是训练常规 forward/backward 路径中的其他生命周期问题。

目前已确认有两个无界增长源，以及一个会在短窗口中表现为增长、但最终有上限的 MindSpore cache。

根因 1：HSDP 每 step 注册 Tensor hook，但丢弃 handle

Hyper-Parallel 每次 forward 都在输出 Tensor 上注册 backward pre-hook：

# hyper_parallel/platform/mindspore/fully_shard/scheduler.py
def _register_backward_pre_hook(self, outputs):
    flat_outputs, _ = tree_flatten(outputs)
    for output in flat_outputs:
        if isinstance(output, ms.Tensor) and output._requires_grad:
            output.register_hook(self._backward_pre_hook)

Tensor.register_hook() 返回 _TensorHookHandle，该 handle 提供显式 remove()，且自身没有析构自动 remove。当前代码直接丢弃 handle。

真实训练中每 step 注册约 182 个此类 hook。Memray 的 10-step/20-step 对照中，Tensor.register_hook 路径的存活 native 分配随窗口长度增长。

CPU 最小 A/B

同一个简单 ms.grad 循环运行 10000 次：

无 hook:
step=1000   anon=922.113 MiB
step=10000  anon=956.762 MiB

注册 hook，丢弃 handle:
step=1000   anon=922.234 MiB
step=10000  anon=957.656 MiB

注册 hook，backward 后 handle.remove():
step=1000   anon=922.160 MiB
step=10000  anon=956.926 MiB

显式 remove 后基本回到无 hook baseline；丢弃 handle 会额外增加约 75 bytes/call。

真实训练 A/B

增加一个仅用于定位的开关，在每个 backward final callback 中 remove 本 step 保存的 hook handles：

HP_DEBUG_REMOVE_BACKWARD_PRE_HOOKS=1

关闭 recompute、只启用该清理后的 rank 0 数据：

step=100  anon=3084.5 MiB
step=150  anon=3085.2 MiB
step=200  anon=3085.8 MiB
step=250  anon=3086.4 MiB
step=300  anon=3086.9 MiB

step 100→300 仅增长 2.4 MiB，约 12.3 KiB/step；相对原始 28 KiB/step 下降约 56%。这证明丢弃 Tensor hook handle 是关闭 recompute 后的主要增长源之一。

根因 2：MindSpore custom Function 未释放 needs_input_grad

该训练每 step 调用 133 次 MindSpore custom autograd Function：

PostBackwardFunction:       122 次/step
_MoEAuxLossAutoScaler:       10 次/step
_ChunkCrossEntropyLoss:       1 次/step

MindSpore FunctionBase 每次 apply 都会创建 needs_input_grad tuple：

PyObject *need_grad_input = PyTuple_New(inputs_size);
ctx->needs_input_grad = need_grad_input;

但当前 FunctionBase_clear() 重复清理了两次 saved_tensors，没有清理 needs_input_grad：

void FunctionBase_clear(FunctionBase *self) {
  Py_CLEAR(self->saved_tensors);
  Py_CLEAR(self->dirty_tensors);
  Py_CLEAR(self->non_differentiable);
  Py_CLEAR(self->saved_tensors);  // 重复
  self->is_tensor_input.clear();
}

相关源码：

• https://gitee.com/mindspore/mindspore/blob/master/mindspore/ccsrc/pynative/backward/hook/function_py.cc

CPU 最小复现 1000 step：

apply-count=0:
step=100   anon=703.6 MiB
step=1000  anon=706.4 MiB

apply-count=133:
step=100   anon=706.1 MiB
step=1000  anon=714.5 MiB

apply-count=133，backward 中 ctx.needs_input_grad=None:
step=100   anon=705.4 MiB
step=1000  anon=708.2 MiB

显式清空后回到 baseline。该问题约贡献 7 KiB/step。

在真实训练中同时启用 HSDP hook 清理和 needs_input_grad 临时清理后：

step=100  anon=3084.0 MiB
step=150  anon=3084.3 MiB
step=200  anon=3084.4 MiB
step=250  anon=3084.7 MiB
step=300  anon=3084.8 MiB

step 100→300 仅增长 0.8 MiB，约 4.1 KiB/step。两个清理合计将原始增长从约 28 KiB/step 降低约 85%。

MindSpore collective event 生命周期问题

Memray 10-step/20-step 聚合对照显示，额外 10 个训练 step 有 654938 bytes、10933 次仍存活的 traced native 分配，约 64 KiB/step。主要路径为：

约 23 KiB/step  Muon norm -> PyNativeExecutor::RunOpStub
约 34 KiB/step  all_gather/reduce_scatter/all_reduce generated PyBoost wrapper
其余             Tensor hook 和少量其他算子

一个 step 的完整 native 栈显示，每个 HSDP collective 都创建一个 event：

AscendResManager::CreateEventWithFlag
CommHandle::CreateEvent
PyboostDistCommAllGatherIntoTensorOp / PyboostDistCommReduceScatterTensorOp /
PyboostDistCommAllReduceOp

MindSpore 源码中，CreateEventWithFlag() 会把 event 强引用保存到 device_events_，只有 CommHandle 析构调用 DestroyEvent() 时才从列表移除：

• https://gitee.com/mindspore/mindspore/blob/master/mindspore/ccsrc/plugin/ascend/res_manager/ascend_res_manager.cc
• https://gitee.com/mindspore/mindspore/blob/master/mindspore/ccsrc/pynative/utils/pyboost/comm_handle.cc
• https://gitee.com/mindspore/mindspore/blob/master/mindspore/ccsrc/pynative/forward/pyboost/template/pyboost_comm_function.tpl

2 卡最小复现

最小脚本见 repro_collective_event_host_memory.py。核心循环为：

result = dist.all_reduce(tensor, async_op=async_op)
if async_op:
    result.wait()
del result

# 每 500 次执行
ms.runtime.synchronize()
gc.collect()

运行 5000 次后，无论 async_op=True 还是 False，rank 0 都增长约 1.2 MiB，约 0.25 KiB/call；synchronize() 和 gc.collect() 均不能使其回落：

async_op=True:
step=500   anon=1380.840 MiB  malloc_inuse=616.624 MiB
step=5000  anon=1382.051 MiB  malloc_inuse=617.866 MiB

async_op=False:
step=500   anon=1382.066 MiB  malloc_inuse=616.622 MiB
step=5000  anon=1383.242 MiB  malloc_inuse=617.876 MiB

因此这不是 Hyper-Parallel 忘记调用 wait()，也不是普通 Python GC 积压。同步通信同样增长，说明 generated communication wrapper 内部的 CommHandle/event 生命周期需要 MindSpore 侧进一步检查。

已排除：RunOpStub 的 node abstract cache 是有界的

Muon norm 的 native 栈落在：

PyNativeExecutor::RunOpStub
Tensor::ToAbstract / MakeAbstractTensor / AbstractTensor

MindSpore InferOperation::node_abs_cache_ 会缓存 run-op output abstract。源码说明纯 run-op 时该 cache 不会主动清空，但超过 10000 项会删除前 5000 项：

if (node_abs_cache_.size() > kCacheThreshold) {
  auto half_it = node_abs_cache_.begin();
  std::advance(half_it, kHalfThreshold);
  node_abs_cache_.erase(node_abs_cache_.begin(), half_it);
}

CPU 连续运行 20000 次 ops.norm 后，malloc_inuse 在阈值附近按设计锯齿回落，anon 最终稳定。因此该路径会放大短窗口增长，但不是无限长稳上涨的根因。

相关源码：

• https://gitee.com/mindspore/mindspore/blob/master/mindspore/ccsrc/pynative/forward/op_graph/do_infer.cc

Expected behavior

1. Hyper-Parallel 应保存 Tensor.register_hook() 返回的 handle，并在对应 backward 完成后 remove，不能每 step 丢弃。
2. MindSpore FunctionBase_clear() 应执行 Py_CLEAR(self->needs_input_grad)。
3. MindSpore collective 在 wait() 或同步调用完成后应及时销毁或复用 CommHandle event；连续同步/异步 collective 不应按调用次数增加 host 内存。
4. 关闭 recompute 后，经过 warm-up，host 匿名内存、RSS 和 malloc_inuse 应达到稳定平台。

建议回归用例

1. HSDP 层连续 forward/backward 1000 step，检查 Tensor hook handle 数量及 host 内存。
2. custom Function 连续 apply/backward，检查 needs_input_grad tuple 是否释放。
3. Ascend 上同步和异步 all_reduce 各循环 10000 次，检查 event/CommHandle 数量以及 host 内存。
4. 真实 8 卡配置关闭 recompute 后运行长稳，至少比较 step 100 和 step 1000。

环境信息

MindSpore:       2.10.0
Python:          3.11
Execution mode:  PYNATIVE_MODE
Device:          8-card Ascend training; 2-card Ascend collective reproducer; CPU unit reproducers
Model:           12-layer DeepSeek V3
Sequence length: 4096
Local batch size: 1
Global batch size: 8
Recompute:       disabled

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

Start with hyper_parallel/platform/mindspore/fully_shard/scheduler.py and the hook-cleanup A/B results, then inspect function_py.cc for FunctionBase_clear. Use repro_collective_event_host_memory.py and the referenced Ascend collective files to separate the three growth paths. Done means the listed hook, custom Function, and collective regression checks show stable host memory after warm-up.