From b230ebe73f3fe7430a9e56199e8ffb54d2325c05 Mon Sep 17 00:00:00 2001 From: Jinsol Park Date: Tue, 22 Jul 2025 23:26:10 -0400 Subject: [PATCH] Add warning for unused GPU when `n_clusters < n_ranks` in batch `all_neighbors` (#1072) In MG all_neighbors, some GPUs stay idle when `n_clusters < n_ranks` (this is the expected behavior). Adding warning to let users know that they are not fully using all GPUs and providing directions for better utilization. Authors: - Jinsol Park (https://github.com/jinsolp) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1072 --- .../all_neighbors/all_neighbors_batched.cuh | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh index 6254b448dd..468d4ee0df 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh @@ -16,6 +16,7 @@ #pragma once #include "all_neighbors_builder.cuh" +#include "raft/core/logger_macros.hpp" #include #include #include @@ -336,6 +337,17 @@ void multi_gpu_batch_build(const raft::resources& handle, size_t base_cluster_idx = rank * clusters_per_rank + std::min((size_t)rank, rem); size_t num_clusters_for_this_rank = (size_t)rank < rem ? clusters_per_rank + 1 : clusters_per_rank; + if (num_clusters_for_this_rank == 0) { + // Happens in the case when num_ranks > n_clusters. + RAFT_LOG_WARN( + "Rank %d is not used for computation. This happens because the total number of ranks (%d) " + "> n_clusters (%lu). Consider increasing n_clusters or reduce the number of GPUs for " + "better utilization.", + rank, + num_ranks, + params.n_clusters); + continue; + } for (size_t p = 0; p < num_clusters_for_this_rank; p++) { num_data_for_this_rank += cluster_sizes(base_cluster_idx + p); }