Add DATASET_COMBOS for grouped calibration datasets

CHANGELOG.rst

@@ -25,6 +25,8 @@ Changelog
 - Add ``--cast_mxfp4_to_nvfp4`` flag to ``examples/llm_ptq/hf_ptq.py`` for closed-form, bit-exact MXFP4 → NVFP4 weight conversion. Supports the GPT-OSS family (``openai/gpt-oss-20b``, ``openai/gpt-oss-120b``). See `examples/llm_ptq/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_ptq#mxfp4--nvfp4-cast-for-gpt-oss>`__ for usage.
 - DeepSeek PTQ (``examples/deepseek/ptq.py``) now defaults to native top-k calibration with post-hoc per-layer peer-max sync of expert ``input_quantizer.amax``; the all-experts path is preserved behind ``--calib_all_experts``.
 - Add NVFP4 W4A16 weight-only quantization (``w4a16_nvfp4``): FP4 weights with group_size=16, BF16 activations, no calibration forward pass required. Use ``mtq.W4A16_NVFP4_CFG`` or ``--qformat w4a16_nvfp4`` in ``hf_ptq.py``. vLLM deployment support is in progress.
+- Add ``DATASET_COMBOS`` to ``modelopt.torch.utils.dataset_utils`` — single ``--dataset`` tokens that fan out to multiple registered datasets; per-entry ``num_samples`` is split evenly across the members. Initial combos: ``cnn_nemotron_v2_mix`` (``cnn_dailymail`` + ``nemotron-post-training-dataset-v2``, used by ``hf_ptq.py`` when no ``--dataset`` is provided) and ``nemotron-post-training-v3`` (the seven ``nvidia/Nemotron-*`` SFT datasets added in #1498, mirroring the `nemotron-post-training-v3 collection <https://huggingface.co/collections/nvidia/nemotron-post-training-v3>`_). Combo names are listed by ``get_supported_datasets()`` and surfaced in ``--dataset`` help. ``get_dataset_dataloader`` rejects inputs that mix a combo with one of its member datasets (e.g. ``cnn_dailymail,cnn_nemotron_v2_mix``) to avoid double-sampling, and ``get_dataset_samples`` rejects combo names so callers route through the dataloader. ``hf_ptq.py`` default ``--calib_size`` is bumped from ``512`` to ``1024`` so the total calibration sample count under the new default combo matches the previous two-dataset fallback.
+- The ``nemotron-sft-agentic-v2`` registered dataset (added in #1498) now uses only the ``search`` split. The previously configured ``interactive_agent`` and ``tool_calling`` splits contain content-level defects (heterogeneous schema and a malformed JSON row, respectively) that cause pyarrow's streaming JSON reader to fail deterministically.
 
 0.44 (2026-05-18)
 ^^^^^^^^^^^^^^^^^

examples/llm_ptq/hf_ptq.py

@@ -531,9 +531,10 @@ def load_model(args: argparse.Namespace):
             language_model = full_model
         else:
             if args.dataset is None:
-                args.dataset = ["cnn_dailymail", "nemotron-post-training-dataset-v2"]
+                args.dataset = ["cnn_nemotron_v2_mix"]
                 warnings.warn(
-                    "No dataset specified. Defaulting to cnn_dailymail and nemotron-post-training-dataset-v2."
+                    "No dataset specified. Defaulting to the 'cnn_nemotron_v2_mix' combo "
+                    "(cnn_dailymail + nemotron-post-training-dataset-v2)."
                 )
             # Adjust calib_size to match dataset length by extending or truncating as needed
             args.calib_size = (args.calib_size + [args.calib_size[-1]] * len(args.dataset))[
@@ -1238,7 +1239,7 @@ def parse_args() -> argparse.Namespace:
             "This argument will be parsed and converted as a list of ints."
         ),
         type=str,
-        default="512",
+        default="1024",
     )
     parser.add_argument(
         "--calib_seq",

modelopt/torch/utils/dataset_utils.py

@@ -109,10 +109,12 @@ def _join_messages_content(sample: dict) -> str:
         "chat_key": "messages",
     },
     "nemotron-sft-agentic-v2": {
-        # Skips ``search`` split: heterogeneous messages schema fails streaming cast.
+        # Only ``search`` streams cleanly: ``interactive_agent`` has a heterogeneous
+        # tools schema (string vs list) that breaks pyarrow JSON inference, and
+        # ``tool_calling`` contains at least one malformed JSON row in a later shard.
         "config": {
             "path": "nvidia/Nemotron-SFT-Agentic-v2",
-            "split": ["interactive_agent", "tool_calling"],
+            "split": ["search"],
         },
         "preprocess": _join_messages_content,
         "chat_key": "messages",
@@ -195,6 +197,42 @@ def _join_messages_content(sample: dict) -> str:
     },
 }
 
+# Named groups of registered datasets, expanded in ``get_dataset_dataloader``.
+# Useful when callers want a single ``--dataset`` token that fans out to several
+# entries; per-dataset ``num_samples`` is split evenly across the members.
+DATASET_COMBOS: dict[str, list[str]] = {
+    "cnn_nemotron_v2_mix": ["cnn_dailymail", "nemotron-post-training-dataset-v2"],
+    "nemotron-post-training-v3": [
+        "nemotron-sft-instruction-following-chat-v2",
+        "nemotron-science-v1",
+        "nemotron-competitive-programming-v1",
+        "nemotron-sft-agentic-v2",
+        "nemotron-math-v2",
+        "nemotron-sft-swe-v2",
+        "nemotron-sft-multilingual-v1",
+    ],
+}
+
+
+def _validate_dataset_combos() -> None:
+    """Validate DATASET_COMBOS at import time: fail loud on typos / collisions."""
+    overlap = set(DATASET_COMBOS) & set(SUPPORTED_DATASET_CONFIG)
+    if overlap:
+        raise ValueError(
+            f"DATASET_COMBOS keys collide with SUPPORTED_DATASET_CONFIG: {sorted(overlap)}"
+        )
+    for combo_name, members in DATASET_COMBOS.items():
+        if not members:
+            raise ValueError(f"DATASET_COMBOS['{combo_name}'] must contain at least one dataset.")
+        unknown = [m for m in members if m not in SUPPORTED_DATASET_CONFIG]
+        if unknown:
+            raise ValueError(
+                f"DATASET_COMBOS['{combo_name}'] references unknown datasets: {unknown}"
+            )
+
+
+_validate_dataset_combos()
+
 __all__ = [
     "create_forward_loop",
     "download_hf_dataset_as_jsonl",
@@ -349,6 +387,13 @@ def get_dataset_samples(
     Returns:
         Samples: The list of samples.
     """
+    if dataset_name in DATASET_COMBOS:
+        raise ValueError(
+            f"'{dataset_name}' is a DATASET_COMBOS entry, not a single dataset. "
+            "Use ``get_dataset_dataloader`` to expand combos, or pass one of "
+            f"its members: {DATASET_COMBOS[dataset_name]}"
+        )
+
     # Local JSONL: load via HF's ``json`` builder and route through the same
     # auto-preprocess path as unregistered HF datasets so chat / prompt / text
     # columns are handled consistently with a downloaded HF dataset.  Never
@@ -560,6 +605,34 @@ def get_dataset_dataloader(
         "dataset_name and num_samples must be the same length"
     )
 
+    # Reject inputs that include both a combo and one of its member datasets
+    # (e.g. ``["cnn_dailymail", "cnn_nemotron_v2_mix"]``), since the combo would sample the
+    # plain entry a second time with a smaller per-member quota.
+    plain_inputs = {n for n in dataset_name if n not in DATASET_COMBOS}
+    for ds_name in dataset_name:
+        if ds_name in DATASET_COMBOS:
+            overlap = plain_inputs & set(DATASET_COMBOS[ds_name])
+            if overlap:
+                raise ValueError(
+                    f"--dataset includes both combo '{ds_name}' and its "
+                    f"member(s) {sorted(overlap)}; remove one to avoid "
+                    "double-sampling."
+                )
+
+    expanded_names: list[str] = []
+    expanded_num_samples: list[int] = []
+    for ds_name, n in zip(dataset_name, num_samples):
+        if ds_name in DATASET_COMBOS:
+            members = DATASET_COMBOS[ds_name]
+            base, remainder = divmod(n, len(members))
+            for i, member in enumerate(members):
+                expanded_names.append(member)
+                expanded_num_samples.append(base + (1 if i < remainder else 0))
+        else:
+            expanded_names.append(ds_name)
+            expanded_num_samples.append(n)
+    dataset_name, num_samples = expanded_names, expanded_num_samples
+
     all_samples = []
     for ds_name, num_sample in zip(dataset_name, num_samples):
         samples = get_dataset_samples(
@@ -617,7 +690,7 @@ def get_supported_datasets() -> list[str]:
 
             print("Supported datasets:", get_supported_datasets())
     """
-    return list(SUPPORTED_DATASET_CONFIG.keys())
+    return list(SUPPORTED_DATASET_CONFIG.keys()) + list(DATASET_COMBOS.keys())
 
 
 @contextmanager

tests/unit/torch/utils/test_dataset_utils.py

@@ -610,6 +610,89 @@ def test_length_mismatch_raises(self, tmp_path, pad_tokenizer):
             )
 
 
+class TestDatasetCombosExpansion:
+    """Combo names in ``--dataset`` fan out to their registered members.
+
+    The combo branch in ``get_dataset_dataloader`` is exercised by mocking
+    ``get_dataset_samples`` so we can assert the post-expansion (name, count)
+    sequence without hitting HF.
+    """
+
+    @staticmethod
+    def _record_calls(monkeypatch):
+        calls: list[tuple[str, int]] = []
+
+        def _fake(name, num_sample, **_kwargs):
+            calls.append((name, num_sample))
+            return [f"{name}-{i}" for i in range(num_sample)]
+
+        from modelopt.torch.utils import dataset_utils
+
+        monkeypatch.setattr(dataset_utils, "get_dataset_samples", _fake)
+        return calls
+
+    def test_combo_expands_evenly(self, monkeypatch, pad_tokenizer):
+        from modelopt.torch.utils.dataset_utils import DATASET_COMBOS
+
+        calls = self._record_calls(monkeypatch)
+        get_dataset_dataloader(
+            dataset_name="cnn_nemotron_v2_mix",
+            tokenizer=pad_tokenizer,
+            num_samples=8,
+            batch_size=1,
+            max_sample_length=16,
+        )
+        members = DATASET_COMBOS["cnn_nemotron_v2_mix"]
+        assert calls == [(members[0], 4), (members[1], 4)]
+
+    def test_combo_remainder_distributed_to_earlier_members(self, monkeypatch, pad_tokenizer):
+        from modelopt.torch.utils.dataset_utils import DATASET_COMBOS
+
+        calls = self._record_calls(monkeypatch)
+        get_dataset_dataloader(
+            dataset_name="nemotron-post-training-v3",
+            tokenizer=pad_tokenizer,
+            num_samples=10,
+            batch_size=1,
+            max_sample_length=16,
+        )
+        members = DATASET_COMBOS["nemotron-post-training-v3"]
+        # 10 / 7 = 1 base, remainder 3 -> first 3 get +1
+        expected_counts = [2, 2, 2, 1, 1, 1, 1]
+        assert calls == list(zip(members, expected_counts))
+
+    def test_plain_and_combo_compose(self, monkeypatch, pad_tokenizer):
+        from modelopt.torch.utils.dataset_utils import DATASET_COMBOS
+
+        calls = self._record_calls(monkeypatch)
+        get_dataset_dataloader(
+            dataset_name=["cnn_dailymail", "nemotron-post-training-v3"],
+            tokenizer=pad_tokenizer,
+            num_samples=[3, 7],
+            batch_size=1,
+            max_sample_length=16,
+        )
+        members = DATASET_COMBOS["nemotron-post-training-v3"]
+        assert calls == [("cnn_dailymail", 3)] + [(m, 1) for m in members]
+
+    def test_combo_overlapping_with_member_raises(self, monkeypatch, pad_tokenizer):
+        self._record_calls(monkeypatch)
+        with pytest.raises(ValueError, match="combo 'cnn_nemotron_v2_mix'"):
+            get_dataset_dataloader(
+                dataset_name=["cnn_dailymail", "cnn_nemotron_v2_mix"],
+                tokenizer=pad_tokenizer,
+                num_samples=[2, 4],
+                batch_size=1,
+                max_sample_length=16,
+            )
+
+    def test_get_dataset_samples_rejects_combo_name(self):
+        from modelopt.torch.utils.dataset_utils import get_dataset_samples
+
+        with pytest.raises(ValueError, match="DATASET_COMBOS"):
+            get_dataset_samples("cnn_nemotron_v2_mix", num_samples=1)
+
+
 # ---------------------------------------------------------------------------
 # Live HF dataset round-trips. ``hf-internal-testing/dataset_with_data_files``
 # is a 10-row x {train,test} fixture maintained by HF for their own CI — tiny