⭐ Football Update: NVIDIA NeMo/Curator/pull/

praateekmahajan · 2026-09-02T21:06:10Z

Summary

Read Parquet footer metadata to group files by actual embedding element counts before loading them with cuDF.
Preallocate KMeans and pairwise embedding arrays from exact row and element counts.
When fit_data_fraction=None, select complete fit files within a budget derived from currently available GPU memory. Explicit fractions retain file-level sampling.
Predict labels for every input row with KMeansMG.predict, including rows not selected for fitting.
Remove embedding_dim from the semantic-deduplication API, benchmarks, tutorials, and documentation.
Keep Parquet writes synchronous; asynchronous writes are scoped to perf(semantic-dedup): overlap parquet writes #2361.

Behavior and constraints

Read groups stay below cuDF's embedding-column element limit. Files are kept intact, so each individual Parquet file must itself fit below that limit.

Benchmark results

EOS benchmark on 8× H100 GPUs. Baseline: main at 07920606; PR measurement: aa7529ea. The current head only simplifies the JSONL fallback and benchmark reporting; the measured Parquet path is unchanged.

20% dataset, full fit — 65.8M rows

Open the run comparison

Metric	Main	PR	Change
End-to-end workflow	218.2 s	199.5 s	18.7 s faster (−8.6%)
KMeans stage	138.0 s	125.5 s	12.5 s faster (−9.1%)
KMeans actor processing	105.4 s	94.3 s	11.1 s faster (−10.5%)
KMeans read	33.9 s	24.2 s	9.7 s faster (−28.7%)
KMeans fit + predict	38.5 s	35.3 s	3.2 s faster (−8.2%)
KMeans output	33.0 s	22.6 s	10.4 s faster (−31.6%)
Pairwise stage	80.2 s	74.0 s	6.2 s faster (−7.8%)
Run-average GPU utilization	28.3%	31.3%	+3.0 pp
Run-average GPU memory	27.8%	19.3%	−8.5 pp (−30.6%)
Peak GPU memory	82.9%	82.4%	Flat (−0.5 pp)

The steady KMeans fit plateau drops from about 66% to 32% GPU memory, showing that exact preallocation removes the previous duplicate full embedding buffer. The similar absolute peak occurs later in distance calculation and partitioned Parquet output.

Full dataset, 20% fit — 329.8M rows

Open the run comparison

Metric	Main	PR	Change
End-to-end workflow	20m 54s	18m 37s	2m 17s faster (−11.0%)
KMeans stage	7m 43s	6m 13s	1m 30s faster (−19.4%)
KMeans actor processing	414.2 s	316.7 s	97.5 s faster (−23.5%)
KMeans read	201.3 s	156.0 s	45.3 s faster (−22.5%)
KMeans fit	40.2 s	35.0 s	5.2 s faster (−12.9%)
KMeans predict + output	172.6 s	110.1 s	62.5 s faster (−36.2%)
Pairwise stage	13m 11s	12m 23s	47.6 s faster (−6.0%)
Run-average GPU utilization	48.2%	55.9%	+7.8 pp
Run-average GPU memory	33.1%	14.3%	−18.8 pp (−56.8%)
Peak GPU memory	58.5%	78.0%	+19.5 pp

The partial-fit run uses substantially less memory on average and keeps GPUs busier, but its instantaneous peak is higher during output. Reducing that output-path peak is intentionally scoped to the asynchronous/bounded writer work in #2361.

Validation

Semantic KMeans, pairwise, and utility GPU tests: 55 passed, 1 skipped
Pre-commit hooks

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

copy-pr-bot · 2026-09-02T21:06:14Z

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-09-02T21:08:24Z

🌿 Preview your docs: https://nvidia-preview-codex-semdedup-bounded-read-memory.docs.buildwithfern.com/nemo/curator

Here are the markdown pages you've updated:

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

praateekmahajan · 2026-09-02T21:48:31Z

/claude review
@greptileai review

praateekmahajan · 2026-09-02T21:48:33Z

/ok to test d25112c

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

praateekmahajan · 2026-09-02T22:29:31Z

/ok to test 05fd74f

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

sarahyurick · 2026-09-03T00:03:38Z

-            return self._process_batch_two_pass(tasks, groups)
-        return self._process_batch_single_pass(tasks, groups)
+            return self._process_batch_two_pass(tasks, all_files)
+        return self._process_batch_single_pass(tasks, all_files)


Why we are keeping both of _process_batch_two_pass and _process_batch_single_pass if the new behavior defaults to "however much fits in 60% of GPU memory" (if I am understanding correctly)?

Is it only because of Parquet vs JSONL?

