diff --git a/docs/oom_fault_injection.md b/docs/oom_fault_injection.md new file mode 100644 index 0000000000..f18207cd9d --- /dev/null +++ b/docs/oom_fault_injection.md @@ -0,0 +1,50 @@ +# Synthetic OOM fault injection + +This repository provides synthetic OOM injection hooks for recovery validation. +The op and allocator hooks are disabled by default; the HCCL hook keeps its +existing count-based default. All hooks throw `OutOfMemoryError` without +consuming device memory, so a serving layer can exercise its fast recovery path +without leaving real HBM pressure behind. + +## Recommended PD inference flow + +Use op-level automatic triggering for both prefill and decode nodes. Add these +values to the MindIE launch environment: + +```text +HCCL_OOM_COUNT=1 +NPU_ALLOCATOR_OOM_TRIGGER_COUNT=0 +HCCL_OOM_TRIGGER_COUNT=0 +``` + +Start the PD inference service normally with this environment. No extra command +is required during inference. The hook counts NPU operations inside each +process; once the configured count is reached, each rank throws a synthetic OOM +once per local device by default. + +This is the preferred mode for Qwen-235B PD separation tests because it does not +depend on hitting a specific HCCL operation and does not require real HBM +pressure. Increase `HCCL_OOM_COUNT` if the fault must be delayed until +after warmup requests. + +## Environment variables + +| Variable | Scope | Default | Meaning | +| --- | --- | --- | --- | +| `NPU_OOM_TRIGGER_FILE` | op, allocator, HCCL fallback | unset | Shared trigger file. If it exists, inject OOM. | +| `NPU_OP_OOM_TRIGGER_FILE` | op | unset | Op-only trigger file. Overrides the shared file for op injection. | +| `NPU_ALLOCATOR_OOM_TRIGGER_FILE` | allocator | unset | Allocator-only trigger file. Overrides the shared file for allocator injection. | +| `HCCL_OOM_TRIGGER_FILE` | HCCL | unset | HCCL-only trigger file. Overrides the shared file for HCCL injection. | +| `NPU_OOM_TRIGGER_COUNT` | op, allocator, HCCL fallback | `0` for op/allocator, HCCL keeps its legacy default if unset | Shared automatic trigger count. | +| `HCCL_OOM_COUNT` | op | `0` | Inject when per-device op count reaches this value. `0` disables count injection. | +| `NPU_ALLOCATOR_OOM_TRIGGER_COUNT` | allocator | `0` | Inject when per-device allocation count reaches this value. `0` disables count injection. | +| `HCCL_OOM_TRIGGER_COUNT` | HCCL | `6000` | Inject when HCCL call count reaches this value. `0` disables count injection. | +| `NPU_OOM_TRIGGER_MODE` | op, allocator, HCCL fallback | `once` | Shared one-shot or repeated mode. | +| `NPU_OP_OOM_TRIGGER_MODE` | op | `once` | Set to `always`, `repeat`, or `1` to inject repeatedly. | +| `NPU_ALLOCATOR_OOM_TRIGGER_MODE` | allocator | `once` | Set to `always`, `repeat`, or `1` to inject repeatedly. | +| `HCCL_OOM_TRIGGER_MODE` | HCCL | `once` | Set to `always`, `repeat`, or `1` to inject repeatedly. | + +For fast recovery validation, keep the default one-shot mode. Repeated mode is +mainly useful for stress tests where the serving process is expected to restart. +Trigger files remain available for manual debugging, but automatic count-based +triggering is the recommended PD inference path. diff --git a/test/npu/test_fault_mode.py b/test/npu/test_fault_mode.py index 90aefdffc2..1befd1970f 100644 --- a/test/npu/test_fault_mode.py +++ b/test/npu/test_fault_mode.py @@ -293,6 +293,43 @@ def test_aclrtMallocAlign32(self): x = torch.randn(2000, 2000, 200, 20, device="npu:0") y = torch.randn(2000, 2000, 200, 20, device="npu:0") + def test_synthetic_allocator_oom_inject(self): + command = [ + "python", + "-c", + "import os; " + "os.environ['NPU_ALLOCATOR_OOM_TRIGGER_COUNT'] = '1'; " + "import torch; import torch_npu; " + "torch.empty((1,), device='npu:0')", + ] + process = subprocess.Popen(command, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + out, error = process.communicate(timeout=1800) + process.stderr.close() + process.stdout.close() + process.terminate() + process.wait() + self.assertNotEqual(process.returncode, 0) + self.assertIn("Injected NPU allocator OOM", error) + + def test_synthetic_shared_oom_count_inject(self): + command = [ + "python", + "-c", + "import os; import torch; import torch_npu; " + "x = torch.empty((1,), device='npu:0'); " + "os.environ['NPU_OOM_TRIGGER_COUNT'] = '1'; " + "(x + 1).cpu()", + ] + process = subprocess.Popen(command, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + out, error = process.communicate(timeout=1800) + process.stderr.close() + process.stdout.close() + process.terminate() + process.wait() + self.assertNotEqual(process.returncode, 0) + self.assertIn("Injected NPU", error) + self.assertIn("OOM", error) + if __name__ == "__main__": run_tests() diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 783d27969a..b5fa8c1271 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -8,6 +8,8 @@ #include #include #include +#include +#include #include #include @@ -27,6 +29,7 @@ #include "torch_npu/csrc/core/npu/NpuVariables.h" #include "torch_npu/csrc/core/npu/GetCANNInfo.h" #include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h" +#include "torch_npu/csrc/core/npu/register/OptionsManager.h" #include "torch_npu/csrc/core/npu/NPUEvent.h" #include "torch_npu/csrc/profiler/npu_profiler.h" #ifndef BUILD_LIBTORCH @@ -108,6 +111,8 @@ 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; +static std::map g_allocator_oom_call_counts; +static std::set g_allocator_oom_triggered_devices; static char SHAREABLE_HANDLE_VERSION = 1; enum ShareableHandleType : char { SHAREABLE_NPU_MALLOC = 'c', @@ -116,6 +121,86 @@ enum ShareableHandleType : char { using StatTypes = std::array(StatType::NUM_TYPES)>; +int64_t getAllocatorOomTriggerCount() +{ + const static int64_t trigger_count = []() -> int64_t { + char *env_val = c10_npu::option::get_and_log_env("NPU_ALLOCATOR_OOM_TRIGGER_COUNT"); + if (env_val == nullptr) { + env_val = c10_npu::option::get_and_log_env("NPU_OOM_TRIGGER_COUNT"); + } + return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : 0; + }(); + return trigger_count; +} + +bool isAllocatorOomTriggerRepeatable() +{ + const static bool repeatable = []() -> bool { + char *env_val = c10_npu::option::get_and_log_env("NPU_ALLOCATOR_OOM_TRIGGER_MODE"); + if (env_val == nullptr) { + env_val = c10_npu::option::get_and_log_env("NPU_OOM_TRIGGER_MODE"); + } + if (env_val == nullptr) { + return false; + } + std::string mode(env_val); + std::transform(mode.begin(), mode.end(), mode.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + return mode == "always" || mode == "repeat" || mode == "1"; + }(); + return repeatable; +} + +std::string getNpuOomTriggerFile() +{ + const static std::string trigger_file = []() -> std::string { + char *env_val = c10_npu::option::get_and_log_env("NPU_ALLOCATOR_OOM_TRIGGER_FILE"); + if (env_val != nullptr) { + return std::string(env_val); + } + env_val = c10_npu::option::get_and_log_env("NPU_OOM_TRIGGER_FILE"); + return (env_val != nullptr) ? std::string(env_val) : std::string(); + }(); + return trigger_file; +} + +bool isTriggerFilePresent(const std::string &path) +{ + if (path.empty()) { + return false; + } + std::ifstream file(path); + return file.good(); +} + +bool shouldThrowAllocatorOom(int device) +{ + const auto trigger_file = getNpuOomTriggerFile(); + const int64_t trigger_count = getAllocatorOomTriggerCount(); + if (trigger_file.empty() && trigger_count <= 0) { + return false; + } + + const bool file_triggered = isTriggerFilePresent(trigger_file); + const int64_t current_count = ++g_allocator_oom_call_counts[device]; + const bool count_triggered = trigger_count > 0 && current_count >= trigger_count; + if (!file_triggered && !count_triggered) { + return false; + } + if (!isAllocatorOomTriggerRepeatable() && g_allocator_oom_triggered_devices.count(device) > 0) { + return false; + } + g_allocator_oom_triggered_devices.insert(device); + return true; +} + +std::string getInjectedAllocatorOomMessage(int device) +{ + return std::string("NPU out of memory. Failed to allocate memory. Injected NPU allocator OOM on device ") + + std::to_string(device) + " after " + std::to_string(g_allocator_oom_call_counts[device]) + + " allocation attempts. This is a synthetic OOM fault for recovery validation."; +} + void update_stat(Stat &stat, int64_t amount) { stat.current += amount; @@ -1167,6 +1252,11 @@ class DeviceCachingAllocator { NPU_CHECK_ERROR(c10_npu::GetDevice(&device)); TORCH_NPU_MEMORY_LOGD("Using device: %d", device); } + if (shouldThrowAllocatorOom(device)) { + auto injected_oom_msg = getInjectedAllocatorOomMessage(device); + TORCH_NPU_MEMORY_LOGE("%s", injected_oom_msg.c_str()); + TORCH_CHECK_WITH(OutOfMemoryError, false, injected_oom_msg.c_str()); + } if (!CachingAllocatorConfig::multi_stream_lazy_reclaim() && C10_LIKELY(captures_underway.empty())) { // Processes end-of-life events for outstanding allocations used on diff --git a/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp b/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp index d6de8a8337..ce94b458a5 100644 --- a/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp +++ b/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #ifndef BUILD_LIBTORCH @@ -78,6 +79,7 @@ constexpr int32_t NSLB_JOBID_OFFSET = 32; static constexpr int CoalActive = 0x01, CoalColl = 0x02, CoalP2P = 0x04; static constexpr int64_t kDefaultHcclOomTriggerCount = 6000; static std::atomic g_hccl_oom_call_count{0}; +static std::atomic g_hccl_oom_triggered{false}; // HCCL ReduceOp mapping std::map hcclOp = { @@ -96,11 +98,60 @@ 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"); + if (env_val != nullptr) { + return strtol(env_val, nullptr, 10); + } + env_val = c10_npu::option::get_and_log_env("NPU_OOM_TRIGGER_COUNT"); return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : kDefaultHcclOomTriggerCount; }(); return trigger_count; } +bool isHcclOomTriggerRepeatable() +{ + const static bool repeatable = []() -> bool { + char *env_val = c10_npu::option::get_and_log_env("HCCL_OOM_TRIGGER_MODE"); + if (env_val == nullptr) { + env_val = c10_npu::option::get_and_log_env("NPU_OOM_TRIGGER_MODE"); + } + if (env_val == nullptr) { + return false; + } + std::string mode(env_val); + std::transform(mode.begin(), mode.end(), mode.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + return mode == "always" || mode == "repeat" || mode == "1"; + }(); + return repeatable; +} + +std::string getHcclOomTriggerFile() +{ + const static std::string trigger_file = []() -> std::string { + char *env_val = c10_npu::option::get_and_log_env("HCCL_OOM_TRIGGER_FILE"); + if (env_val != nullptr) { + return std::string(env_val); + } + env_val = c10_npu::option::get_and_log_env("NPU_OOM_TRIGGER_FILE"); + return (env_val != nullptr) ? std::string(env_val) : std::string(); + }(); + return trigger_file; +} + +bool shouldThrowHcclOom(int64_t current_count, int64_t trigger_count) +{ + const auto trigger_file = getHcclOomTriggerFile(); + const bool file_triggered = !trigger_file.empty() && isFileExists(trigger_file); + const bool count_triggered = trigger_count > 0 && current_count >= trigger_count; + if (!file_triggered && !count_triggered) { + return false; + } + if (!isHcclOomTriggerRepeatable() && g_hccl_oom_triggered.exchange(true)) { + return false; + } + return true; +} + void maybeThrowHcclOom(c10d::OpType opType, c10_npu::CaptureStatus capture_status) { if (capture_status != c10_npu::CaptureStatus::None) { @@ -108,12 +159,8 @@ void maybeThrowHcclOom(c10d::OpType opType, c10_npu::CaptureStatus capture_statu } 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) { + if (shouldThrowHcclOom(current_count, trigger_count)) { 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) + diff --git a/torch_npu/csrc/framework/OpCommand.cpp b/torch_npu/csrc/framework/OpCommand.cpp index 21bea98be0..e7d54487a4 100644 --- a/torch_npu/csrc/framework/OpCommand.cpp +++ b/torch_npu/csrc/framework/OpCommand.cpp @@ -1,6 +1,13 @@ #include #include #include +#include +#include +#include +#include +#include +#include +#include #include "torch_npu/csrc/framework/OpCommand.h" #include "torch_npu/csrc/core/npu/register/OptionsManager.h" @@ -22,6 +29,9 @@ namespace { const uint64_t kStringOffset = 16UL; const std::string kStringDType = "string"; +std::mutex g_op_oom_mutex; +std::map g_op_oom_call_counts; +std::set g_op_oom_triggered_devices; static std::unordered_map> floating_limits_map{ {at::ScalarType::Double, {std::numeric_limits::max(), std::numeric_limits::min()}}, {at::ScalarType::Float, {std::numeric_limits::max(), std::numeric_limits::min()}}, @@ -35,6 +45,89 @@ static std::unordered_map> integral_limits_map {at::ScalarType::Byte, {std::numeric_limits::max(), std::numeric_limits::min()}}, {at::ScalarType::Char, {std::numeric_limits::max(), std::numeric_limits::min()}}, {at::ScalarType::Short, {std::numeric_limits::max(), std::numeric_limits::min()}}}; + +int64_t getOpOomTriggerCount() +{ + const static int64_t trigger_count = []() -> int64_t { + char *env_val = c10_npu::option::get_and_log_env("HCCL_OOM_COUNT"); + if (env_val == nullptr) { + env_val = c10_npu::option::get_and_log_env("NPU_OOM_TRIGGER_COUNT"); + } + return (env_val != nullptr) ? strtol(env_val, nullptr, 10) : 0; + }(); + return trigger_count; +} + +bool isOpOomTriggerRepeatable() +{ + const static bool repeatable = []() -> bool { + char *env_val = c10_npu::option::get_and_log_env("NPU_OP_OOM_TRIGGER_MODE"); + if (env_val == nullptr) { + env_val = c10_npu::option::get_and_log_env("NPU_OOM_TRIGGER_MODE"); + } + if (env_val == nullptr) { + return false; + } + std::string mode(env_val); + std::transform(mode.begin(), mode.end(), mode.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + return mode == "always" || mode == "repeat" || mode == "1"; + }(); + return repeatable; +} + +std::string getOpOomTriggerFile() +{ + const static std::string trigger_file = []() -> std::string { + char *env_val = c10_npu::option::get_and_log_env("NPU_OP_OOM_TRIGGER_FILE"); + if (env_val != nullptr) { + return std::string(env_val); + } + env_val = c10_npu::option::get_and_log_env("NPU_OOM_TRIGGER_FILE"); + return (env_val != nullptr) ? std::string(env_val) : std::string(); + }(); + return trigger_file; +} + +bool isTriggerFilePresent(const std::string &path) +{ + if (path.empty()) { + return false; + } + std::ifstream file(path); + return file.good(); +} + +bool shouldThrowOpOom(int device) +{ + const auto trigger_file = getOpOomTriggerFile(); + const int64_t trigger_count = getOpOomTriggerCount(); + if (trigger_file.empty() && trigger_count <= 0) { + return false; + } + + std::lock_guard guard(g_op_oom_mutex); + const bool file_triggered = isTriggerFilePresent(trigger_file); + const int64_t current_count = ++g_op_oom_call_counts[device]; + const bool count_triggered = trigger_count > 0 && current_count >= trigger_count; + if (!file_triggered && !count_triggered) { + return false; + } + if (!isOpOomTriggerRepeatable() && g_op_oom_triggered_devices.count(device) > 0) { + return false; + } + g_op_oom_triggered_devices.insert(device); + return true; +} + +std::string getInjectedOpOomMessage(const std::string &op_name, int device) +{ + std::lock_guard guard(g_op_oom_mutex); + return std::string("NPU out of memory. Failed to allocate memory. Injected NPU op OOM on device ") + + std::to_string(device) + " after " + std::to_string(g_op_oom_call_counts[device]) + + " NPU operations, op type is " + op_name + + ". This is a synthetic OOM fault for recovery validation."; +} } // namespace std::atomic g_used_aclop{false}; @@ -147,6 +240,13 @@ void OpCommand::Run() const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace(); #endif auto stream = c10_npu::getCurrentNPUStream(); + if (shouldThrowOpOom(stream.device_index())) { + auto injected_oom_msg = getInjectedOpOomMessage(op_name, stream.device_index()); + ASCEND_LOGE("%s", injected_oom_msg.c_str()); + aclCmd->releaseSource(); + aclCmds->Pop(); + TORCH_CHECK_WITH(OutOfMemoryError, false, injected_oom_msg.c_str()); + } if (!stream.isSyncLaunchStream() && c10_npu::option::OptionsManager::GetTaskQueueEnable() && !sync) { RECORD_FUNCTION(op_name, std::vector({})); #ifndef BUILD_LIBTORCH @@ -188,6 +288,11 @@ void OpCommand::RunOpApi(const string &op_name, PROC_FUNC func, bool sync) const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace(); #endif auto stream = c10_npu::getCurrentNPUStream(); + if (shouldThrowOpOom(stream.device_index())) { + auto injected_oom_msg = getInjectedOpOomMessage(op_name, stream.device_index()); + ASCEND_LOGE("%s", injected_oom_msg.c_str()); + TORCH_CHECK_WITH(OutOfMemoryError, false, injected_oom_msg.c_str()); + } if (!stream.isSyncLaunchStream() && c10_npu::option::OptionsManager::GetTaskQueueEnable()) { RECORD_FUNCTION(op_name, std::vector({})); #ifndef BUILD_LIBTORCH @@ -235,6 +340,11 @@ void OpCommand::RunOpApiV2(const string &op_name, const PROC_FUNC &func, bool sy const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace(); #endif auto stream = c10_npu::getCurrentNPUStream(); + if (shouldThrowOpOom(stream.device_index())) { + auto injected_oom_msg = getInjectedOpOomMessage(op_name, stream.device_index()); + ASCEND_LOGE("%s", injected_oom_msg.c_str()); + TORCH_CHECK_WITH(OutOfMemoryError, false, injected_oom_msg.c_str()); + } if (!stream.isSyncLaunchStream() && c10_npu::option::OptionsManager::GetTaskQueueEnable()) { RECORD_FUNCTION(op_name, std::vector({})); #ifndef BUILD_LIBTORCH @@ -288,6 +398,12 @@ void OpCommand::RunOpApiV3(const string &op_name, const PROC_FUNC &func, bool sy const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace(); #endif auto stream = c10_npu::getCurrentNPUStream(); + auto inject_device = task_stream == nullptr ? stream.device_index() : task_stream->device_index(); + if (shouldThrowOpOom(inject_device)) { + auto injected_oom_msg = getInjectedOpOomMessage(op_name, inject_device); + ASCEND_LOGE("%s", injected_oom_msg.c_str()); + TORCH_CHECK_WITH(OutOfMemoryError, false, injected_oom_msg.c_str()); + } if (!stream.isSyncLaunchStream() && c10_npu::option::OptionsManager::GetTaskQueueEnable()) { RECORD_FUNCTION(op_name, std::vector({})); #ifndef BUILD_LIBTORCH