Skip to content

Commit

Permalink
Avoid limited memory adaptor issue in balanced KMeans (#2570)
Browse files Browse the repository at this point in the history
- Switch to the use of `get_large_workspace_resource` instead of `get_workspace_resource`
- Do not use explicit managed memory allocation.

Based on and merge after #2541 ([diff](viclafargue/raft@fix-sparse-utilities...csadorf:raft:sadorf/address-limited-memory-adaptor-issue))

Authors:
  - Simon Adorf (https://github.com/csadorf)

Approvers:
  - Artem M. Chirkin (https://github.com/achirkin)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #2570
  • Loading branch information
csadorf authored Feb 13, 2025
1 parent 0d03e15 commit 842afd7
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions cpp/include/raft/cluster/detail/kmeans_balanced.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -967,9 +967,10 @@ void build_hierarchical(const raft::resources& handle,
IdxT n_mesoclusters = std::min(n_clusters, static_cast<IdxT>(std::sqrt(n_clusters) + 0.5));
RAFT_LOG_DEBUG("build_hierarchical: n_mesoclusters: %u", n_mesoclusters);

// TODO: Remove the explicit managed memory- we shouldn't be creating this on the user's behalf.
// Need to use explicit managed_memory here since corresponding allocations
// must be both host and device accessible.
rmm::mr::managed_memory_resource managed_memory;
rmm::device_async_resource_ref device_memory = resource::get_workspace_resource(handle);
rmm::device_async_resource_ref device_memory = resource::get_large_workspace_resource(handle);
auto [max_minibatch_size, mem_per_row] =
calc_minibatch_size<MathT>(n_clusters, n_rows, dim, params.metric, std::is_same_v<T, MathT>);

Expand Down

0 comments on commit 842afd7

Please sign in to comment.