Add counter-based PTA OOM injection for P/D inference - #6
Draft
yjyang62 wants to merge 1 commit into
Draft
Conversation
Implement PtaOomInjector with a shared counter across tensor allocations and NPU op executions so OOM faults can be triggered during both prefill (P) and decode (D) phases, including graph-replay inference on large models like Qwen-235B. - PTA_OOM_TRIGGER_COUNT: allocation event threshold (default 30004) - PTA_OOM_OP_TRIGGER_COUNT: optional separate op threshold - PTA_OOM_DURING_CAPTURE=1: allow injection during graph capture - HCCL OOM (HCCL_OOM_TRIGGER_COUNT) remains for P2P communication faults Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Restores and improves PTA OOM fault injection so it works during inference on large models (e.g. Qwen-235B) in both Prefill (P) and Decode (D) phases.
The previous branch only incremented
g_malloc_call_countin the internal caching allocatormalloc()without throwing, so OOM faults could not be triggered. This change introduces a dedicatedPtaOomInjectorwith a shared counter across two injection points:NPUCachingAllocator::allocate) — covers prefill and initial memory growthOpCommand::Run/RunOpApi*) — covers decode and graph-replay inference where allocations are reusedHCCL OOM injection (
HCCL_OOM_TRIGGER_COUNT, default 6000) remains unchanged for P2P communication faults between P and D nodes.Environment variables
PTA_OOM_TRIGGER_COUNT30004<=0to disable alloc countingPTA_OOM_OP_TRIGGER_COUNTPTA_OOM_TRIGGER_COUNTPTA_OOM_DURING_CAPTURE01to allow injection during NPU graph captureHCCL_OOM_TRIGGER_COUNT6000Usage example (Qwen-235B)
Files changed
torch_npu/csrc/core/npu/PtaOomInjector.h— new injector APItorch_npu/csrc/core/npu/PtaOomInjector.cpp— counter logic and env var parsingtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp— hook alloc OOM, remove broken countertorch_npu/csrc/framework/OpCommand.cpp— hook op OOM on all Run paths