feat(semdedup): bound pairwise precision and memory by praateekmahajan · Pull Request #2318 · NVIDIA-NeMo/Curator

praateekmahajan · 2026-08-17T20:28:30Z

Summary

Implements NMCUR-318 on current main:

bounds Pairwise similarity memory to a reusable N x B workspace with fixed positive pairwise_batch_size (default 1024);
compares every ranked row only with valid earlier-ranked neighbors, preserving deterministic earliest-rank tie behavior;
keeps arbitrary ranking columns and the existing id, max_id, and cosine_sim_score output contract;
separates KMeans embedding storage precision from Pairwise multiplication precision;
supports opt-in FP16 storage/compute while preserving existing FP32 storage and FP32 compute as the defaults;
retains FP16 through multiplication/reduction and promotes scores to FP32 only for cuDF output;
routes Torch allocations through the RMM resource already used by cuDF/CuPy, without resetting process-global allocator state;
clears failed traceback frames so references to large GPU allocations are released;
records cumulative footer-scan, read, rank, conversion, compute, and write timings.

The ranked-neighbor sidecar remains an extension point only. This PR does not implement NMCUR-242 or NMCUR-319.

EOS benchmark

Compared PR commit 5e5b8248 with the September 4 main nightly on the 329,774,475-document predict-all workload:

Configuration	Workflow	Pairwise	Duplicates	Retained artifacts
Main nightly	1122.8s	740.6s	10,819,452	—
PR FP32 storage / FP32 compute	786.9s	423.0s	10,819,441	1.025 TB
PR FP32 storage / FP16 compute	539.6s	184.9s	10,823,919	1.024 TB
PR FP16 storage / FP16 compute	497.0s	154.3s	10,823,956	498.7 GB

The bounded FP32 path reduced Pairwise wall time by 42.9% with an 11-document duplicate-count difference. FP16 remains opt-in pending duplicate-ID analysis; the benchmark indicates substantial performance and storage potential.

The benchmark's original per-cluster custom metrics were under-aggregated across executor task batches. Workflow and Pairwise wall times above are valid; this PR now aggregates the granular phase timings per executor batch and removes the redundant row-count and resolved-batch metrics.

The six-entry precision matrix completed successfully for this analysis. The committed nightly YAML now retains only the two existing SemDedup baseline entries, explicitly configured for the FP32 storage/FP32 compute defaults.

An allocator-only EOS A/B is in progress for both committed SemDedup entries:

before RMM sharing: pr-2318-2026_09_04__22_24_57_UTC-5d9be38a
with RMM sharing: pr-2318-2026_09_04__22_37_45_UTC-fb6bbdbb

A local 330,000-row × 768-dimension Pairwise workload reduced peak GPU memory from 31.83 GiB to 3.49 GiB with RMM sharing; its median compute time was about 7% slower, which is why the EOS A/B is being run before drawing a performance conclusion.

Validation

At fb6bbdbb, using nvcr.io/nvidian/nemo-curator:nightly-2026-09-04 with GPUs 2 and 3 only:

Pairwise, KMeans, and semantic-dedup workflow tests: 60 passed, 1 skipped
fresh-process RMM allocator and Pairwise functional smoke: passed
Ruff check and format check: passed
scoped pre-commit checks for changed files: passed

copy-pr-bot · 2026-08-17T20:28:33Z

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.

github-actions · 2026-08-17T20:30:59Z

🌿 Preview your docs: https://nvidia-preview-praateekm-nmcur-318-semdedup-pairwise-precision.docs.buildwithfern.com/nemo/curator

Here are the markdown pages you've updated:

nemo/curator/curate-text/process-data/deduplication/semdedup

praateekmahajan · 2026-08-18T18:10:17Z

        return None
+
+
+def release_cached_gpu_memory() -> None:


@ayushdg can you review this, to see if it's not sloppy? I'm wondering if we can then also add this in the minhash teardown part

praateekmahajan · 2026-08-21T01:31:32Z

/ok to test cd18cc1

praateekmahajan · 2026-08-21T18:28:08Z

/ok to test bcb544b

ayushdg · 2026-09-03T17:32:42Z

        eps: Epsilon value for duplicate identification threshold (cosine_sim >= 1-eps)
        which_to_keep: Strategy for ranking within clusters ("hard", "easy", "random")
-        pairwise_batch_size: Batch size for pairwise similarity computation
+        pairwise_compute_dtype: Multiplication precision used by Pairwise


nit: Should we rename this to precision?

Can do it but I think it could become ambiguous.

KMeans outputs FP32 or 16 and this one can perform matmul in FP32 or 16, both of which are "precision" args.

ayushdg

At a high level the concept makes sense to me. Just had high level questions & few nits

ayushdg · 2026-09-03T18:35:28Z

        return None
+
+
+def release_cached_gpu_memory(*, reset_owned_rmm_pool: bool = False) -> None:


Not an expert on memory stuff but I know there's a way to ask cupy and possibly torch to re-use the same rmm pool being used by cudf. Not sure if we use that or if that needs to be handled separately.

Signed-off-by: Praateek <praateekm@gmail.com>

sarahyurick · 2026-09-04T22:39:40Z

+        """Make Torch allocate from the RMM resource already used by cuDF and CuPy."""
+        torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
+
+    def process(self, task: FileGroupTask) -> FileGroupTask:


Nit but do we still need process if we are doing process_batch now?

sarahyurick · 2026-09-04T22:40:45Z

+        except BaseException as exc:
+            # An exception's traceback otherwise keeps the unwound _process
+            # frame—and its large GPU objects—alive until after this finalizer.
+            traceback.clear_frames(exc.__traceback__)


Sorry what is the point of this?

sarahyurick · 2026-09-04T22:43:40Z

+        conversion_start = time.perf_counter()
+        resolved_compute_dtype = _resolve_compute_dtype(cluster_embeddings, self.compute_dtype)
+        storage_was_converted = cluster_embeddings.dtype != resolved_compute_dtype
+        cluster_embeddings = cluster_embeddings.to(dtype=resolved_compute_dtype)


So you are doing _resolve_compute_dtype and to(dtype=...) here and then again inside pairwise_cosine_similarity_batched? Are they the same thing each time?

sarahyurick · 2026-09-04T22:44:24Z

+    ``B``, ``B`` over ``C``, and so on. For example, consider these normalized
+    four-dimensional embeddings::
+
+        X = A [1.00, 0.00, 0.00, 0.00]


This seems like a nice example, is it something we should consider explaining in the docs too?