Avoid large device allocation in UMAP with nndescent #6292
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Currently `NNDescent` returns two arrays: - `graph.graph()`: (n x graph_degree) on host - `graph.distances()`: (n x graph_degree) on device Downstream, the rest of UMAP wants both of these to be device arrays of shape (n x n_neighbors). Currently we copy `graph.graph()` to a temporary device array, then slice and and copy it to the output array `out.knn_indices`. Ideally we'd force `graph_degree = n_neighbors` to avoid the slicing entirely (and reduce the size of the intermediate results). However, it seems like currently there's a bug in `NNDescent` where reducing `graph_degree` to `n_neighbors` causes a significant decrease in result quality. So for now we need to keep the slicing around. We can avoid allocating the temporary device array though, instead doing the slicing on host. Doing this avoids allocating a (n x graph_degree) device array entirely; for large `n` this can be a significant savings (47 GiB on one test problem I was trying). We still should fix the `graph_degree` issue, but for now this should help unblock running UMAP on very large datasets.
jcrist
added
non-breaking
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## branch-25.02 #6292 +/- ##
================================================
- Coverage 68.18% 67.06% -1.13%
================================================
Files 201 202 +1
Lines 13055 13065 +10
================================================
- Hits 8902 8762 -140
- Misses 4153 4303 +150 ☔ View full report in Codecov by Sentry. |
divyegala
approved these changes
Feb 5, 2025
wphicks
approved these changes
Feb 5, 2025
|
/merge |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Currently
NNDescentreturns two arrays:graph.graph(): (n x graph_degree) on hostgraph.distances(): (n x graph_degree) on deviceDownstream, the rest of UMAP wants both of these to be device arrays of shape (n x n_neighbors).
Currently we copy
graph.graph()to a temporary device array, then slice and and copy it to the output arrayout.knn_indices.Ideally we'd force
graph_degree = n_neighborsto avoid the slicing entirely (and reduce the size of the intermediate results). However, it seems like currently there's a bug inNNDescentwhere reducinggraph_degreeton_neighborscauses a significant decrease in result quality. So for now we need to keep the slicing around.We can avoid allocating the temporary device array though, instead doing the slicing on host. Doing this avoids allocating a (n x graph_degree) device array entirely; for large
nthis can be a significant savings (47 GiB on one test problem I was trying).We still should fix the
graph_degreeissue, but for now this should help unblock running UMAP on very large datasets.