已合并
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
已合并
共 1 个文件变更+6-3
| @@ -546,9 +546,12 @@ def _npu_fa_v3_handler(op_call, args, kwargs): | |||
| 546 | 546 | ||
| 547 | # Push (step_caches, is_causal, merged_out, merged_lse) for backward handler. | 547 | # Push (step_caches, is_causal, merged_out, merged_lse) for backward handler. |
| 548 | # Backward needs merged_out to compute D = sum(dout * O_merged); per-step raw output is not enough. | 548 | # Backward needs merged_out to compute D = sum(dout * O_merged); per-step raw output is not enough. |
| 549 | - _step_cache_stack.append( | 549 | + # These forward intermediates need transformation before backward consumes them, so save them here. |
| 550 | - (step_caches, is_causal, attn_output.detach(), merged_lse.detach()) | 550 | + # Skip this during recompute forward: no need to preserve them for the recompute pass. |
| 551 | - ) | 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 | + ) | ||
 | |||
| 552 | 555 | ||
| 553 | # Build v3 6-tuple. softmax_max/sum come from the final step; op_plugin typically only uses attn_output. | 556 | # Build v3 6-tuple. softmax_max/sum come from the final step; op_plugin typically only uses attn_output. |
| 554 | B, N, S, D = attn_output.shape | 557 | B, N, S, D = attn_output.shape |
🟡 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）时：_current_graph_task_id()返回-1→ 入栈 entry_0（含首次 forward 的 seed/offset/merged_out/merged_lse）。_current_graph_task_id()返回>= 0→ 跳过入栈。_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 的过期缓存条目。