Reduce peak memory in UMAP.fit/UMAP.fit_transform
#6323
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
By dropping intermediate objects earlier we can reduce peak memory usage. Also reduces some code duplication between `_get_graph` and `_get_graph_supervised`.
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
jcrist
added
improvement
|
Hold off on merging this - the commits ran fine on a 270 GiB dataset yesterday, but after rebasing on #6314 there's now an underflow happening on large data (I think from |
This is not a bug in the change here, but rather a regression added in #6314. See my comment here: #6314 (comment). |
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.
This reduces the peak GPU memory usage during a
UMAP.fitcall by releasing certain temporaries as soon as possible.This required splitting the
FuzzySimplSet::runfunction into a few sub-functions and intermingling them within_get_graphto let us drop the temporary arrays earlier. While I was at it, I noticed that_get_graphand_get_graph_supervisedwere almost identical, so I've merged these paths a bit to reduce duplication. There's a lot more duplication in here that could be cleaned up (and I went down that rabbit hole a bit before stopping myself), but for now I think this should be sufficient.On an input of 190 GiB this dropped 12 GiB off the peak, and moved the high point into the
SimplSetEmbedroutines instead of inFuzzySimplSet. Further memory improvements will come in a later PR.