From 689c61f8a4af107c5415a5b88a8241a14e948022 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 12 Jun 2026 08:38:25 +0000 Subject: [PATCH 1/6] Inject HCCL OOM after warmup Co-authored-by: yjyang62 --- .../csrc/distributed/ProcessGroupHCCL.cpp | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp b/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp index 005f98fbdf..8eded16cf4 100644 --- a/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp +++ b/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -75,6 +76,8 @@ using hcclUs = std::chrono::steady_clock::time_point; constexpr int32_t MAX_GROUP_NAME_LEN = 128; constexpr int32_t NSLB_JOBID_OFFSET = 32; static constexpr int CoalActive = 0x01, CoalColl = 0x02, CoalP2P = 0x04; +static constexpr int64_t kDefaultHcclOomTriggerCount = 5000; +static std::atomic g_hccl_oom_call_count{0}; // HCCL ReduceOp mapping std::map hcclOp = { @@ -90,6 +93,35 @@ std::map unsupportedOp = { {c10d::ReduceOp::BXOR, "BXOR"} }; +int64_t getHcclOomTriggerCount() +{ + const static int64_t trigger_count = []() -> int64_t { + char *env_val = c10_npu::option::get_and_log_env("HCCL_OOM_TRIGGER_COUNT"); + return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : kDefaultHcclOomTriggerCount; + }(); + return trigger_count; +} + +void maybeThrowHcclOom(c10d::OpType opType, c10_npu::CaptureStatus capture_status) +{ + if (capture_status != c10_npu::CaptureStatus::None) { + return; + } + + const int64_t trigger_count = getHcclOomTriggerCount(); + if (trigger_count <= 0) { + return; + } + + const int64_t current_count = ++g_hccl_oom_call_count; + if (current_count > trigger_count && current_count < trigger_count + 2) { + auto retmsg = std::string("HCCL out of memory. Injected OOM after ") + + std::to_string(current_count) + " HCCL operations, op type is " + + opTypeToString(opType) + "."; + TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str()); + } +} + bool nslb_is_end = false; std::string device_error_msg; bool force_stop_error_flag = false; @@ -3785,6 +3817,7 @@ c10::intrusive_ptr ProcessGroupHCCL::collective( auto key = getKeyFromDevices(devices); HcclCommConfig config = createHcclCommConfigWithOptions(); std::vector> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config); + maybeThrowHcclOom(opType, capture_status); auto& hcclStreams = hcclStreams_[key]; syncStreams(devices, hcclEvents_[key], hcclStreams); @@ -4000,6 +4033,7 @@ c10::intrusive_ptr ProcessGroupHCCL::collectiveCoalesced( NPU_CHECK_ERROR(c10_npu::SetDevice(devices[0].index())); HcclCommConfig config = createHcclCommConfigWithOptions(); std::vector> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config); + maybeThrowHcclOom(opType, capture_status); auto& hcclStreams = hcclStreams_[key]; syncStreams(devices, hcclEvents_[key], hcclStreams); @@ -4236,6 +4270,7 @@ c10::intrusive_ptr ProcessGroupHCCL::pointToPoint( key = getKeyFromDevices(devices); hcclComms = getHCCLComm(key, devices); } + maybeThrowHcclOom(opType, capture_status); // Bump the logical operation counter regardless of whether this op is // coalesced or individual From 251278751ec14c9bc4f19c5e0a37263e4c59ac44 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 12 Jun 2026 09:34:26 +0000 Subject: [PATCH 2/6] Generate ATen bindings for CMake builds Co-authored-by: yjyang62 --- CMakeLists.txt | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 71c7907970..018b0572c8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -223,7 +223,38 @@ if(DEFINED PYTHON_INCLUDE_DIR) else() message(FATAL_ERROR "Cannot find installed Python head file directory") endif() - + +if(DEFINED PYTHON_EXECUTABLE) + set(TORCHNPU_PYTHON_EXECUTABLE ${PYTHON_EXECUTABLE}) +else() + find_package(Python3 COMPONENTS Interpreter REQUIRED) + set(TORCHNPU_PYTHON_EXECUTABLE ${Python3_EXECUTABLE}) +endif() + +if(NOT DEFINED TORCH_VERSION) + execute_process( + COMMAND ${TORCHNPU_PYTHON_EXECUTABLE} -c "import torch; print(torch.__version__.split('+')[0])" + RESULT_VARIABLE _TORCH_VERSION_RESULT + OUTPUT_VARIABLE TORCH_VERSION + OUTPUT_STRIP_TRAILING_WHITESPACE) + if(NOT _TORCH_VERSION_RESULT EQUAL 0) + message(FATAL_ERROR "Cannot determine TORCH_VERSION for ATen binding generation") + endif() + add_definitions(-DPYTORCH_NPU_VERSION="${TORCH_VERSION}") +endif() + +set(TORCHNPU_NATIVE_FUNCTIONS_HEADER + ${PROJECT_SOURCE_DIR}/torch_npu/csrc/aten/NPUNativeFunctions.h) +if(NOT EXISTS ${TORCHNPU_NATIVE_FUNCTIONS_HEADER}) + execute_process( + COMMAND bash ${PROJECT_SOURCE_DIR}/generate_code.sh ${TORCHNPU_PYTHON_EXECUTABLE} ${TORCH_VERSION} + WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} + RESULT_VARIABLE _GENERATE_CODE_RESULT) + if(NOT _GENERATE_CODE_RESULT EQUAL 0) + message(FATAL_ERROR "Failed to generate ATen bindings for torch_npu") + endif() +endif() + # sources set(ATEN_SRCS) set(CORE_SRCS) From fc7761a7bf4d0dcf4cf5a7f77cf1a5bb1b965320 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 06:49:15 +0000 Subject: [PATCH 3/6] Make PTA OOM injection controllable Co-authored-by: yjyang62 --- .../csrc/core/npu/NPUCachingAllocator.cpp | 85 +++++++++++++++++-- 1 file changed, 78 insertions(+), 7 deletions(-) diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 2a4049a0dc..529b0b4194 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -1,5 +1,7 @@ #include +#include #include +#include #include #include #include @@ -23,6 +25,8 @@ #include "torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h" #include "torch_npu/csrc/core/npu/NPURecovery.h" #include "torch_npu/csrc/core/npu/NPUGuard.h" +#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h" +#include "torch_npu/csrc/core/npu/register/OptionsManager.h" #include "NPUBlockHandle.h" #include "torch_npu/csrc/core/npu/NpuVariables.h" #include "torch_npu/csrc/core/npu/GetCANNInfo.h" @@ -107,7 +111,10 @@ const std::string kMinDriverVersion = "25.0.RC1"; // minimum driver version const std::string kCannModule = "CANN"; // cann module name constexpr int kPrecision = 4; // precision of the memory usage information constexpr size_t kLazyQuerySize = 512; // lazy query event size -static int64_t g_malloc_call_count = 0; +constexpr int64_t kDefaultPtaOomTriggerCount = 1000; +constexpr size_t kDefaultPtaOomMinAllocSize = kSmallSize; +static std::atomic g_pta_oom_candidate_count{0}; +static std::atomic g_pta_oom_triggered{false}; static char SHAREABLE_HANDLE_VERSION = 1; enum ShareableHandleType : char { SHAREABLE_NPU_MALLOC = 'c', @@ -116,6 +123,74 @@ enum ShareableHandleType : char { using StatTypes = std::array(StatType::NUM_TYPES)>; +int64_t getPtaOomTriggerCount() +{ + const static int64_t trigger_count = []() -> int64_t { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_COUNT"); + return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : kDefaultPtaOomTriggerCount; + }(); + return trigger_count; +} + +size_t getPtaOomMinAllocSize() +{ + const static size_t min_alloc_size = []() -> size_t { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_MIN_ALLOC_BYTES"); + int64_t env_flag = (env_val != nullptr) ? strtol(env_val, nullptr, 10) : + static_cast(kDefaultPtaOomMinAllocSize); + return env_flag > 0 ? static_cast(env_flag) : 0; + }(); + return min_alloc_size; +} + +int64_t getPtaOomTargetDevice() +{ + const static int64_t target_device = []() -> int64_t { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_DEVICE"); + return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : -1; + }(); + return target_device; +} + +bool isPtaOomEnabled() +{ + const static bool enabled = []() -> bool { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_ENABLE"); + return (env_val == nullptr) || (strtol(env_val, nullptr, 10) != 0); + }(); + return enabled; +} + +void maybeThrowPtaOom(int device, size_t size) +{ + if (!isPtaOomEnabled() || g_pta_oom_triggered.load()) { + return; + } + + const int64_t trigger_count = getPtaOomTriggerCount(); + if (trigger_count <= 0 || size == 0 || size < getPtaOomMinAllocSize()) { + return; + } + + const int64_t target_device = getPtaOomTargetDevice(); + if (target_device >= 0 && target_device != device) { + return; + } + + if (c10_npu::currentStreamCaptureStatus() != c10_npu::CaptureStatus::None) { + return; + } + + const int64_t current_count = ++g_pta_oom_candidate_count; + if (current_count > trigger_count && current_count < trigger_count + 2) { + g_pta_oom_triggered.store(true); + auto retmsg = std::string("NPU out of memory. Injected PTA OOM after ") + + std::to_string(current_count) + " eligible allocations. Tried to allocate " + + format_size(size) + " on NPU " + std::to_string(device) + "."; + TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str()); + } +} + void update_stat(Stat &stat, int64_t amount) { stat.current += amount; @@ -1152,12 +1227,7 @@ class DeviceCachingAllocator { // Thus, do not call a public method from another public method. Block *malloc(int device, size_t orig_size, aclrtStream stream, uint8_t allocator_type = 0) - { - g_malloc_call_count++; - auto retmsg = std::string("NPU out of memory. Tried to allocate more than 1EB memory."); - if (g_malloc_call_count > 30004 && g_malloc_call_count < 30006) { - TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str()); - } + { TORCH_NPU_MEMORY_LOGD("Allocating memory: size=%zu, device=%d", orig_size, device); // done outside the lock because we don't know what locks the recorder needs // to have... @@ -3518,6 +3588,7 @@ class NpuCachingAllocator : public NPUAllocator { int device = 0; NPU_CHECK_ERROR(c10_npu::GetDevice(&device)); + maybeThrowPtaOom(device, size); LazySetDevice(device); void *devPtr = nullptr; void (*deleteFunc)(void *) = &local_raw_delete; From 96115e757dc94e55f1d24b078cd58c124f2a4c86 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 07:11:15 +0000 Subject: [PATCH 4/6] Move HCCL OOM injection after API dispatch Co-authored-by: yjyang62 --- .../csrc/distributed/ProcessGroupHCCL.cpp | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp b/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp index 8eded16cf4..67af53c54f 100644 --- a/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp +++ b/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp @@ -115,9 +115,10 @@ void maybeThrowHcclOom(c10d::OpType opType, c10_npu::CaptureStatus capture_statu const int64_t current_count = ++g_hccl_oom_call_count; if (current_count > trigger_count && current_count < trigger_count + 2) { - auto retmsg = std::string("HCCL out of memory. Injected OOM after ") + - std::to_string(current_count) + " HCCL operations, op type is " + - opTypeToString(opType) + "."; + auto retmsg = std::string("HCCL function error: Failed to allocate memory. " + "Injected HCCL OOM after ") + std::to_string(current_count) + + " HCCL operations, op type is " + opTypeToString(opType) + + ", error code is " + std::to_string(HCCL_E_OOM) + " " + DIST_ERROR(ErrCode::HCCL) + "."; TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str()); } } @@ -3817,7 +3818,6 @@ c10::intrusive_ptr ProcessGroupHCCL::collective( auto key = getKeyFromDevices(devices); HcclCommConfig config = createHcclCommConfigWithOptions(); std::vector> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config); - maybeThrowHcclOom(opType, capture_status); auto& hcclStreams = hcclStreams_[key]; syncStreams(devices, hcclEvents_[key], hcclStreams); @@ -3966,7 +3966,9 @@ c10::intrusive_ptr ProcessGroupHCCL::collective( c10_npu::SetStreamResLimit(hcclStream, c10_npu::acl::ACL_RT_DEV_RES_VECTOR_CORE, current_aiv_num); } hcclUs startut = std::chrono::steady_clock::now(); - HCCL_CHECK_ERROR(fn(inputs[i], outputs[i], hcclComms[i]->getHcclComm(), hcclStream, work->is_dispatched), opTypeToString(opType).c_str()); + auto hcclResult = fn(inputs[i], outputs[i], hcclComms[i]->getHcclComm(), hcclStream, work->is_dispatched); + HCCL_CHECK_ERROR(hcclResult, opTypeToString(opType).c_str()); + maybeThrowHcclOom(opType, capture_status); if (c10_npu::option::OptionsManager::GetMultiStreamMemoryReuse() == c10_npu::option::ERASE_RECORD_STREAM) { work->recorded_outputs_.push_back( std::make_pair(outputs[i].storage().getWeakStorageImpl(), hcclStream)); @@ -4033,7 +4035,6 @@ c10::intrusive_ptr ProcessGroupHCCL::collectiveCoalesced( NPU_CHECK_ERROR(c10_npu::SetDevice(devices[0].index())); HcclCommConfig config = createHcclCommConfigWithOptions(); std::vector> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config); - maybeThrowHcclOom(opType, capture_status); auto& hcclStreams = hcclStreams_[key]; syncStreams(devices, hcclEvents_[key], hcclStreams); @@ -4190,7 +4191,9 @@ c10::intrusive_ptr ProcessGroupHCCL::collectiveCoalesced( c10_npu::SetStreamResLimit(hcclStream, c10_npu::acl::ACL_RT_DEV_RES_VECTOR_CORE, current_aiv_num); } hcclUs startut = std::chrono::steady_clock::now(); - HCCL_CHECK_ERROR(fn(inputs[i], outputs[i], hcclComms[0]->getHcclComm(), hcclStream, work->is_dispatched), opTypeToString(opType).c_str()); + auto hcclResult = fn(inputs[i], outputs[i], hcclComms[0]->getHcclComm(), hcclStream, work->is_dispatched); + HCCL_CHECK_ERROR(hcclResult, opTypeToString(opType).c_str()); + maybeThrowHcclOom(opType, capture_status); if (c10_npu::option::OptionsManager::GetMultiStreamMemoryReuse() == c10_npu::option::ERASE_RECORD_STREAM) { work->recorded_outputs_.push_back( std::make_pair(outputs[i].storage().getWeakStorageImpl(), hcclStream)); @@ -4270,7 +4273,6 @@ c10::intrusive_ptr ProcessGroupHCCL::pointToPoint( key = getKeyFromDevices(devices); hcclComms = getHCCLComm(key, devices); } - maybeThrowHcclOom(opType, capture_status); // Bump the logical operation counter regardless of whether this op is // coalesced or individual @@ -4441,13 +4443,15 @@ c10::intrusive_ptr ProcessGroupHCCL::pointToPoint( }; at_npu::native::OpCommand::RunOpApiV3("hcclGroupStart", hccl_call); } - HCCL_CHECK_ERROR(fn(tensors[i], hcclComms[i]->getHcclComm(), hcclStream, is_dispatched, p2pTargetRank), opTypeToString(opType).c_str()); + auto hcclResult = fn(tensors[i], hcclComms[i]->getHcclComm(), hcclStream, is_dispatched, p2pTargetRank); if (coalescing_state_) { auto hccl_call = [this]() -> HcclResult { return hcclGroupEnd(); }; at_npu::native::OpCommand::RunOpApiV3("hcclGroupEnd", hccl_call); } + HCCL_CHECK_ERROR(hcclResult, opTypeToString(opType).c_str()); + maybeThrowHcclOom(opType, capture_status); } } From e080e199ff8da86d38b24ca986fc2f50f21d52f8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 07:53:55 +0000 Subject: [PATCH 5/6] Narrow PTA OOM recovery trigger scope Co-authored-by: yjyang62 --- .../csrc/core/npu/NPUCachingAllocator.cpp | 42 ++++++++++++++++++- 1 file changed, 40 insertions(+), 2 deletions(-) diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 529b0b4194..55a0596f42 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -112,7 +112,7 @@ const std::string kCannModule = "CANN"; // cann module name constexpr int kPrecision = 4; // precision of the memory usage information constexpr size_t kLazyQuerySize = 512; // lazy query event size constexpr int64_t kDefaultPtaOomTriggerCount = 1000; -constexpr size_t kDefaultPtaOomMinAllocSize = kSmallSize; +constexpr size_t kDefaultPtaOomMinAllocSize = 64 * 1024 * 1024; static std::atomic g_pta_oom_candidate_count{0}; static std::atomic g_pta_oom_triggered{false}; static char SHAREABLE_HANDLE_VERSION = 1; @@ -143,6 +143,16 @@ size_t getPtaOomMinAllocSize() return min_alloc_size; } +size_t getPtaOomMaxAllocSize() +{ + const static size_t max_alloc_size = []() -> size_t { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_MAX_ALLOC_BYTES"); + int64_t env_flag = (env_val != nullptr) ? strtol(env_val, nullptr, 10) : -1; + return env_flag > 0 ? static_cast(env_flag) : 0; + }(); + return max_alloc_size; +} + int64_t getPtaOomTargetDevice() { const static int64_t target_device = []() -> int64_t { @@ -152,11 +162,29 @@ int64_t getPtaOomTargetDevice() return target_device; } +int64_t getPtaOomTargetRank() +{ + const static int64_t target_rank = []() -> int64_t { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_RANK"); + if (env_val != nullptr) { + return strtol(env_val, nullptr, 10); + } + return std::getenv("RANK") != nullptr ? 0 : -1; + }(); + return target_rank; +} + +int64_t getCurrentRank() +{ + char *rank = std::getenv("RANK"); + return rank != nullptr ? strtol(rank, nullptr, 10) : -1; +} + bool isPtaOomEnabled() { const static bool enabled = []() -> bool { char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_ENABLE"); - return (env_val == nullptr) || (strtol(env_val, nullptr, 10) != 0); + return (env_val != nullptr) && (strtol(env_val, nullptr, 10) != 0); }(); return enabled; } @@ -172,11 +200,21 @@ void maybeThrowPtaOom(int device, size_t size) return; } + const size_t max_alloc_size = getPtaOomMaxAllocSize(); + if (max_alloc_size != 0 && size > max_alloc_size) { + return; + } + const int64_t target_device = getPtaOomTargetDevice(); if (target_device >= 0 && target_device != device) { return; } + const int64_t target_rank = getPtaOomTargetRank(); + if (target_rank >= 0 && target_rank != getCurrentRank()) { + return; + } + if (c10_npu::currentStreamCaptureStatus() != c10_npu::CaptureStatus::None) { return; } From 3f5b2f0bf9d1eb095dd5c947a94b15908694841b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 14 Jun 2026 07:55:54 +0000 Subject: [PATCH 6/6] Enable PTA OOM on all ranks by default Co-authored-by: yjyang62 --- torch_npu/csrc/core/npu/NPUCachingAllocator.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 55a0596f42..f656b01519 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -169,7 +169,7 @@ int64_t getPtaOomTargetRank() if (env_val != nullptr) { return strtol(env_val, nullptr, 10); } - return std::getenv("RANK") != nullptr ? 0 : -1; + return -1; }(); return target_rank; }