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?
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