diff --git a/torch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp b/torch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp index 560b3437b3..5d600efb44 100644 --- a/torch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp +++ b/torch_npu/csrc/aten/ops/op_api/CopyKernelOpApi.cpp @@ -23,6 +23,7 @@ #include "torch_npu/csrc/core/npu/CachingHostAllocator.h" #include "torch_npu/csrc/aten/NPUOpApiNativeFunctions.h" #include "torch_npu/csrc/aten/NPUNativeFunctions.h" +#include "torch_npu/csrc/core/npu/interface/AclInterface.h" #include "third_party/op-plugin/op_plugin/utils/op_api_common.h" #ifndef BUILD_LIBTORCH #include "torch_npu/csrc/sanitizer/NPUTrace.h" @@ -49,7 +50,7 @@ void copy_between_host_and_device_opapi(at::Tensor& dst, const at::Tensor& src, void* currentPtr = torch_npu::utils::is_npu(dst) ? src.data_ptr() : dst.data_ptr(); process_non_blocking_copy(storage, currentPtr, stream, kind); } else { - aclError error = aclrtSynchronizeStream(stream); + aclError error = c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream); auto ret = CalcuOpUtil::AclrtMemcpyWithModeSwitch( std::make_pair(dst.storage().unsafeGetStorageImpl(), dst.storage_offset() * dst.itemsize()), nbytes, std::make_pair(src.storage().unsafeGetStorageImpl(), src.storage_offset() * src.itemsize()), nbytes, kind); diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 783d27969a..5abd2b9354 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -107,7 +107,6 @@ 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; static char SHAREABLE_HANDLE_VERSION = 1; enum ShareableHandleType : char { SHAREABLE_NPU_MALLOC = 'c', @@ -1152,9 +1151,8 @@ 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."); + { + maybeThrowPtaOom("malloc", device); 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 @@ -2937,6 +2935,9 @@ class DeviceCachingAllocator { } else { TORCH_NPU_MEMORY_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event.get()); } + if (check_error) { + maybeThrowPtaOom("NPUCachingAllocator::synchronize_and_free_events"); + } #ifndef BUILD_LIBTORCH const c10_npu::impl::PyCallbackTrigger *trigger = c10_npu::impl::NPUTrace::getTrace(); if (C10_UNLIKELY(trigger)) { diff --git a/torch_npu/csrc/core/npu/NPUEvent.cpp b/torch_npu/csrc/core/npu/NPUEvent.cpp index 3a20235889..8082fece88 100644 --- a/torch_npu/csrc/core/npu/NPUEvent.cpp +++ b/torch_npu/csrc/core/npu/NPUEvent.cpp @@ -2,6 +2,7 @@ #include "torch_npu/csrc/core/npu/NPUFunctions.h" #include "torch_npu/csrc/core/npu/NPUGuard.h" #include "torch_npu/csrc/core/npu/NPUException.h" +#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h" #include "torch_npu/csrc/core/npu/NPUEventManager.h" #include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h" #include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h" @@ -202,6 +203,7 @@ void NPUEvent::synchronize() const } NPU_CHECK_ERROR(aclrtSynchronizeEvent(event_)); ASCEND_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", event_); + maybeThrowPtaOom("NPUEvent::synchronize"); #ifndef BUILD_LIBTORCH const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace(); if (C10_UNLIKELY(trigger)) { diff --git a/torch_npu/csrc/core/npu/NPUException.cpp b/torch_npu/csrc/core/npu/NPUException.cpp index e9f8dd356c..4a8f41a0ab 100644 --- a/torch_npu/csrc/core/npu/NPUException.cpp +++ b/torch_npu/csrc/core/npu/NPUException.cpp @@ -1,6 +1,12 @@ +#include +#include +#include +#include + #include "torch_npu/csrc/core/npu/NPUException.h" +#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h" #include "torch_npu/csrc/core/npu/NPUFunctions.h" -#include "torch_npu/csrc/core/npu/NPUStream.h" +#include "torch_npu/csrc/core/npu/NPURecovery.h" #include "torch_npu/csrc/core/npu/NpuVariables.h" #include "torch_npu/csrc/core/npu/register/OptionsManager.h" @@ -433,4 +439,133 @@ bool isCannOOM(const std::string &errMsg) return false; } +namespace { +static std::atomic g_pta_oom_injected{false}; +static std::atomic g_pta_oom_timer_started{false}; +static std::atomic g_pta_oom_timer_expired{false}; +static std::chrono::steady_clock::time_point g_pta_oom_start_time; +static constexpr int64_t kDefaultTriggerAfterSeconds = 360; + +bool isPtaOomDebugEnabled() +{ + const static bool enabled = []() -> bool { + char *env_val = c10_npu::option::get_and_log_env("PTA_OOM_DEBUG"); + return (env_val != nullptr) && (strtol(env_val, nullptr, 10) != 0); + }(); + return enabled; +} + +int64_t getTriggerAfterSeconds() +{ + const static int64_t trigger_seconds = []() -> int64_t { + char *seconds_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_AFTER_SECONDS"); + if (seconds_val != nullptr) { + return strtol(seconds_val, nullptr, 10); + } + char *minutes_val = c10_npu::option::get_and_log_env("PTA_OOM_TRIGGER_AFTER_MINUTES"); + if (minutes_val != nullptr) { + return strtol(minutes_val, nullptr, 10) * 60; + } + char *timer_mode = c10_npu::option::get_and_log_env("PTA_OOM_TIMER"); + if (timer_mode != nullptr && strtol(timer_mode, nullptr, 10) != 0) { + return kDefaultTriggerAfterSeconds; + } + return 0; + }(); + return trigger_seconds; +} + +void throwFullCardPtaOom(const char *context, int device); + +void ensurePtaOomTimerStarted() +{ + if (g_pta_oom_timer_started.load()) { + return; + } + const int64_t trigger_seconds = getTriggerAfterSeconds(); + if (trigger_seconds <= 0) { + return; + } + bool expected = false; + if (!g_pta_oom_timer_started.compare_exchange_strong(expected, true)) { + return; + } + g_pta_oom_start_time = std::chrono::steady_clock::now(); + ASCEND_LOGI("PTA OOM timer started, will trigger after %lld seconds", + static_cast(trigger_seconds)); + std::thread([trigger_seconds]() { + std::this_thread::sleep_for(std::chrono::seconds(trigger_seconds)); + if (!g_pta_oom_injected.load()) { + g_pta_oom_timer_expired.store(true); + ASCEND_LOGI("PTA OOM timer expired after %lld seconds, pending throw on next hook", + static_cast(trigger_seconds)); + } + }).detach(); +} + +int64_t getElapsedSecondsSincePtaOomStart() +{ + const auto now = std::chrono::steady_clock::now(); + return std::chrono::duration_cast(now - g_pta_oom_start_time).count(); +} + +const bool kPtaOomTimerArmed = []() { + if (getTriggerAfterSeconds() > 0) { + ensurePtaOomTimerStarted(); + } + return true; +}(); + +void throwFullCardPtaOom(const char *context, int device) +{ + if (g_pta_oom_injected.exchange(true)) { + return; + } + + if (device < 0) { + NPU_CHECK_ERROR(c10_npu::GetDevice(&device)); + } + + NPUCachingAllocator::markAllBlockUnsafe(device); + c10_npu::set_npu_data_unsafe_flag(true); + + const int64_t trigger_seconds = getTriggerAfterSeconds(); + auto retmsg = std::string("NPU out of memory. Injected full-card PTA OOM on NPU ") + + std::to_string(device) + + ". All existing tensors on this device are marked unsafe. " + "Triggered after " + std::to_string(trigger_seconds) + " seconds"; + if (context != nullptr && context[0] != '\0') { + retmsg += ", context is "; + retmsg += context; + } + retmsg += ". "; + retmsg += PTA_ERROR(ErrCode::MEMORY); + TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str()); +} +} // namespace + +void maybeThrowPtaOom(const char *context, int device) +{ + const int64_t trigger_seconds = getTriggerAfterSeconds(); + if (trigger_seconds <= 0 || g_pta_oom_injected.load()) { + return; + } + ensurePtaOomTimerStarted(); + const bool expired = g_pta_oom_timer_expired.load() || + (g_pta_oom_timer_started.load() && getElapsedSecondsSincePtaOomStart() >= trigger_seconds); + if (!expired) { + if (isPtaOomDebugEnabled()) { + const int64_t elapsed = getElapsedSecondsSincePtaOomStart(); + if (elapsed > 0 && elapsed % 60 == 0) { + ASCEND_LOGI("PTA OOM timer: elapsed=%lld/%lld seconds context=%s", + static_cast(elapsed), + static_cast(trigger_seconds), + context != nullptr ? context : ""); + } + } + return; + } + throwFullCardPtaOom(context, device); +} + } // namespace c10_npu diff --git a/torch_npu/csrc/core/npu/NPUException.h b/torch_npu/csrc/core/npu/NPUException.h index 987d9fa6a1..c44a26c793 100644 --- a/torch_npu/csrc/core/npu/NPUException.h +++ b/torch_npu/csrc/core/npu/NPUException.h @@ -336,6 +336,8 @@ std::string handleSuspectRemoteError(int errorCode); bool isCannOOM(const std::string &errMsg); +void maybeThrowPtaOom(const char *context = nullptr, int device = -1); + bool ShouldAppendDeviceErrorVerbose(); void clear_device_error_info(); diff --git a/torch_npu/csrc/core/npu/NPUFunctions.cpp b/torch_npu/csrc/core/npu/NPUFunctions.cpp index 02cbc299c0..ac38d5c9be 100644 --- a/torch_npu/csrc/core/npu/NPUFunctions.cpp +++ b/torch_npu/csrc/core/npu/NPUFunctions.cpp @@ -6,6 +6,7 @@ #include "torch_npu/csrc/core/npu/NPUAffinityController.h" #include "torch_npu/csrc/core/npu/register/OptionsManager.h" #include "torch_npu/csrc/core/npu/GetCANNInfo.h" +#include "torch_npu/csrc/core/npu/interface/AclInterface.h" #include "third_party/acl/inc/acl/acl_rt.h" #ifndef BUILD_LIBTORCH #include "torch_npu/csrc/sanitizer/NPUTrace.h" @@ -363,7 +364,7 @@ void stream_synchronize(aclrtStream stream) trigger->traceNpuStreamSynchronization(reinterpret_cast(stream)); } #endif - NPU_CHECK_ERROR(aclrtSynchronizeStream(stream)); + NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream)); } aclError SetDeviceResLimit(int32_t device, int32_t type, uint32_t value) diff --git a/torch_npu/csrc/core/npu/NPUQueue.cpp b/torch_npu/csrc/core/npu/NPUQueue.cpp index 4d820bbf9b..e775386c25 100644 --- a/torch_npu/csrc/core/npu/NPUQueue.cpp +++ b/torch_npu/csrc/core/npu/NPUQueue.cpp @@ -8,6 +8,7 @@ #include "torch_npu/csrc/framework/OpCommand.h" #include "torch_npu/csrc/core/npu/register/OptionsManager.h" #include "torch_npu/csrc/core/npu/NPUEventManager.h" +#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h" #include "torch_npu/csrc/logging/LogContext.h" #ifndef BUILD_LIBTORCH @@ -348,6 +349,9 @@ NPUStatus Repository::MakeSureQueueEmpty(bool check_error) throw std::runtime_error(runtime_error); } } + if (check_error) { + maybeThrowPtaOom("MakeSureQueueEmpty", device_idx); + } logger->debug("MakeSureQueueEmpty: clearing successful, device = %d, write_idx = %u, read_idx = %u, status = %d", device_idx, write_idx.idx, read_idx.idx, GetStatus()); diff --git a/torch_npu/csrc/core/npu/NPUStream.cpp b/torch_npu/csrc/core/npu/NPUStream.cpp index ad554f5694..1ff6f1bd5f 100644 --- a/torch_npu/csrc/core/npu/NPUStream.cpp +++ b/torch_npu/csrc/core/npu/NPUStream.cpp @@ -14,6 +14,7 @@ #include "torch_npu/csrc/core/npu/NPUGuard.h" #include "torch_npu/csrc/core/npu/NPUQueue.h" #include "torch_npu/csrc/core/npu/NPUException.h" +#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h" #include "torch_npu/csrc/core/npu/register/OptionsManager.h" #include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h" #include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h" @@ -508,6 +509,10 @@ NPUStatus emptyAllNPUStream(bool check_error) } } + if (check_error) { + maybeThrowPtaOom("emptyAllNPUStream"); + } + return NPU_STATUS_SUCCESS; } @@ -559,6 +564,9 @@ bool npuSynchronizeDevice(bool check_error) ASCEND_LOGE("MakeSureQueueEmpty fail, ret: %s", ret.c_str()); } } + if (check_error) { + maybeThrowPtaOom("npuSynchronizeDevice"); + } auto acl_ret = c10_npu::acl::AclrtSynchronizeDeviceWithTimeout(); if (acl_ret != ACL_ERROR_NONE) { CHECK_AND_THROW_ERROR_WITH_SPECIFIC_MESSAGE(acl_ret); diff --git a/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp b/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp index 52411a045d..77cad1980d 100644 --- a/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp +++ b/torch_npu/csrc/core/npu/impl/NPUGuardImpl.cpp @@ -12,6 +12,8 @@ #include "torch_npu/csrc/core/NPUSerialization.h" #include "torch_npu/csrc/core/npu/NPUHooksInterface.h" #include "torch_npu/csrc/core/npu/NPUEventManager.h" +#include "torch_npu/csrc/core/npu/NPUException.h" +#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h" #ifndef BUILD_LIBTORCH #include "torch_npu/csrc/sanitizer/NPUTrace.h" @@ -221,6 +223,7 @@ void NPUGuardImpl::synchronizeEvent(void* event) const NPU_CHECK_ERROR_WITHOUT_UCE(aclrtSynchronizeEvent(npu_event)); ASCEND_LOGI("Event: aclrtSynchronizeEvent is successfully executed, event=%p", npu_event); + maybeThrowPtaOom("NPUGuardImpl::synchronizeEvent"); #ifndef BUILD_LIBTORCH const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace(); if (C10_UNLIKELY(trigger)) { diff --git a/torch_npu/csrc/core/npu/interface/AclInterface.cpp b/torch_npu/csrc/core/npu/interface/AclInterface.cpp index 1c9f1793ef..75c5c611d4 100644 --- a/torch_npu/csrc/core/npu/interface/AclInterface.cpp +++ b/torch_npu/csrc/core/npu/interface/AclInterface.cpp @@ -8,6 +8,7 @@ #include "torch_npu/csrc/core/npu/register/OptionsManager.h" #include "torch_npu/csrc/core/npu/NPUException.h" #include "torch_npu/csrc/core/npu/NPUFunctions.h" +#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h" #include "torch_npu/csrc/core/npu/GetCANNInfo.h" #ifndef BUILD_LIBTORCH #include "torch_npu/csrc/sanitizer/NPUTrace.h" @@ -594,8 +595,9 @@ aclError AclrtSynchronizeStreamWithTimeout(aclrtStream stream) { typedef aclError (*AclrtSynchronizeStreamWithTimeout)(aclrtStream, int32_t); static AclrtSynchronizeStreamWithTimeout func = (AclrtSynchronizeStreamWithTimeout)GET_FUNC(aclrtSynchronizeStreamWithTimeout); int32_t timeout = c10_npu::option::OptionsManager::GetACLExecTimeout(); + aclError ret = ACL_ERROR_NONE; if (func != nullptr) { - return func(stream, timeout); + ret = func(stream, timeout); } else { TORCH_NPU_WARN_ONCE(func, "Failed to find function", "aclrtSynchronizeStreamWithTimeout"); typedef aclError (*AclrtSynchronizeStream)(aclrtStream); @@ -604,8 +606,10 @@ aclError AclrtSynchronizeStreamWithTimeout(aclrtStream stream) { func_backup = (AclrtSynchronizeStream)GET_FUNC(aclrtSynchronizeStream); } TORCH_CHECK(func_backup, "Failed to find function", "aclrtSynchronizeStreamWithTimeout and aclrtSynchronizeStream", PROF_ERROR(ErrCode::NOT_FOUND)); - return func_backup(stream); + ret = func_backup(stream); } + c10_npu::maybeThrowPtaOom("AclrtSynchronizeStreamWithTimeout"); + return ret; } aclError AclrtDestroyStreamForce(aclrtStream stream) { @@ -992,8 +996,9 @@ aclError AclrtSynchronizeDeviceWithTimeout(void) typedef aclError (*AclrtSynchronizeDeviceWithTimeout)(int32_t); static AclrtSynchronizeDeviceWithTimeout func = (AclrtSynchronizeDeviceWithTimeout)GET_FUNC(aclrtSynchronizeDeviceWithTimeout); int32_t timeout = c10_npu::option::OptionsManager::GetACLDeviceSyncTimeout(); + aclError ret = ACL_ERROR_NONE; if (func != nullptr) { - return func(timeout); + ret = func(timeout); } else { if (timeout > 0) { TORCH_NPU_WARN_ONCE("The ACL_DEVICE_SYNC_TIMEOUT does not take effect. If you want to enable this env, please upgrade CANN to the matching version."); @@ -1004,8 +1009,10 @@ aclError AclrtSynchronizeDeviceWithTimeout(void) func_backup = (AclrtSynchronizeDevice)GET_FUNC(aclrtSynchronizeDevice); } TORCH_CHECK(func_backup, "Failed to find function ", "aclrtSynchronizeDeviceWithTimeout and aclrtSynchronizeDevice", PTA_ERROR(ErrCode::NOT_FOUND)); - return func_backup(); + ret = func_backup(); } + c10_npu::maybeThrowPtaOom("AclrtSynchronizeDeviceWithTimeout"); + return ret; } aclError AclrtEventGetTimestamp(aclrtEvent event, uint64_t *timestamp) @@ -1076,6 +1083,7 @@ aclError AclmdlRIDebugPrint(aclmdlRI modelRI) aclError AclmdlRIExecuteAsync(aclmdlRI modelRI, aclrtStream stream) { ACL_CALL_LOG("aclmdlRIExecuteAsync", "modelRI=" << modelRI << ", stream=" << stream); + c10_npu::maybeThrowPtaOom("AclmdlRIExecuteAsync"); typedef aclError (*AclmdlRIExecuteAsync)(aclmdlRI, aclrtStream); static AclmdlRIExecuteAsync func = nullptr; if (func == nullptr) { @@ -1574,6 +1582,9 @@ aclError AclrtMemcpyAsyncWithCondition(void *dst, size_t destMax, const void *sr func = (AclrtMemcpyAsyncWithConditionFunc)GET_FUNC(aclrtMemcpyAsyncWithCondition); } TORCH_CHECK(func, "Failed to find function ", "aclrtMemcpyAsyncWithCondition", PROF_ERROR(ErrCode::NOT_FOUND)); + if (kind == aclrtMemcpyKind::ACL_MEMCPY_DEVICE_TO_HOST) { + c10_npu::maybeThrowPtaOom("AclrtMemcpyAsyncWithCondition"); + } return func(dst, destMax, src, count, kind, stream); } diff --git a/torch_npu/csrc/framework/OpParamMaker.cpp b/torch_npu/csrc/framework/OpParamMaker.cpp index 6fd381f8e8..963e7bcdb6 100644 --- a/torch_npu/csrc/framework/OpParamMaker.cpp +++ b/torch_npu/csrc/framework/OpParamMaker.cpp @@ -3,6 +3,8 @@ #include "torch_npu/csrc/core/npu/CachingHostAllocator.h" #include "torch_npu/csrc/core/npu/NPUEventManager.h" +#include "torch_npu/csrc/core/npu/NPUException.h" +#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h" #include "torch_npu/csrc/core/npu/NPUQueue.h" #include "torch_npu/csrc/core/npu/interface/AsyncTaskQueueInterface.h" #include "torch_npu/csrc/distributed/HCCLUtils.hpp" @@ -273,6 +275,8 @@ aclError OpCommandImpl::InnerRun( aclError OpCommandImpl::InnerRunOpApi(const string &op_name, PROC_FUNC func) { + c10_npu::maybeThrowPtaOom("InnerRunOpApi"); + aclError ret; auto stream = c10_npu::getCurrentNPUStream(); if (stream.getRepoStopFlag()) { @@ -332,6 +336,7 @@ bool ContainsAny(const std::string& str, std::initializer_list patt int ExecFunc(c10_npu::queue::QueueParas *in, aclrtStream stream) { auto cur_paras = static_cast(in->paramVal); + c10_npu::maybeThrowPtaOom("ExecFunc"); ASCEND_LOGD("Op %s Run.", cur_paras->opType); logger->debug("ExecFunc: Op %s Run.", cur_paras->opType); aclError ret; @@ -426,6 +431,7 @@ int ExecFunc(c10_npu::queue::QueueParas *in, aclrtStream stream) int ExecFuncOpApi(c10_npu::queue::QueueParas *in, aclrtStream stream) { auto cur_paras = static_cast(in->paramVal); + c10_npu::maybeThrowPtaOom("ExecFuncOpApi"); ASCEND_LOGD("Op %s Run.", cur_paras->opType); logger->debug("ExecFuncOpApi: Op %s Run.", cur_paras->opType); aclError ret; @@ -468,6 +474,7 @@ int ExecFuncOpApi(c10_npu::queue::QueueParas *in, aclrtStream stream) int MemcopyAsyncFunc(c10_npu::queue::QueueParas *in, aclrtStream stream) { auto cur_paras = static_cast(in->paramVal); + c10_npu::maybeThrowPtaOom("MemcopyAsyncFunc"); logger->debug("MemcopyAsyncFunc Run."); aclError ret; bool flag;