r2.1.0-beta1
r2.1.0-beta1 的 PyNative Expert Parallel dispatch hook 将路由概率、恢复索引、通信 split 等状态保存在每层长期存在的 self.ctx 中。combine hook 解包并使用这些状态后没有清空成员引用。
self.ctx
因此每个 MoE layer 在一次前向/反向完成后,仍持续持有最后一次调用的 probs、unsort_token_indices_experts 以及 fusion restore map 等 Tensor;最后一步会一直保留到模型销毁。这不是逐步无限增长,但会形成按 MoE layer 累积的常驻激活/索引内存。
probs
unsort_token_indices_experts
期望行为:combine 将状态转移到局部变量后清空 self.ctx,让自动求导仅保留反向真正需要的 Tensor,并允许额外成员引用在 combine 返回后释放。
建议将同一修复同步到 r2.1.0-beta1 的普通 EP 和 Deredundancy EP,并增加弱引用生命周期测试。
该问题本质是 Python 成员引用生命周期问题,不依赖特定 NPU 型号。
隔离验证中,一个仅由 self.ctx 持有的 64 MiB NPU Tensor 在 combine 后可回收,memory_allocated 从 67,109,376 字节降至 0;普通前向/反向和 non-reentrant recompute 梯度保持正确。
memory_allocated
MindFormers r2.1.0-beta1。
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该问题本质是 Python 成员引用生命周期问题,不依赖特定 NPU 型号。
其他辅助信息
隔离验证中,一个仅由
self.ctx持有的 64 MiB NPU Tensor 在 combine 后可回收,memory_allocated从 67,109,376 字节降至 0;普通前向/反向和 non-reentrant recompute 梯度保持正确。版本信息
MindFormers
r2.1.0-beta1。