Unknown memory type error for alltoallv with torch.empty tensors

[rickybalin]

Summary

In Pytorch, we can replicate an alltoallv with some non-zero and some zero-sized buffers (similar to the use case described in #174) by passing torch.empty(0) tensors to the list of send and receive buffers. This type of collective breaks on Intel GPU on Aurora with the oneCCL backend with CCL_ALLTOALLV=topo and CCL_ALLTOALLV_MONOLITHIC_KERNEL=0. It works however with CCL_ALLTOALLV_MONOLITHIC_KERNEL=1, although it is VERY slow, and it also works with other alltoallv algorithms (like direct and naive). The error with the topology-aware algorithm is shown below.

|CCL_ERROR| worker.cpp:338 ccl_worker_func: worker 0 caught internal exception: oneCCL: ze_handle_manager.cpp:226 get_ptr: EXCEPTION: unknown memory type
terminate called after throwing an instance of 'ccl::v1::exception'
  what():  oneCCL: ze_handle_manager.cpp:226 get_ptr: EXCEPTION: unknown memory type

Version and environment

Sunspot system at ALCF.

• oneCCL release: 2021.17
• MPI version: aurora_test branch @ 3c70a61
• Compiler type and version: Intel compiler 2024.3.2
• PyTorch: 2.10.0
• OS name and version:
• GPU driver information:
• Hardware configuration: Aurora/Sunspot @ ALCF configuration

Reproducer

The reproducer is found here (includes a correctness check too):
https://github.com/argonne-lcf/nekRS-ML/blob/alcf4/3rd_party/dist-gnn/all2all_bench.py

And run instructions for Aurora are found here:
https://github.com/argonne-lcf/nekRS-ML/blob/alcf4/3rd_party/dist-gnn/run_all2all_bench.sh

Logs

Will produce this.

Expected behavior

The expected behavior is for the alltoallv collective to work with this setup with the topo algorithm and CCL_ALLTOALLV_MONOLITHIC_KERNEL=0.

Observed behavior

The observed behavior is the following error

|CCL_ERROR| worker.cpp:338 ccl_worker_func: worker 0 caught internal exception: oneCCL: ze_handle_manager.cpp:226 get_ptr: EXCEPTION: unknown memory type
terminate called after throwing an instance of 'ccl::v1::exception'
  what():  oneCCL: ze_handle_manager.cpp:226 get_ptr: EXCEPTION: unknown memory type

Existing workarounds

Using other alltoallv algorithms and setting CCL_ALLTOALLV_MONOLITHIC_KERNEL=1. The latter is especially undesirable due to the performance penalty.

Affected projects

This bug affects scaling of a graph neural network on Aurora, which a critical workload for modeling mesh-based PDE based systems.

[rickybalin]

Working off of oneCCL 2021.17.2 (6649993).

Using some AI agents, I was able to get rid of the bug and then address a correctness issue when performing the collective on multiple nodes. The first error (unknown memory type) appears to come from not properly tracking and handling zero-size buffers through the IPC handle exchange and copy operations when torch.empty(0, ...) tensors are used. After addressing this issue, a correctness check identified that on >1 nodes, ranks on nodes other than 0 would get the wrong data in their receive buffers. This second issue seemed to have been caused by a use of the local rank instead of the global rank number in the alltoallv topo algorithm.

Please note that these are not meant to be suggested fixes, but rather they hopefully can provide helpful information on what the causes of the issues were.

diff --git a/src/coll/algorithms/alltoallv.cpp b/src/coll/algorithms/alltoallv.cpp
index 3e83329..915fa2d 100644
--- a/src/coll/algorithms/alltoallv.cpp
+++ b/src/coll/algorithms/alltoallv.cpp
@@ -644,12 +644,23 @@ ccl::status ccl_coll_build_topo_alltoallv(ccl_sched* main_sched,
                     auto peer_rank = (card_idx * tile_count + tile_idx);
                     if (peer_rank == comm->rank())
                         continue;
-                    if (counts[peer_rank] == 0)
+                    // Store local peer rank for IPC handle lookup
+                    auto local_peer_rank = peer_rank;
+
+                    // Adjust peer_rank to global rank for multi-node BEFORE checking counts
+                    // because counts/bufs arrays are indexed by global rank
+                    if (!is_single_node) {
+                        peer_rank = peer_rank + (r2r_comm->rank() * node_comm->size());
+                    }
+
+                    // Now check counts with the correct (global) peer_rank
+		    if (counts[peer_rank] == 0)
                         continue;

                     copy_attr attr{};
-                    attr.peer_rank = peer_rank;
-                    attr.peer_buf_idx = start_buf_idx + offset;
+                    //attr.peer_rank = peer_rank;
+                    attr.peer_rank = local_peer_rank;  // Use local rank for IPC handle lookup
+		    attr.peer_buf_idx = start_buf_idx + offset;
                     attr.map_comm = comm;
                     auto copy_engine_idx = card_idx * 2;
                     if (ccl::global_data::env().type2_mode == ccl::type2_tune_mode::detected ||
@@ -659,10 +670,10 @@ ccl::status ccl_coll_build_topo_alltoallv(ccl_sched* main_sched,
                     attr.hint_queue_index = copy_engine_idx;
                     attr.direction = copy_direction::c2c;

-                    if (!is_single_node) {
-                        // in order to get the correct offset for peer rank for example(2)
-                        peer_rank = peer_rank + (r2r_comm->rank() * node_comm->size());
-                    }
+                    //if (!is_single_node) {
+                    //    // in order to get the correct offset for peer rank for example(2)
+                    //    peer_rank = peer_rank + (r2r_comm->rank() * node_comm->size());
+                    //}

                     auto src = bufs[peer_rank];
                     auto dst = ccl_buffer();
diff --git a/src/sched/entry/ze/ze_copy_entry.cpp b/src/sched/entry/ze/ze_copy_entry.cpp
index e322def..6f5aac3 100644
--- a/src/sched/entry/ze/ze_copy_entry.cpp
+++ b/src/sched/entry/ze/ze_copy_entry.cpp
@@ -53,6 +53,15 @@ void ze_copy_entry::init_ze_hook() {
         }
     }

+    // Skip copy operation if count is zero or if buffers are null (empty/zero-size buffers)
+    if (count == 0 || !in_buf.get_ptr() || !out_buf.get_ptr()) {
+        // For zero-count or null buffer cases, just signal the event without doing any copy
+        ze_command_list_handle_t list =
+            ze_base_entry::get_copy_list(attr.direction, attr.hint_queue_index, attr.force_queue_type);
+        ZE_APPEND_CALL(ze_cmd_barrier, list, ze_base_entry::entry_event, wait_events);
+        return;
+    }
+
     void* dst = static_cast<char*>(out_buf.get_ptr()) + attr.out_buf_offset * dtype.size();
     void* src = static_cast<char*>(in_buf.get_ptr()) + attr.in_buf_offset * dtype.size();

diff --git a/src/sched/entry/ze/ze_handle_exchange_entry.cpp b/src/sched/entry/ze/ze_handle_exchange_entry.cpp
index 2509add..474ee79 100644
--- a/src/sched/entry/ze/ze_handle_exchange_entry.cpp
+++ b/src/sched/entry/ze/ze_handle_exchange_entry.cpp
@@ -116,11 +116,17 @@ void ze_handle_exchange_entry::create_local_ipc_handles() {
     for (size_t buf_idx = 0; buf_idx < in_buffers.size(); buf_idx++) {
         int mem_handle = ccl::utils::invalid_mem_handle;
         auto mem_ptr = in_buffers[buf_idx].first;
-        if (mem_ptr == nullptr) {
-            continue;
+        auto mem_type = in_buffers[buf_idx].second;  // Moved BEFORE null check
+	if (mem_ptr == nullptr) {
+	    // For null buffers (zero-size), still set mem_type to avoid
+            // "unknown memory type" errors when the handle is accessed later
+            int handle_idx = get_handle_idx(sched->coll_param.ctype, rank);
+            handles[handle_idx][buf_idx].mem_type = mem_type;
+            handles[handle_idx][buf_idx].is_null_buffer = true;
+	    continue;
         }

-        auto mem_type = in_buffers[buf_idx].second;
+        //auto mem_type = in_buffers[buf_idx].second;
         mem_info_t mem_info{};

         ze_ipc_mem_handle_t ipc_handle{};
@@ -225,6 +231,7 @@ void ze_handle_exchange_entry::fill_payload(payload_t& payload, size_t buf_idx)
     payload.handle_id = this_handle.handle_id;
     payload.remote_pid = getpid();
     payload.remote_ptr = this_handle.remote_ptr;
+    payload.is_null_buffer = this_handle.is_null_buffer;

     auto mem_ptr = in_buffers[buf_idx].first;
     if (mem_ptr == nullptr) {
@@ -272,6 +279,7 @@ void ze_handle_exchange_entry::fill_remote_handle(const payload_t& payload,
     this_handle.remote_device_id = payload.remote_device_id;
     this_handle.device_fd = payload.device_fd;
     this_handle.remote_mem_alloc_id = payload.remote_mem_alloc_id;
+    this_handle.is_null_buffer = payload.is_null_buffer;

     LOG_DEBUG("get IPC handle: { peer: ",
               idx,
diff --git a/src/sched/entry/ze/ze_handle_exchange_entry.hpp b/src/sched/entry/ze/ze_handle_exchange_entry.hpp
index f66ea5a..6776d7f 100644
--- a/src/sched/entry/ze/ze_handle_exchange_entry.hpp
+++ b/src/sched/entry/ze/ze_handle_exchange_entry.hpp
@@ -102,6 +102,7 @@ private:
         ssize_t remote_context_id{ ccl::utils::invalid_context_id };
         ssize_t remote_device_id{ ccl::utils::invalid_device_id };
         int device_fd{ ccl::utils::invalid_fd };
+        bool is_null_buffer{ false };
     };

     void fill_payload(payload_t& payload, size_t buf_idx);
diff --git a/src/sched/ze/ze_handle_manager.cpp b/src/sched/ze/ze_handle_manager.cpp
index a0346a1..d2c48f1 100644
--- a/src/sched/ze/ze_handle_manager.cpp
+++ b/src/sched/ze/ze_handle_manager.cpp
@@ -228,6 +228,14 @@ void* ipc_handle_manager::get_ptr(int rank,

     LOG_DEBUG("context: ", context, ", device: ", device, ", rank: ", rank, ", buf_idx: ", buf_idx);

+    // Check if this is a null/empty buffer (zero-size buffer from the remote rank).
+    // For null buffers, is_null_buffer flag is set during handle creation/exchange.
+    // In this case, we should return nullptr without trying to open an IPC handle.
+    if (handle_info.is_null_buffer) {
+        LOG_DEBUG("null/empty buffer detected for rank: ", rank, ", buf_idx: ", buf_idx);
+        return nullptr;
+    }
+
     if (mem_ptr == nullptr) {
         if (mem_type == ipc_mem_type::memory) {
             open_handle(handle_info, &mem_ptr, to_cache);
diff --git a/src/sched/ze/ze_handle_manager.hpp b/src/sched/ze/ze_handle_manager.hpp
index 0355b68..04c6793 100644
--- a/src/sched/ze/ze_handle_manager.hpp
+++ b/src/sched/ze/ze_handle_manager.hpp
@@ -57,6 +57,7 @@ struct ipc_handle_desc {
     int device_fd{ ccl::utils::invalid_fd };

     bool is_cached = false;
+    bool is_null_buffer = false;

     ipc_handle_desc();
     ipc_handle_desc(const ipc_handle_desc&) = default;

[srirajpaul]

@rickybalin Thanks for the triage. it pinpoints the issue correctly. For now, we are not so much trying to change ipc exchange phase so much. So I am passing a dummy pointer instead of nullptr to avoid the issues during ipc exchange. Also the issue with incorrect results, i used the same fix, but used the logic for global_peer instead of local_peer, since it seems more appropriate to me, essentially keeping peer_rank based on the comm that is available.

diff --git a/src/coll/algorithms/alltoallv.cpp b/src/coll/algorithms/alltoallv.cpp
index 295b56357..0fd404729 100644
--- a/src/coll/algorithms/alltoallv.cpp
+++ b/src/coll/algorithms/alltoallv.cpp
@@ -366,6 +366,8 @@ ccl::status ccl_coll_build_topo_alltoallv(ccl_sched* main_sched,
     std::vector<size_t> send_counts, recv_counts, send_offsets, recv_offsets;
     size_t total_send_count = 0, total_recv_count = 0;
     size_t total_send_bytes = 0, total_recv_bytes = 0;
+    static ccl_buffer dummy_buf =
+        sched->alloc_buffer({ 1, ccl::buffer_type::ze, ccl::buffer_place::device });

     ccl_coll_calculate_alltoallv_counts(coll_param,
                                         send_counts,
@@ -385,7 +387,7 @@ ccl::status ccl_coll_build_topo_alltoallv(ccl_sched* main_sched,
         CCL_THROW_IF_NOT(send_counts == recv_counts, "unexpected send_counts");
         for (int idx = 0; idx < comm_size; idx++) {
             if (recv_counts[idx] == 0) {
-                recv_bufs[idx].set(nullptr, 0, 0);
+                recv_bufs[idx] = dummy_buf;
             }
             else {
                 recv_bufs[idx].set(coll_param.get_send_buf(), total_recv_bytes, recv_offsets[idx]);
@@ -396,7 +398,7 @@ ccl::status ccl_coll_build_topo_alltoallv(ccl_sched* main_sched,
         tmp_bufs.resize(comm_size);
         for (int idx = 0; idx < comm_size; idx++) {
             if (send_counts[idx] == 0) {
-                tmp_bufs[idx].set(nullptr, 0, 0);
+                tmp_bufs[idx] = dummy_buf;
             }
             else {
                 ccl::alloc_param alloc_param(send_counts[idx] * dtype.size(),
@@ -411,14 +413,14 @@ ccl::status ccl_coll_build_topo_alltoallv(ccl_sched* main_sched,
                          "unexpected send_counts");
         for (int idx = 0; idx < comm_size; idx++) {
             if (send_counts[idx] == 0) {
-                send_bufs[idx].set(nullptr, 0, 0);
+                send_bufs[idx] = dummy_buf;
             }
             else {
                 send_bufs[idx].set(coll_param.get_send_buf(), total_send_bytes, send_offsets[idx]);
             }

             if (recv_counts[idx] == 0) {
-                recv_bufs[idx].set(nullptr, 0, 0);
+                recv_bufs[idx] = dummy_buf;
             }
             else {
                 recv_bufs[idx].set(coll_param.get_recv_buf(), total_recv_bytes, recv_offsets[idx]);
@@ -629,7 +631,9 @@ ccl::status ccl_coll_build_topo_alltoallv(ccl_sched* main_sched,
                     auto peer_rank = (card_idx * tile_count + tile_idx);
                     if (peer_rank == comm->rank())
                         continue;
-                    if (counts[peer_rank] == 0)
+
+                    auto global_peer_rank = comm->get_global_rank(peer_rank);
+                    if (counts[global_peer_rank] == 0)
                         continue;

                     copy_attr attr{};
@@ -644,21 +648,16 @@ ccl::status ccl_coll_build_topo_alltoallv(ccl_sched* main_sched,
                     attr.hint_queue_index = copy_engine_idx;
                     attr.direction = copy_direction::c2c;

-                    if (!is_single_node) {
-                        // in order to get the correct offset for peer rank for example(2)
-                        peer_rank = peer_rank + (r2r_comm->rank() * node_comm->size());
-                    }
-
-                    auto src = bufs[peer_rank];
+                    auto src = bufs[global_peer_rank];
                     auto dst = ccl_buffer();
                     if (is_read) {
                         src = ccl_buffer();
-                        dst = bufs[peer_rank];
+                        dst = bufs[global_peer_rank];
                         LOG_DEBUG("ze_copy: read copy is enabled")
                     }

                     auto entry = entry_factory::create<ze_copy_entry>(
-                        sched, src, dst, counts[peer_rank], dtype, attr, wait_events);
+                        sched, src, dst, counts[global_peer_rank], dtype, attr, wait_events);
                     parallel_copy_events.push_back(entry->entry_event);
                 }
             }

[colleeneb]

Thanks @srirajpaul ! Do you know when this fix would be merged in oneCCL?

[srirajpaul]

@colleeneb it is merged to the next release.