diff --git a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp index 257e0764bf..783d27969a 100644 --- a/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp +++ b/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp @@ -1155,9 +1155,7 @@ class DeviceCachingAllocator { { g_malloc_call_count++; auto retmsg = std::string("NPU out of memory. Tried to allocate more than 1EB memory."); - if (g_malloc_call_count > 60003 && g_malloc_call_count < 60005) { - 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... diff --git a/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp b/torch_npu/csrc/distributed/ProcessGroupHCCL.cpp index 005f98fbdf..d6de8a8337 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 = 6000; +static std::atomic g_hccl_oom_call_count{0}; // HCCL ReduceOp mapping std::map hcclOp = { @@ -89,7 +92,35 @@ std::map unsupportedOp = { {c10d::ReduceOp::BOR, "BOR"}, {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; @@ -3785,7 +3816,6 @@ c10::intrusive_ptr ProcessGroupHCCL::collective( auto key = getKeyFromDevices(devices); HcclCommConfig config = createHcclCommConfigWithOptions(); std::vector> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config); - auto& hcclStreams = hcclStreams_[key]; syncStreams(devices, hcclEvents_[key], hcclStreams); @@ -3933,7 +3963,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)); @@ -4000,7 +4032,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); - auto& hcclStreams = hcclStreams_[key]; syncStreams(devices, hcclEvents_[key], hcclStreams); @@ -4156,7 +4187,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)); @@ -4236,7 +4269,6 @@ c10::intrusive_ptr ProcessGroupHCCL::pointToPoint( key = getKeyFromDevices(devices); hcclComms = getHCCLComm(key, devices); } - // Bump the logical operation counter regardless of whether this op is // coalesced or individual op_id_++; @@ -4406,13 +4438,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); } }