已合并
fix: resolve memory leak issue during recompute forward in context parallel attention #39895
pengjingyou创建于 7月2日
fix: resolve memory leak issue during recompute forward in context parallel attention #39895
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 pengjingyou 的贡献）7月2日 创建了 pull request，commit b12db823
ascend-robot
7月2日 评论：
7月2日 评论：
7月2日 添加了label：ascend-cla/yes
7月2日 添加了label：needs-issue
ascend-robot
7月2日 评论：
7月2日 评论：
Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| torch_npu/distributed | ✅ renyujin, wjlflyer (2/2) | ✅ renyujin (1/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
pengjingyou, thanks for your pull request. All authors of the commits have signed the CLA. 👍
ascend-robot
7月2日 评论：
7月2日 评论：
Linking Issue Notice
@pengjingyou , the pull request must be linked to at least one issue.
If an issue has already been linked, but the needs-issue label remains, you can remove the label by commenting /check-issue .
ascend-robot
7月2日 评论：
7月2日 评论：
当前仓库存在以下 保护分支 ：
| Protected Branch | Version | Release |
|---|---|---|
| master | ||
| v2.11.0 | ||
| v2.7.1 | ||
| v2.7.1-26.1.0 | ||
| v2.10.0 | ||
| v2.12.0 | ||
| v2.9.0 | ||
| v2.12.0-26.1.0 | ||
| v2.11.0-26.1.0 | ||
| v2.10.0-26.1.0 | ||
| v2.9.0-26.1.0 | ||
| ci-test | ||
| sync-pr28113--to-v2.9.0 |
评论 /sync <branch1> <branch2> ... 可将当前 PR 修改同步到其它分支（创建同步 PR）：
a) 如果当前 PR 是 Open 状态，同步操作将延迟到 PR 被合并时执行
b) 如果当前 PR 已经 Merged，将立即执行同步操作
注意：
- /sync 命令可以指定同步到多个分支，仅最后一个 /sync 命令生效
- 如果创建的同步 PR 不正确，可通过向同步 PR 的源分支提交轻量级 PR 完善，或使用 /close 命令关闭
7月2日 添加了label：ci-pipeline-running
7月2日 关联了issue：[Bug]: torchtitan + torch_npu Context Parallel 显存泄漏问题
7月2日 删除了label：needs-issue
ascend-robot
7月2日 评论：
7月2日 评论：
ascend docs pipeline is running...
7月2日 添加了label：docs-ci-pipeline-running
ascend-robot
7月2日 评论：
7月2日 评论：
✅ 跳过 docs ci 检查，没有需要检查的文档文件
7月2日 删除了label：docs-ci-pipeline-running
7月2日 添加了label：docs-ci-pipeline-success
7月2日 删除了label：ci-pipeline-running
7月2日 添加了label：ci-pipeline-passed
ascend-robot
7月2日 评论：
7月2日 评论：
流水线 PR-pipeline_pytorch#40739 [ commitID：3fa37c03 ] 已完成
>>>代码风格自动修复执行成功（无修复内容）
| 阶段 | 任务名 | 状态 | 详情 |
|---|---|---|---|
| 编译构建 | Build_X86 | ✅ | >>> |
| Build_ARM | ✅ | >>> | |
| Build_LibTorch_x86 | ✅ | >>> | |
| Build_LibTorch_ARM | ✅ | >>> | |
| Build_X86_torchair | 🛑 | >>> | |
| Build_ARM_torchair | 🛑 | >>> | |
| patch_test | 🛑 | >>> | |
| 恶意代码检查 | Antipoison | ✅ | >>> |
| 编码安全与规范检查 | CodeCheck | ✅ | >>> |
| check_error | ✅ | >>> | |
| CodeCheck_lintrunner | ✅ | >>> | |
| 开源片段检查 | SCA | ✅ | >>> |
| 开发者测试 | UT_X86_Part_01 | 🛑 | >>> |
| UT_X86_Part_02 | 🛑 | >>> | |
| UT_ARM_A3_Part_01 | 🛑 | >>> | |
| UT_ARM_A3_Part_02 | 🛑 | >>> | |
| UT_ARM_A2_Part_01 | ✅ | >>> | |
| UT_ARM_A2_Part_02 | ✅ | >>> | |
| UT_ARM_A2_Part_03 | ✅ | >>> | |
| UT_inductor_Part_01 | 🛑 | >>> | |
| UT_inductor_Part_02 | 🛑 | >>> | |
| UT_inductor_Part_03 | 🛑 | >>> | |
| UT_inductor_Part_04 | 🛑 | >>> | |
| UT_DIST_ARM_Part_01 | ✅ | >>> | |
| UT_DIST_ARM_Part_02 | ✅ | >>> | |
| UT_DIST_ARM_Part_03 | ✅ | >>> | |
| UT_DIST_ARM_Part_04 | ✅ | >>> | |
| UT_ARM_A2_Select_Part_01 | ✅ | >>> | |
| UT_ARM_A2_Select_Part_02 | ✅ | >>> | |
| 流水线 | PR-pipeline_pytorch | ✅ | >>> |
- compile、compile_inductor、compile_torchair : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线
atomgit-bot
7月2日 评论：
7月2日 评论：
变更摘要
此 PR 修复了 context parallel attention 在 recompute forward 过程中的内存泄漏问题。通过在 _npu_fa_v3_handler 中引入 torch._C._current_graph_task_id() 判断，跳过重计算前向阶段对全局变量 _step_cache_stack 的 push 操作，确保 push/pop 平衡，避免重复累积导致的显存泄漏。
主要改动
_npu_fa_v3_handler增加 recompute 判断：在向_step_cache_stack追加(step_caches, is_causal, attn_output.detach(), merged_lse.detach())之前，新增if torch._C._current_graph_task_id() < 0条件，当 graph task ID 为负数（即当前执行的是 recompute forward）时跳过 push，仅保留原始正向传播时的 push 操作。- 全局变量
_step_cache_stack的 push/pop 平衡修复：避免 recompute forward 重复向_step_cache_stack压入中间结果，导致 backward 阶段 pop 无法完全清空栈而产生内存泄漏。
atomgit-bot
7月2日 评论：
7月2日 评论：
代码审查
我已经完整审查了此 diff 中唯一的变更文件 _npu_attention.py。以下是详细分析：
变更分析
变更内容：在 _npu_fa_v3_handler 函数中，将对 _step_cache_stack.append() 的调用包裹在 if torch._C._current_graph_task_id() < 0: 条件判断中。
变更意图：修复 context parallel attention 在开启重计算（activation checkpointing）时的显存泄漏问题。
审查结论
核心逻辑验证
-
torch._C._current_graph_task_id()的语义：- 在 autograd 引擎未运行时返回
-1 - 在 backward 过程中返回非负整数（0, 1, 2, ...）
- 重计算 forward 发生在 backward 过程中，因此
_current_graph_task_id() >= 0
- 在 autograd 引擎未运行时返回
-
条件
< 0的正确性：- 正常 forward：
-1 < 0→ True → 执行 push ✓ - 重计算 forward：
N >= 0→< 0为 False → 跳过 push ✓
- 正常 forward：
-
栈平衡验证：
- 正常 forward → push，backward → pop → 平衡 ✓
- 带重计算：原始 forward → push，重计算 forward → 跳过 push，backward → pop → 平衡 ✓
-
_step_cache_stack使用链路：- 仅在本文件内使用（定义、push、pop、清理）
- 不涉及跨文件或跨进程的状态共享
npu_disable_cp_dtensor_dispatcher中有栈清理逻辑作为兜底
-
函数可用性：
torch._C._current_graph_task_id()是 PyTorch C++ autograd 引擎的 Python 绑定，在 PyTorch 2.10.0（目标分支）中稳定可用。
边界情况检查
- 嵌套 checkpoint：栈的 LIFO 性质保证 push/pop 配对正确
- 推理模式（无 backward）：行为与修改前一致（栈累积但不 pop），
npu_disable_cp_dtensor_dispatcher会清理 torch.compile/ JIT：CP attention 场景不使用编译路径，不适用- 多线程：PyTorch autograd 默认单线程，无竞争条件
未发现任何逻辑错误、安全风险、可靠性问题或 Breaking Change。此变更干净、正确。
审查总结
- P0: 0
- P1: 0
- P2: 0
- P3: 0
已审查文件：
torch_npu/distributed/tensor/experimental/_context_parallel/_npu_attention.py— 无问题
整体风险判断：低风险。此变更是对显存泄漏问题的精确修复，逻辑正确，栈平衡得到保证，无边界情况或回归风险。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 0 |
| 🟡 建议 | 1 |
💬 仅评论
torch_npu/distributed/tensor/experimental/_context_parallel/_npu_attention.py
@@ -555,0 +551,4 @@
551+ if torch._C._current_graph_task_id() < 0:
552+ _step_cache_stack.append(
553+ (step_caches, is_causal, attn_output.detach(), merged_lse.detach())
554+ )
🟡 Medium Priority
变更行 551-554：在 _npu_fa_v3_handler 中新增了条件 if torch._C._current_graph_task_id() < 0:，使得 recompute forward 时跳过 _step_cache_stack.append()。
受影响的行为/契约：
_npu_fa_grad_v3_handler在 backward 时出栈并使用这些缓存。
失败模式：
当同时使用梯度检查点（gradient checkpointing / recompute）和 dropout（dropout_p > 0）时：
- 首次 forward：
_current_graph_task_id()返回-1→ 入栈 entry_0（含首次 forward 的 seed/offset/merged_out/merged_lse）。 - Recompute forward：
_current_graph_task_id()返回>= 0→ 跳过入栈。 - Recompute backward：
_npu_fa_grad_v3_handler出栈 → 拿到 entry_0（首次 forward 的缓存）。
但 entry_0 中的 seed/offset 对应首次 forward 的 dropout 随机掩码，merged_out 也是首次 forward 中 dropout 后的输出。而 recompute forward 的 dropout 因 RNG 状态已推进而使用了不同的随机掩码。Backward kernel（npu_fusion_attention_grad）使用 out 参数（即 merged_out）和 seed/offset 来计算梯度，此时两者不匹配，导致 梯度计算结果错误，且该错误是静默的——不会报错，但梯度值不正确。
当 dropout_p == 0 时，首次 forward 与 recompute forward 的计算结果完全相同，entry_0 的缓存有效，此问题不触发。
证据链：
建议：短期方案：添加 dropout_p 的运行时校验，当同时检测到 recompute（task_id >= 0）且 dropout_p > 0 时，记录 warning 或 raise，防止静默梯度错误。长期方案：重构缓存管理，使 recompute forward 的缓存能够被 recompute backward 正确消费，同时清理首次 forward 的过期缓存条目。
7月3日 修改了pull request 的描述
7月6日 添加了label：approved
7月7日 添加了label：lgtm
7月7日 删除了label：ci-pipeline-passed
ascend-robot
7月7日 评论：
7月7日 评论：
The following label is not ready.
ci-pipeline-passed: The ci-pipeline-passed label is expired. Please compile again.
pengjingyou
7月8日 评论：
7月8日 评论：
compile
7月8日 添加了label：ci-pipeline-running
ascend-robot
7月8日 评论：
7月8日 评论：
流水线 PR-pipeline_pytorch#42278 [ commitID：3fa37c03 ] 运行中
>>>代码风格自动修复执行成功（无修复内容）
| 阶段 | 任务名 | 状态 | 详情 |
|---|---|---|---|
| 编译构建 | Build_X86 | 🕚 | >>> |
| Build_ARM | 🕚 | >>> | |
| Build_LibTorch_x86 | 🕚 | >>> | |
| Build_LibTorch_ARM | 🕚 | >>> | |
| Build_X86_torchair | 🕚 | >>> | |
| Build_ARM_torchair | 🕚 | >>> | |
| patch_test | 🕚 | >>> | |
| 恶意代码检查 | Antipoison | 🕚 | >>> |
| 编码安全与规范检查 | CodeCheck | 🕚 | >>> |
| check_error | 🕚 | >>> | |
| CodeCheck_lintrunner | 🕚 | >>> | |
| 开源片段检查 | SCA | 🕚 | >>> |
| 开发者测试 | UT_X86_Part_01 | 🕚 | >>> |
| UT_X86_Part_02 | 🕚 | >>> | |
| UT_ARM_A3_Part_01 | 🕚 | >>> | |
| UT_ARM_A3_Part_02 | 🕚 | >>> | |
| UT_ARM_A2_Part_01 | 🕚 | >>> | |
| UT_ARM_A2_Part_02 | 🕚 | >>> | |
| UT_ARM_A2_Part_03 | 🕚 | >>> | |
| UT_inductor_Part_01 | 🕚 | >>> | |
| UT_inductor_Part_02 | 🕚 | >>> | |
| UT_inductor_Part_03 | 🕚 | >>> | |
| UT_inductor_Part_04 | 🕚 | >>> | |
| UT_DIST_ARM_Part_01 | 🕚 | >>> | |
| UT_DIST_ARM_Part_02 | 🕚 | >>> | |
| UT_DIST_ARM_Part_03 | 🕚 | >>> | |
| UT_DIST_ARM_Part_04 | 🕚 | >>> | |
| UT_ARM_A2_Select_Part_01 | 🕚 | >>> | |
| UT_ARM_A2_Select_Part_02 | 🕚 | >>> | |
| 流水线 | PR-pipeline_pytorch | 🕚 | >>> |
- compile、compile_inductor、compile_torchair : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线
ascend-robot
7月8日 评论：
7月8日 评论：
ascend docs pipeline is running...
7月8日 删除了label：docs-ci-pipeline-success
7月8日 添加了label：docs-ci-pipeline-running
ascend-robot
7月8日 评论：
7月8日 评论：
✅ 跳过 docs ci 检查，没有需要检查的文档文件
7月8日 删除了label：docs-ci-pipeline-running
7月8日 添加了label：docs-ci-pipeline-success
pengjingyou
7月13日 评论：
7月13日 评论：
retry
ascend-robot
7月13日 评论：
7月13日 评论：
重试流水线失败：PR-pipeline_pytorch 仅失败状态流水线支持重试操作。
7月13日 删除了label：ci-pipeline-running
7月13日 添加了label：ci-pipeline-passed
ascend-robot
7月13日 评论：
7月13日 评论：
流水线 PR-pipeline_pytorch#42278 （重试第1次） [ commitID：3fa37c03 ] 已完成
>>>代码风格自动修复执行成功（无修复内容）
| 阶段 | 任务名 | 状态 | 详情 |
|---|---|---|---|
| 编译构建 | Build_X86 | ✅ | >>> |
| Build_ARM | ✅ | >>> | |
| Build_LibTorch_x86 | ✅ | >>> | |
| Build_LibTorch_ARM | ✅ | >>> | |
| Build_X86_torchair | 🛑 | >>> | |
| Build_ARM_torchair | 🛑 | >>> | |
| patch_test | 🛑 | >>> | |
| 恶意代码检查 | Antipoison | ✅ | >>> |
| 编码安全与规范检查 | CodeCheck | ✅ | >>> |
| check_error | ✅ | >>> | |
| CodeCheck_lintrunner | ✅ | >>> | |
| 开源片段检查 | SCA | ✅ | >>> |
| 开发者测试 | UT_X86_Part_01 | 🛑 | >>> |
| UT_X86_Part_02 | 🛑 | >>> | |
| UT_ARM_A3_Part_01 | 🛑 | >>> | |
| UT_ARM_A3_Part_02 | 🛑 | >>> | |
| UT_ARM_A2_Part_01 | ✅ | >>> | |
| UT_ARM_A2_Part_02 | ✅ | >>> | |
| UT_ARM_A2_Part_03 | ✅ | >>> | |
| UT_inductor_Part_01 | 🛑 | >>> | |
| UT_inductor_Part_02 | 🛑 | >>> | |
| UT_inductor_Part_03 | 🛑 | >>> | |
| UT_inductor_Part_04 | 🛑 | >>> | |
| UT_DIST_ARM_Part_01 | ✅ | >>> | |
| UT_DIST_ARM_Part_02 | ✅ | >>> | |
| UT_DIST_ARM_Part_03 | ✅ | >>> | |
| UT_DIST_ARM_Part_04 | ✅ | >>> | |
| UT_ARM_A2_Select_Part_01 | ✅ | >>> | |
| UT_ARM_A2_Select_Part_02 | ✅ | >>> | |
| 流水线 | PR-pipeline_pytorch | ✅ | >>> |
- compile、compile_inductor、compile_torchair : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线
7月13日 合入了pull request
ascend-robot
7月13日 评论：
7月13日 评论：
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
ascend-robot
7月13日 评论：
7月13日 评论：
流水线 pytorch_gitcode_PR_multiVersion#12446 [ commitID：3fa37c03 ] 已完成
【合入来源】
【修改方案】
_npu_fa_v3_handler()：_npu_fa_grad_v3_handler()：这个设计需要保证 forward 和 backward 严格 1:1：每 push 一次，必须有且仅有一次对应的 backward pop。
但重计算场景下，forward 和 backward 2:1，导致全局变量_step_cache_stack 存在内存泄漏
用
torch._C._current_graph_task_id()标记 recompute，修复全局变量 _step_cache_stack 的push/pop平衡在保存forward中间结果时，进行是否为重计算的判断，如果当前执行的是recompute forward，则不需要执行push操作
【资料变更】
【接口变更】
【功能验证】
执行训练任务，开启cp和重计算，内存占用稳定：
context_parallel ut结果：
【CheckList】