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.

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

Reduce peak memory usage in single vector case #2

Draft
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

ml-evs
Copy link

@ml-evs ml-evs commented Mar 22, 2023

This PR updates both CPU and GPU versions of the dist matrix calculation such that the pairwise distances are accumulated in a 1D array of size (m(m-1) // 2 - m) rather than the full m**2 array, reducing peak memory usage by around half.

Also includes all the changes from #1, which should be merged first. The commit d5034b9 contains all the real changes of this PR.

Outstanding questions:

  • This basically changes the API of the single case to follow scipy.spatial.pdist rather than cdist. It would be simple enough to then create the big 2D array afterwards (by default?) if you want to keep the same API, although of course in the CPU case this requires even larger peak memory than the current status!
  • Similarly, if you want to make this optional to avoid any slowdowns in the low memory case, I'd be happy to make those changes. It could also be made dynamic by querying the memory available and choosing the appropriate method.
  • I may investigate an even more dynamic memory allocation approach, where chunks of the pairwise dist matrix are computed on the GPU and then saved into RAM (or disk) to allow for arbitrarily large matrices. This will depend on whether I can fit everything I personally want to do in my current 16 GB VRAM.

There's one outstanding wrinkle where the CUDA sparse calculation works but dense has some issue, will investigate soon.

@ml-evs ml-evs marked this pull request as draft March 22, 2023 10:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant