diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7dd63c29ecd..c0c18e797f5 100755 --- a/CHANGELOG.rst +++ b/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 `__ 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 `_). 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) ^^^^^^^^^^^^^^^^^ diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index d52c3ee40bb..0c17558d2e3 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/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", diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 4329f15488f..80ed8f9abdd 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/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 diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 7566088a1ad..812d2cd9c3b 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/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