Skip to content

Add counter-based PTA OOM injection for P/D inference - #6

Draft
yjyang62 wants to merge 1 commit into
cursor/pta-oom-2.9.1-fbcafrom
cursor/inference-pd-oom-counter-a4ee
Draft

Add counter-based PTA OOM injection for P/D inference#6
yjyang62 wants to merge 1 commit into
cursor/pta-oom-2.9.1-fbcafrom
cursor/inference-pd-oom-counter-a4ee

Conversation

@yjyang62

Copy link
Copy Markdown
Owner

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_count in the internal caching allocator malloc() without throwing, so OOM faults could not be triggered. This change introduces a dedicated PtaOomInjector with a shared counter across two injection points:

  1. Tensor allocation (NPUCachingAllocator::allocate) — covers prefill and initial memory growth
  2. NPU op execution (OpCommand::Run/RunOpApi*) — covers decode and graph-replay inference where allocations are reused

HCCL OOM injection (HCCL_OOM_TRIGGER_COUNT, default 6000) remains unchanged for P2P communication faults between P and D nodes.

Environment variables

Variable Default Description
PTA_OOM_TRIGGER_COUNT 30004 Shared PTA event threshold for allocations; set <=0 to disable alloc counting
PTA_OOM_OP_TRIGGER_COUNT same as above Optional separate threshold for op events; defaults to PTA_OOM_TRIGGER_COUNT
PTA_OOM_DURING_CAPTURE 0 Set to 1 to allow injection during NPU graph capture
HCCL_OOM_TRIGGER_COUNT 6000 HCCL communication OOM threshold (unchanged)

Usage example (Qwen-235B)

# Prefill node (P) — allocation-heavy phase
export PTA_OOM_TRIGGER_COUNT=30004

# Decode node (D) — op-heavy phase with graph replay
export PTA_OOM_TRIGGER_COUNT=5000   # tune based on warmup op count

# P2P communication fault between P and D
export HCCL_OOM_TRIGGER_COUNT=6000

Files changed

  • torch_npu/csrc/core/npu/PtaOomInjector.h — new injector API
  • torch_npu/csrc/core/npu/PtaOomInjector.cpp — counter logic and env var parsing
  • torch_npu/csrc/core/npu/NPUCachingAllocator.cpp — hook alloc OOM, remove broken counter
  • torch_npu/csrc/framework/OpCommand.cpp — hook op OOM on all Run paths
Open in Web Open in Cursor 

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants