Repository navigation
Conversation
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
…modal_hybrid_real.yaml Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
 [Medium impact] The change reaches a new, operator-run multimodal curation workflow. Fix the parallel startup failure, preserve shortened text in exports, and bound recall memory use before merging. Findings
SummaryAdds an Omni-Fuse tutorial with SNS, expert embeddings, projection training, and query-based dataset selection.
Diagram%%{init: {'theme': 'neutral'}}%%
flowchart LR
A[Paired records] --> B[SNS]
B --> C[Expert embeddings]
C --> D[Projection training]
D --> E[Query ranking]
E --> F[Ranked and selected records]
Reviews (11) · Last reviewed commit: "Merge branch 'NVIDIA-NeMo:main' into mai..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| """Small compatibility layer around NeMo Curator APIs.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import copy | ||
| from typing import Any | ||
|
|
||
| import pandas as pd | ||
| from nemo_curator.pipeline import Pipeline as CuratorPipeline | ||
| from nemo_curator.tasks import DocumentBatch, EmptyTask | ||
|
|
||
|
|
||
| def records_from_task(task: Any) -> list[dict[str, Any]]: | ||
| """Return task data as records from a Curator task.""" | ||
|
|
||
| data = task.data | ||
| if hasattr(data, "to_dict"): | ||
| try: | ||
| return [dict(row) for row in data.to_dict(orient="records")] | ||
| except TypeError: | ||
| return [dict(row) for row in data.to_dict("records")] | ||
| if isinstance(data, list): | ||
| return [dict(row) for row in data] | ||
| raise TypeError(f"Unsupported task data type: {type(data)!r}") | ||
|
|
||
|
|
||
| def make_document_batch( | ||
| task_id: str, | ||
| dataset_name: str, | ||
| records: list[dict[str, Any]], | ||
| metadata: dict[str, Any] | None = None, | ||
| stage_perf: list[Any] | None = None, | ||
| ) -> DocumentBatch: | ||
| """Construct a NeMo Curator DocumentBatch.""" | ||
|
|
||
| return DocumentBatch( | ||
| task_id=task_id, | ||
| dataset_name=dataset_name, | ||
| data=pd.DataFrame.from_records(records), | ||
| _metadata=metadata or {}, | ||
| _stage_perf=stage_perf or [], | ||
| ) | ||
|
|
||
|
|
||
| def make_empty_task() -> EmptyTask: | ||
| if callable(EmptyTask): | ||
| try: | ||
| return EmptyTask(task_id="empty", dataset_name="omnifuse", data=None) | ||
| except TypeError: | ||
| return EmptyTask() | ||
| return copy.deepcopy(EmptyTask) | ||
|
|
||
|
|
||
| def make_curator_pipeline(name: str, stages: list[Any], description: str | None = None) -> CuratorPipeline: | ||
| return CuratorPipeline(name=name, description=description, stages=stages) |
There was a problem hiding this comment.
ProcessingStage, Resources, and re-export of EmptyTask
Every stage file (stages/eee.py, stages/sns.py, stages/projection.py, stages/datablend.py, stages/reader.py) imports ProcessingStage and Resources from this module. Neither is defined here or re-exported from it. stages/reader.py also imports EmptyTask directly. All five stage modules will fail immediately with ImportError when loaded, making the entire pipeline non-runnable.
These classes need to be either imported from nemo_curator and re-exported, or implemented as shim classes within this compat module.
|
|
||
| from omnifuse_tutorial.compat.curator import make_curator_pipeline, make_empty_task | ||
| from omnifuse_tutorial.config.models import ExperimentConfig | ||
| from omnifuse_tutorial.data.io import write_json |
There was a problem hiding this comment.
omnifuse_tutorial/data/ package
omnifuse_tutorial/data/io.py and omnifuse_tutorial/data/loader.py are imported by at least eight files — pipeline.py, utils.py, stages/eee.py, stages/projection.py, stages/datablend.py, stages/reader.py, sns/backends.py, and projection/trainer.py — but the entire omnifuse_tutorial/data/ directory does not exist in the repository. Every entry point to the tutorial will fail at import time with ModuleNotFoundError: No module named 'omnifuse_tutorial.data'.
There was a problem hiding this comment.
Added in recent commit. Was .gitignored previously.
| def _audio_format(path: Path) -> str: | ||
| audio_format = path.suffix.lower().lstrip(".") | ||
| if audio_format == "mpeg": | ||
| audio_format = "mp3" | ||
| if audio_format not in {"wav", "mp3"}: | ||
| raise ValueError(f"NVIDIA audio descriptions support wav/mp3 only, got {path.suffix} for {path}") | ||
| return audio_format |
There was a problem hiding this comment.
_audio_format raises for extensions declared as supported
AUDIO_EXTENSIONS and MIME_TYPES both include .flac, .m4a, .aac, and .ogg, advertising them as valid inputs. But _audio_format raises ValueError for any extension other than .wav and .mp3. The API description path in NvidiaApiEEEBackend._describe_raw → _describe_file → describe_file_with_nvidia_api calls this function unconditionally for audio records, so a user who provides a .flac or .m4a file will receive a confusing ValueError rather than a clear unsupported-format message.
Reconcile by either removing the unsupported extensions from AUDIO_EXTENSIONS/MIME_TYPES, converting to a supported format before sending, or raising a clear error with guidance early in the audio code path.
| "pandas>=2.0", | ||
| "pyyaml>=6.0", | ||
| "requests>=2.31", | ||
| "spacy>=3.8.14", |
There was a problem hiding this comment.
spacy constraint — the same package appears twice with conflicting lower bounds (>=3.8.14 on line 17 and >=3.8,<4 on line 35). Keep a single, unambiguous entry.
| "spacy>=3.8.14", | |
| "spacy>=3.8.14,<4", |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| "qwen-vl-utils>=0.0.8", | ||
| "sentence-transformers>=2.2,<4", | ||
| "soundfile>=0.12", | ||
| "spacy>=3.8,<4", |
There was a problem hiding this comment.
spacy entry.
| "spacy>=3.8,<4", | |
| # spacy consolidated above as spacy>=3.8.14,<4 |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| path = Path(value) | ||
| except OSError: | ||
| return None | ||
| return path if path.exists() else path |
There was a problem hiding this comment.
_path_or_none always returns a Path, never None — return path if path.exists() else path is a tautology; both branches return the same value. The | None return type annotation is misleading. Either return None when the path does not exist, or simplify to return path.
| return path if path.exists() else path | |
| return path |
|
|
||
| ```bash | ||
| uv sync --extra dev |
There was a problem hiding this comment.
The instruction says to run uv sync --extra dev "from the Curator repository root", but uv sync reads pyproject.toml from the current directory. No uv workspace configuration ties this tutorial to the repository root, so running uv sync from the repo root installs the main Curator package rather than this tutorial's dependencies. Users should cd tutorials/multimodal/omni-fuse-data-curation/ first.
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
| return DocumentBatch( | ||
| task_id=task_id, | ||
| dataset_name=dataset_name, | ||
| data=pd.DataFrame.from_records(records), | ||
| _metadata=metadata or {}, | ||
| _stage_perf=stage_perf or [], | ||
| ) |
There was a problem hiding this comment.
task_id is declared as field(init=False, default="") in the Task base class, so it is intentionally excluded from DocumentBatch.__init__. Passing task_id=task_id here raises TypeError: __init__() got an unexpected keyword argument 'task_id' at runtime. Because every stage calls make_document_batch, the entire pipeline fails before any data is processed. The task_id is managed by the framework (set via _set_task_id() at each stage boundary) and must not be supplied to the constructor.
| return DocumentBatch( | |
| task_id=task_id, | |
| dataset_name=dataset_name, | |
| data=pd.DataFrame.from_records(records), | |
| _metadata=metadata or {}, | |
| _stage_perf=stage_perf or [], | |
| ) | |
| return DocumentBatch( | |
| dataset_name=dataset_name, | |
| data=pd.DataFrame.from_records(records), | |
| _metadata=metadata or {}, | |
| _stage_perf=stage_perf or [], | |
| ) |
There was a problem hiding this comment.
You're right to revert. My original comment was incorrect — task_id is defined as field(init=False, default="") on the Task base class, which means Python's @dataclass machinery excludes it from __init__ entirely. Passing task_id=task_id will indeed raise TypeError: __init__() got an unexpected keyword argument 'task_id' at runtime.
The current code in the PR still has this bug. The fix from commit 21b8091 (which you reverted) was correct — task_id should be omitted from the DocumentBatch(...) constructor call and set afterwards via _set_task_id() if needed, or simply left to the framework to assign.
Here's the correct version:
| return DocumentBatch( | |
| task_id=task_id, | |
| dataset_name=dataset_name, | |
| data=pd.DataFrame.from_records(records), | |
| _metadata=metadata or {}, | |
| _stage_perf=stage_perf or [], | |
| ) | |
| return DocumentBatch( | |
| dataset_name=dataset_name, | |
| data=pd.DataFrame.from_records(records), | |
| _metadata=metadata or {}, | |
| _stage_perf=stage_perf or [], | |
| ) |
If the caller needs to set a specific task_id, it should be done after construction:
batch = make_document_batch(...)
batch._set_task_id("", task_id) # or however it fits the pipeline DAGBut in most cases, the framework handles task_id assignment automatically — make_document_batch doesn't need the task_id parameter at all.
Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
| task = run_curator_step( | ||
| name=f"{config.experiment_id}-2-embed", | ||
| stages=[EEEEmbeddingStage(config=config)], | ||
| initial_task=load_sns_task(config), | ||
| ) | ||
| records_path = config.run_dir / "embeddings" / "records.jsonl" | ||
| write_jsonl(records_path, records_from_task(task)) | ||
| metadata = dict(getattr(task, "_metadata", {}) or {}) | ||
| metadata["embedding_records_path"] = str(records_path) | ||
| task._metadata = metadata |
There was a problem hiding this comment.
task_id is not a parameter of make_document_batch
load_sns_task, load_embedding_task, and load_projection_task all call make_document_batch(task_id=..., dataset_name=..., records=..., metadata=...), but the current make_document_batch signature in compat/curator.py is (dataset_name, records, metadata=None, stage_perf=None) — it has no task_id parameter. Every one of these three loaders will raise TypeError: make_document_batch() got an unexpected keyword argument 'task_id' as soon as it is called, breaking all three resume paths (2_embed.py, 3_project.py, 4_datablend.py). Remove the task_id= keyword from all three call sites.
DocumentBatch requires task_id This reverts commit 21b8091. Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
| selected = ranker.select_top(ranked) | ||
|
|
||
| output_dir = self.config.run_dir / "datablend" | ||
| ranked_path = write_jsonl(output_dir / "datablend_ranked.jsonl", ranked) | ||
| topk_path = write_jsonl(output_dir / "datablend_topk.jsonl", selected) | ||
| metadata.update( | ||
| { | ||
| "datablend_ranked_path": str(ranked_path), | ||
| "datablend_topk_path": str(topk_path), | ||
| "datablend_size": len(selected), | ||
| } | ||
| ) |
There was a problem hiding this comment.
datablend_size metadata contradicts the batch record count
make_document_batch is called with records=ranked (all N ranked records), but datablend_size in metadata is set to len(selected) (the top-k count). Any downstream code that consumes the output task via records_from_task(task) receives all N records, not datablend_size records — the two values will silently disagree. If the intent is for the curated output batch to contain only the selected subset, records=selected should be passed instead. If the intent is to carry all rankings forward, datablend_size should reflect len(ranked) (or be renamed to datablend_topk_size) so the contract is unambiguous.
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
| def _describe_chat_completion_file( | ||
| path: Path, | ||
| model: str, | ||
| content_type: str, | ||
| prompt: str, | ||
| api_base_url: str, | ||
| headers: dict[str, str], | ||
| timeout: int, | ||
| ) -> str: | ||
| encoded = base64.b64encode(path.read_bytes()).decode("utf-8") | ||
| mime = MIME_TYPES.get(path.suffix.lower(), "application/octet-stream") | ||
| if content_type == "input_audio": | ||
| media_content = {"type": "input_audio", "input_audio": {"data": encoded, "format": _audio_format(path)}} | ||
| else: | ||
| media_content = {content_type: {"url": f"data:{mime};base64,{encoded}"}, "type": content_type} | ||
| url = f"{api_base_url}/chat/completions" | ||
| payload = { | ||
| "model": model, | ||
| "messages": [{"role": "user", "content": [{"type": "text", "text": prompt}, media_content]}], | ||
| "max_tokens": 512, | ||
| "temperature": 0.2, | ||
| "stream": False, | ||
| } | ||
| response = _post_nvidia_json_with_retries(url, headers, payload, timeout) | ||
| response = _resolve_nvidia_response(response, api_base_url, headers, timeout, model, url) | ||
| return _response_text(response.json(), model, url) |
There was a problem hiding this comment.
_describe_chat_completion_file reads the entire file with path.read_bytes() and embeds it as a base64 data-URL in the JSON body for both image_url and video_url content types. The audio code path has an explicit NVCF_ASSET_UPLOAD_THRESHOLD_BYTES = 180 * 1024 check and falls back to NVCF asset upload for large files. No equivalent guard exists here. A typical 10-second video file is easily 10–50 MB before encoding; base64 adds ~33 %, producing a JSON body the NVIDIA API will reject with HTTP 413 or a generic error — the same limit that drove the audio threshold. This silently breaks the video modality for any non-trivial input. Consider adding a size check and routing large image/video files through _upload_nvcf_asset (or raising an early ValueError with a clear message) just as audio already does.
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
Signed-off-by: Harshil Kotamreddy <hkotamreddy@nvidia.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
|
/review model=codex |
|
/review model=codex |
svcnvidia-nemo-ci
left a comment
svcnvidia-nemo-ci
left a comment
There was a problem hiding this comment.
What this change does: Adds a standalone tutorial that refines multimodal pairs, builds expert embeddings, trains a projection, and exports query-ranked records. It also adds a sharded SNS-to-EEE pipeline for overlapping work across GPUs.
Reviewed all 47 changed files and traced their Curator contracts and artifact consumers. Assumed one tutorial user, using either the numbered scripts or the documented two-GPU profile.
Must fix
Findings 1–4 below cover package installation, query embedding selection, exported text, and disabled-stage validation.
Should fix
Findings 5–6 cover regression tests and an unavailable documented configuration.
The current source addresses the earlier missing imports/data package, constructor compatibility, selected-record count, and large-image/video upload findings. Those are not repeated here.
Verdict: fix 1–4 first.
Not checked: dependency installation, tests, hosted endpoints, or GPU execution. All patches and acceptance checks below are proposed, not executed; verify them before adopting.
Note: formal approval is not available for this request, so this review is posted as a comment only.
Review generated by codex. This was a full review. Request a full review with /review reset=true.
| en-core-web-sm = { url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl" } | ||
|
|
||
| [tool.setuptools] | ||
| py-modules = [] |
There was a problem hiding this comment.
- Must fix: Include the tutorial package in the installation
- What this is: The build configuration explicitly sets
py-modules = []but never selects any packages. - Problem: An explicit module list disables setuptools' automatic discovery, so the built distribution contains no
omnifuse_tutorialpackage. Importing the installed tutorial outside its source directory therefore fails withModuleNotFoundError; running scripts from the checkout can hide this packaging failure. - Fix: Keep the empty top-level module list and add this package selection immediately afterward:
This must install every tutorial subpackage without packaging datasets, model checkouts, or numbered scripts. Proposed acceptance: inspect the built wheel for the package and its subpackages, then install it with dependencies and run
[tool.setuptools.packages.find] include = ["omnifuse_tutorial", "omnifuse_tutorial.*"]
import omnifuse_tutorialfrom a temporary directory with no checkout onPYTHONPATH; retain successful imports when running the documented scripts from the tutorial directory. - If we skip it: The distribution installs successfully but does not provide its implementation.
| backend: EEEBackend | ||
|
|
||
| def rank(self, records: list[dict[str, Any]], projection: ProjectionResult) -> list[dict[str, Any]]: | ||
| query_embedding = self.backend.embed_query(self.config.query, expert="text-based") |
There was a problem hiding this comment.
- Must fix: Embed the query with the projection's anchor expert
- What this is: Both projection trainers choose an
anchor_expert, falling back to the first configured expert whentext-basedis absent. - Problem: With a supported nonempty expert list such as
["fusion"], the torch projection targets LanguageBind annotation embeddings, but ranking always compares its output with a Nemotrontext-basedquery. Equal vector lengths conceal that these are different embedding spaces, producing unrelated similarity scores rather than ranking against the trained target space. - Fix: In
DatablendRanker.rank, replace the hardcoded selection with:Both trainers write this field, and the resume loader preserves it. Proposed acceptance: use a fake backend whose text-based query isanchor_expert = projection.model["anchor_expert"] query_embedding = self.backend.embed_query(self.config.query, expert=anchor_expert)
[1, 0]and fusion query is[0, 1]; withanchor_expert="fusion"and projected rows[[0, 1], [1, 0]], the first record must rank first, including after artifact reload; retain text-based query selection for the default expert list and usee2efor an e2e-only anchor. - If we skip it: Removing the text-based expert silently invalidates the final ranking.
| "pair_id": record["pair_id"], | ||
| "pool": record["pool"], | ||
| "modality": record["modality"], | ||
| "raw_path": record["raw_path"], |
There was a problem hiding this comment.
- Must fix: Preserve the refined raw text in the exported records
- What this is: Text forward extraction stores its retained sentences in
sns_raw_text, while leavingraw_pathpointed at the original file. - Problem: If the original file contains
A cat. A truck.and SNS retains onlyA cat., EEE embeds the shorter text, but this export drops both text fields and retains only the original file path. The resulting datablend cannot recover the refined sample that was ranked. - Fix: Before appending each row, preserve the effective text explicitly:
Document this exported field as the effective text value; retain
if record["modality"] == "text": row["raw_text"] = record.get("sns_raw_text") or record.get("raw_text")
raw_pathfor provenance and keep media rows pointing to their actual media files, not descriptions. Proposed acceptance: assert that the example exportsraw_text == "A cat."; disabled or rejected text extraction must export the original text, and image/audio/video rows must retain their existing paths and annotations; verify these content guarantees withinclude_metadataboth true and false. - If we skip it: Text curation is lost at the final export even though the embeddings used the curated content.
| "bidirectional", | ||
| } | ||
| if ( | ||
| effective_sns_backend in {"hybrid", "local"} |
There was a problem hiding this comment.
- Must fix: Do not require a video extraction checkpoint when SNS is disabled
- What this is: Input validation treats any video pool with a forward direction as requiring CG-DETR.
- Problem: Set
sns.enabled: falsein the supplied hybrid config and provide valid data, an API key, and LanguageBind, but no CG-DETR checkpoint: Step 0 still raises a missing-assets error. The disabled processor branch returns the original pair without ever using that checkpoint, so validation rejects a runnable configuration. - Fix: Add
config.sns.enabledto this checkpoint requirement:Proposed acceptance: with one valid video mapping and the other required assets available, disabled SNS plus a missing checkpoint must pass validation; enabled hybrid/local forward or bidirectional SNS must still reject the missing checkpoint, while backward-only and API SNS must remain exempt.if ( config.sns.enabled and effective_sns_backend in {"hybrid", "local"} and uses_video_forward and not config.sns.cg_detr_checkpoint.exists() ):
- If we skip it: Users cannot use the advertised disable switch to bypass local video extraction setup.
| return summary | ||
|
|
||
|
|
||
| def merge_parallel_outputs(config: ExperimentConfig, output_tasks: list[Any]) -> dict[str, Any]: |
There was a problem hiding this comment.
- Should fix: Add a CPU regression test for shard merging
- What this is: The new merge path sorts completed tasks, combines manifests, and reconstructs alternating raw/annotation arrays.
- Problem: This PR adds no test for those contracts. A change to sorting or alternating-row slicing can still produce valid-shaped files while assigning one record another record's embedding.
- Fix: Add a CPU-only test using fake SNS/EEE backends and temporary artifact directories; the invariant is that record order, manifest order, and each expert's raw/annotation rows stay aligned. Proposed acceptance: shard three records
a,b,cwith size two, return completed shards in reverse order, and assert canonical record/manifest IDs area,b,c; with raw vectors[i,0]and annotation vectors[0,i]fori=1,2,3, assert the merged raw and annotation arrays preserve those exact sequences; retain rejection of duplicate/missing shards and incorrect array dimensions, including the one-record final shard. - If we skip it: The most important parallel-output contract remains unprotected without running expensive models.
| ``` | ||
|
|
||
| For the five-dataset benchmark used in this tutorial, the matching two-GPU, | ||
| 4,000-epoch config is `configs/omni_fuse_1000_4000ep_parallel.yaml`. |
There was a problem hiding this comment.
- Should fix: Remove the unavailable benchmark configuration reference
- What this is: The benchmark instructions name a five-dataset, 4,000-epoch configuration.
- Problem: That file is absent from the head tree; the configuration directory contains only the hybrid tutorial YAML. Passing the named benchmark profile to a tutorial command therefore fails while loading the configuration.
- Fix: Delete these two benchmark-reference lines until that profile is included, retaining the existing parallel command and tuning examples. Proposed acceptance: inspect each configuration filename named by the README and confirm that it exists; the default sequential and parallel examples must still name the supplied hybrid configuration.
- If we skip it: Readers are directed to a benchmark setup they cannot reproduce from this checkout.
svcnvidia-nemo-ci
added
the
waiting-on-customer
| if self.worker_count is not None: | ||
| spec["num_workers"] = self.worker_count |
There was a problem hiding this comment.
The parallel path always sets worker_count, so SNSStage.xenna_stage_spec() returns num_workers. The current XennaExecutor.execute() rejects that key with ValueError before starting workers. EEEEmbeddingStage does the same thing.
Override num_workers() to return worker_count in both stages, and keep only slots_per_actor in xenna_stage_spec().
| "pair_id": record["pair_id"], | ||
| "pool": record["pool"], | ||
| "modality": record["modality"], | ||
| "raw_path": record["raw_path"], |
There was a problem hiding this comment.
For text records, forward SNS saves the shortened text only in sns_raw_text; raw_path still points to the original file. DatablendRanker.rank() exports that original path but drops sns_raw_text. Users loading the final datablend get the unfiltered text, even though the embeddings and ranking used the shortened text.
Include the refined text in the export, or write it to a new file and export that path.
| "raw_path": record["raw_path"], | |
| "raw_path": record["raw_path"], | |
| "sns_raw_text": record.get("sns_raw_text"), |
| similarities = annotations @ projected.T | ||
| targets = torch.arange(similarities.shape[0], device=similarities.device).unsqueeze(1) | ||
| annotation_to_raw = (similarities.topk(k, dim=1).indices == targets).any(dim=1).float().mean() | ||
| raw_to_annotation = (similarities.T.topk(k, dim=1).indices == targets).any(dim=1).float().mean() |
There was a problem hiding this comment.
_torch_recall_at_k() builds an all-record similarity matrix after training. Its memory use grows with the square of the record count, regardless of projection.batch_size. At 150,000 records, this float32 matrix alone needs about 90 GB. Projection therefore cannot finish on the documented 46 GiB GPUs even if training batches fit, and it fails before saving the model or projected embeddings.
Compute recall in chunks. The final model forward should also respect the configured batch size.
| For the five-dataset benchmark used in this tutorial, the matching two-GPU, | ||
| 4,000-epoch config is `configs/omni_fuse_1000_4000ep_parallel.yaml`. |
There was a problem hiding this comment.
The README points users to configs/omni_fuse_1000_4000ep_parallel.yaml, but that file is not included. Users cannot run the described benchmark with the supplied files. Add the config or remove the reference.
| For the five-dataset benchmark used in this tutorial, the matching two-GPU, | |
| 4,000-epoch config is `configs/omni_fuse_1000_4000ep_parallel.yaml`. |
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
Description
This PR adds a tutorial for Omni-Fuse. It follows a similar structure to existing tutorials and builds a curator pipeline to perform retrieval-based data curation for multimodal datasets.
Usage
Instructions to run the tutorial can be found in the README.md file. Open to feedback/adding additional instructions.
Checklist