Skip to content

Import pta_oom_2.9.1 HCCL OOM fault injection from gitcode #4

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 3 commits into
base: cursor/hccl-oom-fault-2bd5
Choose a base branch
from
Draft

Import pta_oom_2.9.1 HCCL OOM fault injection from gitcode #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d60e0af
hccl oom
yjyang62 Jun 12, 2026
f7f4411
hccl_6000
yjyang62 Jun 14, 2026
d8450a4
hccloom
yjyang62 Jun 14, 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
4 changes: 1 addition & 3 deletions torch_npu/csrc/core/npu/NPUCachingAllocator.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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...
Expand Down
46 changes: 40 additions & 6 deletions torch_npu/csrc/distributed/ProcessGroupHCCL.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#include <ATen/record_function.h>
#include <algorithm>
#include <atomic>
#include <map>
#include <tuple>
#include <unordered_set>
Expand Down Expand Up @@ -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<int64_t> g_hccl_oom_call_count{0};

// HCCL ReduceOp mapping
std::map<c10d::ReduceOp, HcclReduceOp> hcclOp = {
Expand All @@ -89,7 +92,35 @@ 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 @@ -3785,7 +3816,6 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collective(
auto key = getKeyFromDevices(devices);
HcclCommConfig config = createHcclCommConfigWithOptions();
std::vector<std::shared_ptr<HCCLComm>> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config);

auto& hcclStreams = hcclStreams_[key];
syncStreams(devices, hcclEvents_[key], hcclStreams);

Expand Down Expand Up @@ -3933,7 +3963,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));
Expand Down Expand Up @@ -4000,7 +4032,6 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::collectiveCoalesced(
NPU_CHECK_ERROR(c10_npu::SetDevice(devices[0].index()));
HcclCommConfig config = createHcclCommConfigWithOptions();
std::vector<std::shared_ptr<HCCLComm>> hcclComms = getHCCLComm(key, devices, HcclCommType::DEFAULT, &config);

auto& hcclStreams = hcclStreams_[key];
syncStreams(devices, hcclEvents_[key], hcclStreams);

Expand Down Expand Up @@ -4156,7 +4187,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));
Expand Down Expand Up @@ -4236,7 +4269,6 @@ c10::intrusive_ptr<c10d::Work> 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_++;
Expand Down Expand Up @@ -4406,13 +4438,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);
}
}

Expand Down