Fix UMAP sparse input crash #5494
Conversation
| size_t requirements = factor * sizeof(float) * index_batch_size * query_batch_size; | ||
| if (requirements > free_mem){ | ||
| index_batch_size = free_mem / (query_batch_size * factor * sizeof(float)); | ||
| } |
There was a problem hiding this comment.
Fix looks good to me, can you just add a description of the fix to the PR description and fix the style issues?
There was a problem hiding this comment.
Sorry guys, this still isn't quite right.
There was a problem hiding this comment.
@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?
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
@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.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
cpp/src/umap/knn_graph/algo.cuh
Outdated
| @@ -101,6 +101,16 @@ inline void launcher(const raft::handle_t& handle, | |||
| const ML::UMAPParams* params, | |||
| cudaStream_t stream) | |||
| { | |||
| size_t free_mem, total_mem; | |||
| cudaMemGetInfo(&free_mem, &total_mem); | |||
There was a problem hiding this comment.
This is, unfortunately, not a reliable way to query and compute the memory available for batching. It also won't work well with rmm pools.
The challenge here is that the total memory space can quickly become very fragmented, which prevents the reporter free memory from being allocated in a single chunk. Further, if 90% of the memory is allocated to a pool, this is going to report almost no memory available even though the pool could be largely empty.
A better way to approach this problem will be to use the limiting memory resource once its available. This new resource will be applied on top of the workspace resource and will provide a safe and consistent mechanism for us to handle implicit batching on the users behalf in a way that also plays nicely with RMM pools.
There was a problem hiding this comment.
@cjnolet, the current approach of this PR was to first compute the memory requirements, and then adjust the index batch size if necessary. I understand that getting the free available memory directly from CUDA API is not a good solution as a RMM pool that could/should be used might have already been pre-allocated.
The limiting_resource_adaptor might provide a way to limit downstream memory usage of an upstream device_memory_resource. However, I do not see how this will help solving this issue. The issue is that a complete sparse operation consumes too much memory at once and should be split in multiple operations through batching instead.
Could you give more precision as to how it will provide a safe and consistent mechanism for us to handle implicit batching on the users behalf in a way that also plays nicely with RMM pools?
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
github-actions
bot
added
conda
viclafargue
force-pushed
the
fix-umap-sparse-input
branch
from
082ec5a to
6e93b15
Compare
github-actions
bot
removed
conda
Answers #4984