已开启
[Bug]: MindSpore 2.10 非重入 recompute 每 step 持续残留 Tensor,导致 host 内存线性增长 #300
niujunhao创建于  7月23日
niujunhao
7月23日 创建

[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:

  • 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 修复如下:

该改动在 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:

该 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
likedislike