已开启
[Bug]: [r2.1.0-beta1] PyNative EP combine 后仍持有 dispatch 上下文 #2528
husichao创建于  8月19日
husichao成员
8月19日 创建

Checklist

🐞 问题详细描述

r2.1.0-beta1 的 PyNative Expert Parallel dispatch hook 将路由概率、恢复索引、通信 split 等状态保存在每层长期存在的 self.ctx 中。combine hook 解包并使用这些状态后没有清空成员引用。

因此每个 MoE layer 在一次前向/反向完成后,仍持续持有最后一次调用的 probs、unsort_token_indices_experts 以及 fusion restore map 等 Tensor;最后一步会一直保留到模型销毁。这不是逐步无限增长,但会形成按 MoE layer 累积的常驻激活/索引内存。

期望行为:combine 将状态转移到局部变量后清空 self.ctx,让自动求导仅保留反向真正需要的 Tensor,并允许额外成员引用在 combine 返回后释放。

建议将同一修复同步到 r2.1.0-beta1 的普通 EP 和 Deredundancy EP,并增加弱引用生命周期测试。

详细的环境信息描述

  • 分支:r2.1.0-beta1
  • 执行模式:MindSpore PyNative
  • 并行模式:Expert Parallel
  • 验证硬件:Ascend 910B2

该问题本质是 Python 成员引用生命周期问题,不依赖特定 NPU 型号。

其他辅助信息

隔离验证中,一个仅由 self.ctx 持有的 64 MiB NPU Tensor 在 combine 后可回收,memory_allocated 从 67,109,376 字节降至 0;普通前向/反向和 non-reentrant recompute 梯度保持正确。

版本信息

MindFormers r2.1.0-beta1。

likedislike