adressing review

Author: viclafargue

cpp/include/cuvs/cluster/kmeans.hpp

@@ -125,6 +125,10 @@ struct params : base_params {
    * When set to 0 (default) with host data uses `min(3 * n_clusters, n_samples)`
    * as a default.
    *
+   * In Batched multi-GPU host-data fits, the effective KMeansPlusPlus initialization
+   * sample is materialized on device and seeding is run on rank 0. Rank 0
+   * must have enough GPU memory for this sample and the seeding workspace.
+   *
    * Default: 0.
    */
   int64_t init_size = 0;
@@ -186,6 +190,10 @@ enum class kmeans_type { KMeans = 0, KMeansBalanced = 1 };
  *     this worker's partition, and RAFT communicators are used for collectives.
  *   - Otherwise: single-GPU batched k-means.
  *
+ * With `params.init == InitMethod::KMeansPlusPlus` in multi-GPU mode, the
+ * effective initialization sample must fit in rank 0 GPU memory because
+ * seeding is run on rank 0 before centroids are broadcast.
+ *
  * @code{.cpp}
  *   #include <raft/core/resources.hpp>
  *   #include <cuvs/cluster/kmeans.hpp>

cpp/src/cluster/detail/kmeans_mg_batched.cuh

@@ -318,10 +318,7 @@ void mnmg_fit(const raft::resources& handle,
                              true);
   }
 
-  bool need_compute_norms = metric == cuvs::distance::DistanceType::L2Expanded ||
-                            metric == cuvs::distance::DistanceType::L2SqrtExpanded;
-  auto h_norm_cache =
-    raft::make_pinned_vector<T, IdxT>(dev_res, (need_compute_norms && has_data) ? n_local : 0);
+  auto h_norm_cache = raft::make_pinned_vector<T, IdxT>(dev_res, has_data ? n_local : 0);
   bool norms_cached = false;
 
   auto mnmg_allreduce = [&](auto* sendbuf, auto* recvbuf, size_t count) {
@@ -407,21 +404,19 @@ void mnmg_fit(const raft::resources& handle,
           auto L2NormBatch_view =
             raft::make_device_vector_view<T, IdxT>(L2NormBatch.data_handle(), current_batch_size);
 
-          if (need_compute_norms) {
-            auto batch_offset = static_cast<IdxT>(data_batch.offset());
-            if (!norms_cached) {
-              raft::linalg::norm<raft::linalg::L2Norm, raft::Apply::ALONG_ROWS>(
-                dev_res, batch_data_view, L2NormBatch_view);
-              raft::copy(h_norm_cache.data_handle() + batch_offset,
-                         L2NormBatch.data_handle(),
-                         current_batch_size,
-                         stream);
-            } else {
-              raft::copy(L2NormBatch.data_handle(),
-                         h_norm_cache.data_handle() + batch_offset,
-                         current_batch_size,
-                         stream);
-            }
+          auto batch_offset = static_cast<IdxT>(data_batch.offset());
+          if (!norms_cached) {
+            raft::linalg::norm<raft::linalg::L2Norm, raft::Apply::ALONG_ROWS>(
+              dev_res, batch_data_view, L2NormBatch_view);
+            raft::copy(h_norm_cache.data_handle() + batch_offset,
+                       L2NormBatch.data_handle(),
+                       current_batch_size,
+                       stream);
+          } else {
+            raft::copy(L2NormBatch.data_handle(),
+                       h_norm_cache.data_handle() + batch_offset,
+                       current_batch_size,
+                       stream);
           }
 
           auto L2NormBatch_const = raft::make_const_mdspan(L2NormBatch_view);
@@ -447,7 +442,7 @@ void mnmg_fit(const raft::resources& handle,
             raft::make_device_scalar_view(clustering_cost.data_handle()),
             batch_workspace);
         }
-        if (need_compute_norms) { norms_cached = true; }
+        norms_cached = true;
       }
 
       // Phase 2: grouped allreduce

cpp/src/cluster/detail/kmeans_mg_batched_init.cuh

@@ -232,6 +232,8 @@ void init_centroids_for_mg_batched(raft::resources const& handle,
       handle, params, X_local, init_sample_size, rank, root, rank_counts, allreduce, bcast);
 
     if (rank == root) {
+      // KMeans++ seeding runs on root, so root must have enough device memory
+      // for the full initialization sample plus the seeding workspace.
       auto init_view = raft::make_const_mdspan(init_sample.view());
       if (params.oversampling_factor == 0) {
         cuvs::cluster::kmeans::detail::kmeansPlusPlus<DataT, IndexT>(