Investigate memory footprint of UMAP

[jinsolp]

Investigate the memory footprint of UMAP with the goal of:

• Determining the largest dataset size that can be processed with the current implementation
• Identifying key memory bottlenecks and outlining actionable plans to reduce or eliminate them

[jinsolp]

This is not an actionable release item, but adding it to the Project board for tracking purposes

[jinsolp]

Memory analysis of UMAP as of Dec 5, 2025.
This is for the fit_transform function using init="random", and assumes that kNN construction is not the memory bottleneck (by using the batching feature).

Peak Memory Analysis for fit_transform

There are two possible peak memory points depending on UMAP parameters.

Peak1: coo_sort after symmetrize

cuml/cpp/src/umap/runner.cuh

Lines 207 to 216 in 7a8a9a8

	raft::sparse::COO<value_t> fss_graph(stream);
	{
	raft::sparse::COO<value_t> strengths(stream);
	_get_strengths<value_idx, value_t, umap_inputs, nnz_t, TPB_X>(
	handle, inputs, params, strengths);
	FuzzySimplSetImpl::symmetrize<value_t>(strengths, fss_graph, params->set_op_mix_ratio, stream);
	} // end strengths scope
	/* Canonicalize output graph */
	raft::sparse::op::coo_sort<value_t>(&fss_graph, stream);

The only device memory at this step is the fss_graph. However, coo_sort function internally uses a buffer that results in a spike in memory usage.

Peak2: Optimize Embedding

cuml/cpp/src/umap/runner.cuh

Lines 318 to 329 in 7a8a9a8

	// Allocate embeddings buffer just before initialization
	std::size_t embeddings_size =
	static_cast<std::size_t>(inputs.n) * params->n_components * sizeof(value_t);
	if (!embeddings) { embeddings = std::make_unique<rmm::device_buffer>(embeddings_size, stream); }
	value_t* embeddings_ptr = static_cast<value_t*>(embeddings->data());
	/**
	* Run initialization method
	*/
	raft::common::nvtx::push_range("umap::embedding");
	InitEmbed::run<value_t, nnz_t>(
	handle, inputs.n, inputs.d, &graph, params, embeddings_ptr, stream, params->init);

Device memory at this step includes 1) the final graph 2) final embedding 2) some additional arrays for tracking.
This is less likely to be the peak memory unless n_components is really big.

Script for Computing Peak Memory

Ratios 0.8 and 0.99 are dependent on the data. change to larger values to be conservative.

type_to_bytes = {"char":1, "half": 2, "int": 4, "float": 4, "int64": 8, "uint64": 8}

CooIdxT = "int"
ValueT = "float"
IdxT = "int64"
nnzT = "uint64"

def get_peak_mem_bytes_new(n_rows, n_neighbors, n_components, init = "random"):
    nnz = n_rows * n_neighbors
    symmetric_nnz = nnz * 2
    after_remove_zeros_nnz = symmetric_nnz * 0.8
    after_trimming_nnz = after_remove_zeros_nnz * 0.99
    coo_byte_per_nnz = 2 * type_to_bytes[CooIdxT] + type_to_bytes[ValueT]
    if init == "random":
        # first peak coming from coo_sort after symmetrization
        coo_sort_overhead = nnz * 24
        symmetric_graph_size = symmetric_nnz * coo_byte_per_nnz
        peak1 = coo_sort_overhead + symmetric_graph_size

        # second peak coming from simplicial_set
        embedding_size = n_rows * n_components * type_to_bytes[ValueT]
        epochs_per_sample_size = after_trimming_nnz * 3 * type_to_bytes[ValueT]
        final_graph_size = after_trimming_nnz * coo_byte_per_nnz

        peak2 = embedding_size + epochs_per_sample_size + final_graph_size
        return max(peak1, peak2)

Notes

• current peak is max(peak1, peak2). If we manage to tackle peak1, this would become max(remove zeros after sort, peak2).