Skip to content
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

Fix UMAP sparse input crash #5494

Open
wants to merge 2 commits into
base: branch-23.12
Choose a base branch
from
+16 −3
Open

Fix UMAP sparse input crash #5494

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6e93b15
Fix issue with limit resource adapter
viclafargue Oct 30, 2023
531757b
update
viclafargue Nov 1, 2023
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
19 changes: 16 additions & 3 deletions cpp/src/umap/knn_graph/algo.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,19 @@ inline void launcher(const raft::handle_t& handle,
const ML::UMAPParams* params,
cudaStream_t stream)
{
raft::resources tmp_handle(handle);
auto mr = raft::resource::get_workspace_resource(tmp_handle);
size_t free_size = raft::resource::get_workspace_free_bytes(tmp_handle);

double factor = 4.0;
size_t index_batch_size = inputsA.n;
size_t query_batch_size = inputsB.n;
size_t requirements = factor * sizeof(float) * index_batch_size * query_batch_size;

if (requirements > free_size) {
index_batch_size = free_size / (query_batch_size * factor * sizeof(float));
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fix looks good to me, can you just add a description of the fix to the PR description and fix the style issues?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry guys, this still isn't quite right.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@cjnolet, this is an older comment regarding the previous version. The current version should make use of the workspace_resource_factory to create a limiting_memory_resource with the default amount (total_size / 2 ?) and adjust the index batch size according to available memory. There may totally be things I missed, but could you give more precision to what is required here?

Copy link
Contributor Author

@viclafargue viclafargue Nov 1, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I updated the code slightly. But, since it's a RAFT call, the limiting_memory_resource cannot be used directly. The only thing that is doable is to update the handle, so that it has a resource factory that will in time be used to produce a workspace within the RAFT sparse kNN code. We can produce an estimate of the memory that will be available on the cuML side by calling resource::get_workspace_free_bytes on the handle or get_allocation_limit on the LMR. But, that's about it, isn't it?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@achirkin can you offer suggestions here? The main goal is to emable smart batching without relying on a query to get the free memory from the GPU (which is not acceptable practice). We started using the limiting_memory_resource_adapter so that we could do this while allowing the user to control this.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm trying to understand the context here. In isolation, this change sets the batch size based on the available memory. If the raft::sparse::selection::brute_force_knn properly uses the workspace memory resource (which I hope it does but cannot check on my phone atm), then the solution looks good to me. Indeed, the workspace resource is always a limiting memory resource (the type is erased) and resource::get_workspace_free_bytes is a correct way to gets its remaining limit.

But I fail to see the connection to the linked issue. There, an invalid memory access error is reported, as suggested, likely due to integer overflow. If the workspace size is larger than 4gb, that still could be an issue, right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for taking a look @achirkin. The changes in rapidsai/raft#1640 in addition to these changes should solve the issue. However, to allow any batch size to be used, it would be necessary to apply modifications to ensure that the integer values of csr_batcher_t and distances_config_t are also stored on 64 bits in order to guarantee that no operation would be made on 32 bits. This was much a more involved task, this is why I left it for later.


raft::sparse::selection::brute_force_knn(inputsA.indptr,
inputsA.indices,
inputsA.data,
Expand All @@ -115,9 +128,9 @@ inline void launcher(const raft::handle_t& handle,
out.knn_indices,
out.knn_dists,
n_neighbors,
handle,
ML::Sparse::DEFAULT_BATCH_SIZE,
ML::Sparse::DEFAULT_BATCH_SIZE,
tmp_handle,
index_batch_size,
query_batch_size,
params->metric,
params->p);
}
Expand Down