- auto data =
- saved_tensor_hook_ ? CommonUtils::ShallowCopyAndDetachForTensor(saved_tensor_hook_->RunUnpackHook()) : data_;
+ auto data =
+ CommonUtils::ShallowCopyAndDetachForTensor(saved_tensor_hook_ ? saved_tensor_hook_->RunUnpackHook() : data_);
[Bug]: MindSpore 2.10 非重入 recompute 每 step 持续残留 Tensor,导致 host 内存线性增长
Checklist
🐛 Describe the bug
在 MindFormers pynative 长稳训练中,通过 HyperParallel 启用 full recompute 后,host 匿名内存会随训练 step 缓慢、持续上涨,无法达到稳定平台。
进一步统计 Python GC 可见对象后发现,每个启用 recompute 的 transformer layer 每 step 会残留一个
Tensor:Tensor;Tensor;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不能作为临时规避方案,因为会在当前场景报错:真实训练 A/B 结果
训练配置的模型为 12 层 DeepSeek V3,8 卡 Ascend,pynative 模式,序列长度 4096,local batch size 为 1。
12 层启用 full recompute 时,连续两个 20-step 采样点之间均增加 240 个
Tensor:临时将 full recompute 缩小到 6 层后,增长速度同步减半:
12 层 full recompute 对应的进程匿名内存采样如下:
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()运行命令:
实测输出中的 Tensor 数量持续增长:
启用
--gc后增长趋势不变,且gc_collect=0。Expected behavior
一次 forward/backward 完成后,非重入 recompute 为该 step 保存的 Tensor、hook 和 autograd graph 引用应被释放。经过少量 warm-up 后:
Tensor数量应保持稳定;建议:
Additional context
高度相关的 MindSpore 修复如下:
1ca0512208657730336641865f6de4556481b235fix memory leak and circular references when unpack saved output tensor该改动在
SavedTensor::UnWrapToTensor中对 saved tensor unpack 的结果统一执行ShallowCopyAndDetachForTensor,并说明未 detach 时会形成以下循环引用:补丁的关键变化为:
- 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:
该 issue 是 fully_shard 参数 view 污染非重入 recompute placeholder 序列后导致 unpack 报错;本 issue 是不依赖 fully_shard、可在 CPU 最小用例稳定复现的跨 step Tensor/host 内存泄漏。
Environment info