Yes - these are now exclusively the JSONL paths. None preserves JSONL’s full single-pass behavior (because we can't infer the number of rows without reading the files), while an explicit fraction requires two passes. I’ll rename them to make that distinction obvious.

sarahyurick · 2026-09-03T00:07:13Z

+        gc.collect()
+        cp.get_default_memory_pool().free_all_blocks()
+
+        if remaining_info:


Nit but to me remaining_info wasn't clear that it is "the remaining data that we didn't fit with", could we rename and/or add comments about that?

prediction_only_info sounds good?

sarahyurick · 2026-09-03T00:11:41Z

+        fit_frames = iter(self._iter_parquet_frames(fit_info, columns))
+        first_fit_frame = next(fit_frames)
+        embedding_width = get_array_from_df(first_fit_frame, self.embedding_field).shape[1]
+        fit_embeddings = cp.empty((fit_rows, embedding_width), dtype=cp.float32)
+        sampled_chunks: list[tuple[cudf.DataFrame, int, int]] = []
+        read_time = 0.0
+        offset = 0
+        read_start = time.perf_counter()


I thought the read starts at first_fit_frame = ...? Shouldn't this timer be up there?

Good catch! It misses that one read.

sarahyurick · 2026-09-03T00:12:37Z

+            stop = offset + len(df)
+            embeddings = fit_embeddings[offset:stop]
+            embeddings[...] = get_array_from_df(df, self.embedding_field)
+            embeddings /= cp.linalg.norm(embeddings, axis=1, keepdims=True)


Sorry why don't we want normalize_embeddings_col_in_df here?

That helper creates another full normalized Torch array and assigns it back into cuDF. Whereas the current proposed path copies into the exactly allocated CuPy fit buffer and normalizes each slice in place, which is what removes the ~2× fit-memory plateau.
I'll see if we can get rid of the normalize_embeddings_col_in_df or rather put this code there

sarahyurick · 2026-09-03T00:15:29Z

+                read_start = time.perf_counter()
+            read_time += time.perf_counter() - read_start
+
+        del self.kmeans


Why do we delete here? Should we do it for the JSONL path too?

Actually no need to remove it, because the actor pool anyway dies after this. So can remove it

sarahyurick · 2026-09-03T00:30:06Z

-        # Default aggressive assumption of 1024 dimensional embedding
-        embedding_dim = 1024
+def _root_column(path: str | list[str]) -> str:
+    return path[0] if isinstance(path, list) else path.split(".", maxsplit=1)[0]


Couldn't this fail if a real column actually has a dot in it?

Yea, also TBH i don't find this code super clear, so will update it to be clearer

sarahyurick · 2026-09-03T00:30:30Z

-    max_files_per_subgroup = max(1, max_files_per_subgroup)  # Ensure at least 1 file per subgroup
+    retained = set(retained_columns or [])
+    try:
+        if storage_options or get_protocol(files[0]) != "file":


Curious have you done an S3 benchmark?

Nope sadly not

sarahyurick · 2026-09-03T00:34:54Z

+                    column.meta_data.total_uncompressed_size
+                    for row_group in footer.row_groups
+                    for column in row_group.columns
+                    if _root_column(column.meta_data.path_in_schema) in retained


Nit but it does seem like we could reduce read_parquet_file_info's complexity here? Like maybe I am reading it wrong but the _root_column logic here and in the previous code block are both running the same logic right (in terms of nested for loops and if blocks)?

Yea, sadly we use fsspec for remote and pylibcudf for local which both differ in the metadata's object that they return. So math might be the same however attribute handling still remains different

sarahyurick · 2026-09-03T00:36:47Z

+                    if _root_column(column.meta_data.path_in_schema) in retained
+                )
+                embedding_elements = sum(
+                    column.meta_data.num_values


Not too familiar with what is the difference between using column.meta_data.num_values here and column.num_values above?

Same as above, different metadata objects, PyArrow returns ColumnChunkMetaData directly; pylibcudf returns a ColumnChunk whose metadata is under .meta_data, can rename them though

sarahyurick · 2026-09-03T00:36:59Z

+                result.append(ParquetFileInfo(path, footer.num_rows, metadata_bytes, embedding_elements))
+        return result  # noqa: TRY300
+    except Exception as error:
+        msg = f"Failed to read Parquet footer metadata for one of {len(files)} files"


Should we try to say which file(s) failed?

The batch might have 64 files which is too many to report, it might be an overkill to find which of the 64 failed. lmk what you think

Yeah I guess not. I am just wondering about when this error is raised, what is a good way for the user to figure out how to fix their file(s) and/or file(s) to exclude in the next run?

Today we have this issue even with parquet reads, a corrupt file doesn't share which file caused the failure. So I think we can let it be, added a TODO tho in case we run into it

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

praateekmahajan · 2026-09-03T16:47:16Z

/claude review
@greptileai review

sarahyurick · 2026-09-03T17:44:09Z

+                result.append(ParquetFileInfo(path, footer.num_rows, metadata_bytes, embedding_elements))
+        return result  # noqa: TRY300
+    except Exception as error:
+        msg = f"Failed to read Parquet footer metadata for one of {len(files)} files"


Yeah I guess not. I am just wondering about when this error is raised, what is a good way for the user to figure out how to fix their file(s) and/or file(s) to exclude in the next run?

sarahyurick · 2026-09-03T17:46:00Z

+
+
+@pytest.mark.gpu
+def test_parquet_file_info_uses_fsspec_for_remote_uri_without_storage_options() -> None:


I was wondering if we can remove the GPU coupling here yet or not? Even though it is used for semantic dedup only, it is still IO.

Yea we have a todo below for TODO : Remove this once we figure out how to import semantic on CPU

Maybe I should make it top-level..

The whole imports cause this issue, and would prefer not to resolve that in this PR

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

greptile-apps · 2026-09-03T19:59:07Z

-    subgroups = [files[i : i + max_files_per_subgroup] for i in range(0, len(files), max_files_per_subgroup)]
+    retained = set(retained_columns or [])
+    try:
+        if storage_options or get_protocol(files[0]) != "file":


Mixed URI footers bypass fsspec

When a list-valued input combines local and remote Parquet locations without storage options and a local file comes first, this condition selects the pylibcudf backend for the entire actor batch. Later remote files are then passed to SourceInfo instead of fsspec, causing the mandatory footer scan to fail before KMeans runs.

perf(semantic-dedup): bound parquet reads and memory

e26b506

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

praateekmahajan mentioned this pull request Sep 2, 2026

perf(semantic-dedup): overlap parquet writes #2361

Closed

praateekmahajan added 6 commits September 2, 2026 21:15

refactor(semantic-dedup): sample fit rows and use KMeans predict

33b6459

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

refactor(semantic-dedup): infer embedding size from parquet

b9a276c

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

refactor(semantic-dedup): name discovered embedding width

e5fa5c2

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

revert(semantic-dedup): retain file-fraction sampling

2eb0889

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

docs(semantic-dedup): remove embedding dimension tuning

4560c25

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

refactor(semantic-dedup): use ordinary bounded parquet reads

d25112c

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

copy-pr-bot Bot temporarily deployed to public September 2, 2026 21:48 Inactive

copy-pr-bot Bot temporarily deployed to test September 2, 2026 21:49 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci September 2, 2026 21:49 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci September 2, 2026 21:49 Error

copy-pr-bot Bot temporarily deployed to nemo-ci September 2, 2026 21:49 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci September 2, 2026 22:09 Inactive

copy-pr-bot Bot temporarily deployed to public September 2, 2026 22:11 Inactive

copy-pr-bot Bot temporarily deployed to public September 2, 2026 22:12 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci September 2, 2026 22:20 Error

copy-pr-bot Bot temporarily deployed to nemo-ci September 2, 2026 22:20 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci September 2, 2026 22:20 Error

copy-pr-bot Bot temporarily deployed to nemo-ci September 2, 2026 22:20 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci September 2, 2026 22:20 Failure

copy-pr-bot Bot had a problem deploying to nemo-ci September 2, 2026 22:20 Error

test(semantic-dedup): exercise all kmeans paths

05fd74f

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

copy-pr-bot Bot temporarily deployed to public September 2, 2026 22:29 Inactive

copy-pr-bot Bot temporarily deployed to test September 2, 2026 22:30 Inactive

refactor(semdedup): simplify bounded parquet paths

efeaf2a

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

sarahyurick reviewed Sep 3, 2026

View reviewed changes

fix(semdedup): address review feedback

2a8e193

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

Merge branch 'main' into codex/semdedup-bounded-read-memory

78d74bf

sarahyurick reviewed Sep 3, 2026

View reviewed changes

docs(semdedup): clarify format-specific fitting

f339b52

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

sarahyurick reviewed Sep 3, 2026

View reviewed changes

docs(semdedup): align fit defaults with input format

953e235

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

greptile-apps Bot reviewed Sep 3, 2026

View reviewed changes

sarahyurick approved these changes Sep 3, 2026

View reviewed changes

sarahyurick mentioned this pull request Sep 29, 2026

docs: draft 26.09 release and migration notes #2433

Merged



		@pytest.mark.gpu
		def test_parquet_file_info_uses_fsspec_for_remote_uri_without_storage_options() -> None:

Conversation

praateekmahajan commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Behavior and constraints

Benchmark results

20% dataset, full fit — 65.8M rows

Full dataset, 20% fit — 329.8M rows

Validation

Uh oh!

copy-pr-bot Bot commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

praateekmahajan commented Sep 2, 2026

Uh oh!

praateekmahajan commented Sep 2, 2026

Uh oh!

praateekmahajan commented Sep 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

praateekmahajan Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

praateekmahajan commented Sep 3, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

praateekmahajan commented Sep 2, 2026 •

edited

Loading

github-actions Bot commented Sep 2, 2026 •

edited

Loading

praateekmahajan Sep 3, 2026 •

edited

Loading