feat: inject HCCL OOM after warmup

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)

torch_npu/csrc/core/npu/NPUCachingAllocator.cpp

@@ -1,5 +1,7 @@
 #include <algorithm>
+#include <atomic>
 #include <bitset>
+#include <cstdlib>
 #include <deque>
 #include <map>
 #include <memory>
@@ -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 = 64 * 1024 * 1024;
+static std::atomic<int64_t> g_pta_oom_candidate_count{0};
+static std::atomic<bool> g_pta_oom_triggered{false};
 static char SHAREABLE_HANDLE_VERSION = 1;
 enum ShareableHandleType : char {
     SHAREABLE_NPU_MALLOC = 'c',
@@ -116,6 +123,112 @@ enum ShareableHandleType : char {
 
 using StatTypes = std::array<bool, static_cast<size_t>(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<int64_t>(kDefaultPtaOomMinAllocSize);
+        return env_flag > 0 ? static_cast<size_t>(env_flag) : 0;
+    }();
+    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<size_t>(env_flag) : 0;
+    }();
+    return max_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;
+}
+
+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 -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 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 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;
+    }
+
+    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 +1265,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 +3626,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;

torch_npu/csrc/distributed/ProcessGroupHCCL.cpp

@@ -1,5 +1,6 @@
 #include <ATen/record_function.h>
 #include <algorithm>
+#include <atomic>
 #include <map>
 #include <tuple>
 #include <unordered_set>
@@ -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<int64_t> g_hccl_oom_call_count{0};
 
 // HCCL ReduceOp mapping
 std::map<c10d::ReduceOp, HcclReduceOp> hcclOp = {
@@ -90,6 +93,36 @@ std::map<c10d::ReduceOp, std::string> 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 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());
+    }
+}
+
 bool nslb_is_end = false;
 std::string device_error_msg;
 bool force_stop_error_flag = false;
@@ -3933,7 +3966,9 @@ c10::intrusive_ptr<c10d::Work> 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));
@@ -4156,7 +4191,9 @@ c10::intrusive_ptr<c10d::Work> 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));
@@ -4406,13 +4443,15 @@ c10::intrusive_ptr<c10d::Work> 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);
         }
     }