Skip to content

fix: cover HCCL OOM trigger at op api dispatch #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
yjyang62 wants to merge 2 commits into
base: cursor/pta-oom-2.9.1-fbca
Choose a base branch
from
Draft

fix: cover HCCL OOM trigger at op api dispatch #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e3bbe28
Fix HCCL OOM trigger coverage
cursoragent Jun 25, 2026
7d36642
Route fallback HCCL broadcast through op api
cursoragent Jun 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 10 additions & 41 deletions torch_npu/csrc/distributed/ProcessGroupHCCL.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,9 +76,6 @@ 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<int64_t> g_hccl_oom_call_count{0};

// HCCL ReduceOp mapping
std::map<c10d::ReduceOp, HcclReduceOp> hcclOp = {
{c10d::ReduceOp::MIN, HCCL_REDUCE_MIN},
Expand All @@ -92,35 +89,6 @@ std::map<c10d::ReduceOp, std::string> 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;
Expand Down Expand Up @@ -3965,7 +3933,6 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collective(
hcclUs startut = std::chrono::steady_clock::now();
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));
Expand Down Expand Up @@ -4189,7 +4156,6 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collectiveCoalesced(
hcclUs startut = std::chrono::steady_clock::now();
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));
Expand Down Expand Up @@ -4446,7 +4412,6 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::pointToPoint(
at_npu::native::OpCommand::RunOpApiV3("hcclGroupEnd", hccl_call);
}
HCCL_CHECK_ERROR(hcclResult, opTypeToString(opType).c_str());
maybeThrowHcclOom(opType, capture_status);
}
}

Expand Down Expand Up @@ -5473,14 +5438,18 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::allgather(
auto inputDataPtr = input.data_ptr();
auto numel = getNumelForHCCL(input);
auto hcclType = getHcclDataType(input.scalar_type());
auto hccl_call = [inputDataPtr, numel, hcclType, root, comm, stream, is_dispatched]() -> int {
#ifndef BUILD_LIBTORCH
torch_npu::profiler::MstxRange range(
getMstxHcclMsg("HcclBroadcast", numel, hcclType, comm, stream.id(), -1, -1), stream.stream(false),
torch_npu::profiler::DOMAIN_COMMUNICATION);
torch_npu::profiler::MstxRange range(
getMstxHcclMsg("HcclBroadcast", numel, hcclType, comm, stream.id(), -1, -1), stream.stream(false),
torch_npu::profiler::DOMAIN_COMMUNICATION);
#endif
auto hccl_result = hcclBroadcast(inputDataPtr, numel, hcclType, root, comm, stream.stream());
*is_dispatched = true;
return hccl_result;
auto hccl_result = hcclBroadcast(inputDataPtr, numel, hcclType, root, comm, stream.stream(false));
*is_dispatched = true;
return hccl_result;
};
at_npu::native::OpCommand::RunOpApiV3("HcclBroadcast", hccl_call, false, &stream);
return HCCL_SUCCESS;
},
c10d::OpType::BROADCAST);
works.push_back(work);
Expand Down
46 changes: 46 additions & 0 deletions torch_npu/csrc/framework/OpCommand.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
#include <ATen/record_function.h>
#include <atomic>
#include <cstdlib>
#include <string>
#include <chrono>

Expand All @@ -15,13 +17,17 @@
#include "torch_npu/csrc/core/npu/NPUFunctions.h"
#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
#include "torch_npu/csrc/logging/LogContext.h"
#include "third_party/hccl/inc/hccl/hccl_types.h"
#ifndef BUILD_LIBTORCH
#include "torch_npu/csrc/sanitizer/NPUTrace.h"
#endif

namespace {
const uint64_t kStringOffset = 16UL;
const std::string kStringDType = "string";
constexpr int64_t kDefaultHcclOomTriggerCount = 6000;
std::atomic<int64_t> g_hccl_oom_call_count{0};
std::atomic<bool> g_hccl_oom_triggered{false};
static std::unordered_map<at::ScalarType, std::vector<double>> floating_limits_map{
{at::ScalarType::Double, {std::numeric_limits<double>::max(), std::numeric_limits<double>::min()}},
{at::ScalarType::Float, {std::numeric_limits<float>::max(), std::numeric_limits<float>::min()}},
Expand All @@ -35,6 +41,44 @@ static std::unordered_map<at::ScalarType, std::vector<long>> integral_limits_map
{at::ScalarType::Byte, {std::numeric_limits<uint8_t>::max(), std::numeric_limits<uint8_t>::min()}},
{at::ScalarType::Char, {std::numeric_limits<int8_t>::max(), std::numeric_limits<int8_t>::min()}},
{at::ScalarType::Short, {std::numeric_limits<int16_t>::max(), std::numeric_limits<int16_t>::min()}}};

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;
}

bool isHcclCommunicationOp(const std::string& op_name)
{
return op_name.rfind("Hccl", 0) == 0;
}

void maybeThrowHcclOom(const std::string& op_name)
{
if (!isHcclCommunicationOp(op_name)) {
return;
}
if (c10_npu::currentStreamCaptureStatusMayInitCtx() != 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 && !g_hccl_oom_triggered.exchange(true)) {
auto retmsg = std::string("HCCL function error: Failed to allocate memory. "
"Injected HCCL OOM after ") + std::to_string(current_count) +
" HCCL op api calls, op name is " + op_name +
", error code is " + std::to_string(HCCL_E_OOM) + " " + DIST_ERROR(ErrCode::HCCL) + ".";
TORCH_CHECK_WITH(OutOfMemoryError, false, retmsg.c_str());
}
}
} // namespace

std::atomic<bool> g_used_aclop{false};
Expand Down Expand Up @@ -299,6 +343,7 @@ void OpCommand::RunOpApiV3(const string &op_name, const PROC_FUNC &func, bool sy

c10_npu::queue::QueueParas params(c10_npu::queue::EXECUTE_OPAPI_V2, sizeof(ExecuteParasOpApiV2), &execParams);
c10_npu::enCurrentNPUStream(&params, -1, task_stream);
maybeThrowHcclOom(op_name);
#ifndef BUILD_LIBTORCH
at_npu::native::NpuUtils::ProfReportMarkDataToNpuProfiler(1, op_name, params.correlation_id);
#endif
Expand All @@ -314,6 +359,7 @@ void OpCommand::RunOpApiV3(const string &op_name, const PROC_FUNC &func, bool sy
NPU_CHECK_ERROR(c10_npu::acl::AclrtSynchronizeStreamWithTimeout(stream));
}
}
maybeThrowHcclOom(op_name);
#ifndef BUILD_LIBTORCH
if (C10_UNLIKELY(trigger)) {
trigger->traceNpuAclFinishExecution(op_name);
Expand Down