diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 2a4049a0dc..fad36bae04 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,9 @@ 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 = 30004; +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 +122,40 @@ 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; +} + +void maybeThrowPtaOom(int device, size_t size) +{ + if (g_pta_oom_triggered.load() || size == 0) { + return; + } + + const int64_t trigger_count = getPtaOomTriggerCount(); + if (trigger_count <= 0) { + 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) + " normal tensor 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 +1192,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 +3553,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;