[meta] UMAP fixes for 25.04

[jcrist]

Opening a meta issue for centralizing PRs and issues we want resolved before the 25.04 release (across all rapids repos). Please feel free to add new PRs/issues to the list as things come up.

cuml

• Correct UMAP dispatch trigger condition #6330
• UMAP fixes #6316 (this may require more investigation, non-controversial bits split out into Fix forwarding of initial_alpha and learning_rate in UMAP #6417)
• Use L2Expanded/L2SqrtExpanded instead of the *Unexpanded variants #6303
• Plumb metric and metric_kwds through to UMAP with nn_descent #6304
• Reduce peak memory in UMAP.fit/UMAP.fit_transform #6323
• Make UMAP callback pickleable #6402
• Fix forwarding of initial_alpha and learning_rate in UMAP #6417
• Fix UMAP transform #6449

raft

• Swap blocks and threads_per_block in compute_graph_laplacian raft#2597
• Fix COO symmetrization raft#2582 (this requires more investigation and is currently deemed non-necessary)
• Fix duplicate indices in batch NN Descent raft#2586

[jcrist]

(assigning myself just as the person tracking to ensure all these get in and we're in a good state pre-25.04 release, not as the person who should be doing all the work here :)).

[csadorf]

We should evaluate whether to fix #5282 as part of this release or future.

[jcrist]

I actually have a (mostly) finished fix for that. I don't think it's critical by any means, but it is easy to do. Will push up today.

[csadorf]

Ok, sounds good. Would be great if we could track that as part of the meta issue then.

[csadorf]

Let's add #6449 to this issue.

[jcrist]

Alright, all UMAP fixes are in.

To re-evaluate where we stand on memory usage, I ran UMAP on a 130M x 512 dataset (250 GiB) on an A100 with 80 GiB device memory. To do this successfully I had to enable managed memory via rmm. Eliding some details, the benchmark looked roughly like:

import rmm
import numpy as np
from cuml import UMAP

X = np.load("/path/to/data.npy", mmap_mode='r')
rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource())

model = UMAP(
    n_components=32,
    n_neighbors=15,
    build_aglo="nn_descent",
    init="random",
    build_kwds={"nnd_n_clusters": 150},
)
model.fit(X, data_on_host=True)

I instrumented the run with rmmscope to track all device memory allocations. This resulted in the following plot of memory usage (with annotations):

The full run took ~2 hours and memory usage peaked at ~130 GiB of managed memory. UMAP has a few stages, with the allocation logs from rmmscope the plot above can be broken down into:

• a to b: building the KNN graph.
• b to c: converting the KNN graph to a fuzzy simplicial set.
• c to d: creating the final embeddings

The peak memory usage (130 GiB) happens in the third stage during the generation of the embeddings. Currently peak memory usage can be computed roughly as:

def estimate_peak_memory(n_rows, n_neighbors, n_components):
    embedding = n_rows * n_components * 4
    # The true ratio here is structure dependent, but in my tests 0.85 was a decent appx
    ratio = 0.85
    graph = 3 * 2 * (n_rows * n_neighbors * 4)
    simpl_set_overhead = graph + (2 * graph // 3) * ratio
    return (embedding + graph + simpl_set_overhead) // (2 ** 30)

For this test run, the estimated peak is 127 GiB, which is pretty close to what we actually saw (130 GiB). As can be seen in the formula, only n_rows, n_neighbors, and n_components have a significant impact on device memory usage at this stage. Notably, the number of columns in the input data set is irrelevant, as is the memory of the original dataset - at this point we're only working with the constructed fuzzy simplicial set. I have a few ideas for reducing the memory usage of this stage a bit further (I think we can shave off another 3 * 2 * (n_rows * n_neighbors * 4) * ratio (~37 GiB in this run)), but at some point some of this is just algorithmically required. The good news is that using managed memory seems hold up quite well for this algorithm, so scaling up to some small factor larger than your true GPU device memory is possible.

A few recommendations for users running UMAP on large data:

• Use build_algo="nn_descent". This algorithm is approximate (rather than the default of brute_force_knn, which is exact), and is much more efficient for a large number of samples. It also supports batching, which can have a major impact on memory usage (more on this later).
• Load your data as a host array (e.g. numpy) and pass data_on_host=True to fit. This tells cuml.UMAP to keep the original dataset on host, only moving data to device as needed.
• If needed, you can also load your dataset as a mmap to avoid using too much host memory. cuml.UMAP can handle np.memmap objects transparently. Note that the performance of your disk or SSD then becomes more critical, so only do this if necessary.
• Set nnd_n_clusters in build_kwds. The default is 1, setting to a larger number enables batched generation of the knn graph. Larger numbers result in smaller batches and lower memory usage. Too large though and the quality of the results dips (we recommend no larger than 100, but smaller is better). The batching is data dependent, but you can loosely estimate the size of each batch as 2 * data_nbytes / nnd_n_clusters.

We're starting to think about the ergonomics of these hyperparameters and may make some changes in upcoming releases to automatically infer some settings for performance and quality without the user's input. But for now this should hopefully be enough to go on.

Since we can now demonstrate UMAP running on 100M+ datasets on an 80 GiB GPU, I'm going to close this out.