mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Add DATASET_COMBOS for grouped calibration datasets (#1508)
## Summary - Add ``DATASET_COMBOS`` to ``modelopt.torch.utils.dataset_utils`` — single ``--dataset`` tokens that fan out to several entries in ``SUPPORTED_DATASET_CONFIG``. The per-entry ``num_samples`` is split evenly across the members inside ``get_dataset_dataloader``. - Two initial combos: - ``cnn_nemotron_v2_mix`` → ``cnn_dailymail`` + ``nemotron-post-training-dataset-v2``. Replaces the hardcoded two-element fallback list in ``hf_ptq.py`` when ``--dataset`` is omitted. - ``nemotron-post-training-v3`` → the seven ``nvidia/Nemotron-*`` SFT datasets registered in #1498 (mirroring the upstream [`nemotron-post-training-v3` collection](https://huggingface.co/collections/nvidia/nemotron-post-training-v3)). - ``get_supported_datasets()`` now appends combo names so they show up in ``--dataset`` help. - ``hf_ptq.py``'s default ``--calib_size`` bumped from ``512`` to ``1024`` so the ``cnn_nemotron_v2_mix`` combo's even split preserves the previous total sample count (was 512 per-dataset × 2 datasets = 1024; now 1024 split → 512 per-dataset × 2). ``--calib_size`` now denotes the total calibration budget regardless of combo cardinality. - Reject mixing a combo with one of its member datasets in the same ``--dataset`` list (e.g. ``cnn_dailymail,cnn_nemotron_v2_mix``) — combo would otherwise double-sample the explicit member with a smaller per-member quota. - Reject combo names in ``get_dataset_samples``; combos are dataloader-only. The error message points callers to ``get_dataset_dataloader``. - Validate ``DATASET_COMBOS`` at import time: empty member lists, name collisions with ``SUPPORTED_DATASET_CONFIG``, and references to unknown datasets raise ``ValueError`` up front. ## Test plan End-to-end validated against ``/hf-local/Qwen/Qwen3.5-0.8B`` via ``get_dataset_dataloader`` on the actual streamed data, plus 5 new unit tests in ``TestDatasetCombosExpansion`` (all 44 tests in ``test_dataset_utils.py`` pass with no regressions). - [x] ``python -c "from modelopt.torch.utils.dataset_utils import DATASET_COMBOS, get_supported_datasets; assert 'cnn_nemotron_v2_mix' in get_supported_datasets() and 'nemotron-post-training-v3' in get_supported_datasets()"`` - [x] ``hf_ptq.py`` with no ``--dataset`` flag still calibrates on cnn_dailymail + nemotron-post-training-dataset-v2 with the same total sample count as before. - [x] ``--dataset nemotron-post-training-v3 --calib_size 1024`` allocates 146 per member across the seven Nemotron datasets; full 1022-sample dataloader builds without error. - [x] ``--dataset cnn_dailymail,nemotron-post-training-v3 --calib_size 256,1024`` composes correctly: 256 from cnn_dailymail (as a plain entry) plus the 7-way split from the combo. (The earlier ``cnn_dailymail,cnn_nemotron_v2_mix`` example is rejected by design since ``cnn_dailymail`` is a member of that combo.) - [x] ``--dataset cnn_dailymail,cnn_nemotron_v2_mix`` raises ``ValueError`` with a clear message. - [x] ``get_dataset_samples("cnn_nemotron_v2_mix", ...)`` raises ``ValueError`` pointing to ``get_dataset_dataloader``. - [x] Unit tests: ``pytest tests/unit/torch/utils/test_dataset_utils.py`` — 44 passed. ## Post-validation fix End-to-end testing surfaced that the original ``nemotron-sft-agentic-v2`` entry kept the two splits (``interactive_agent``, ``tool_calling``) that pyarrow's streaming JSON reader cannot parse, and excluded ``search`` which is the only clean split. Failures reproduce deterministically across cache wipes with ``force_redownload``, so they are content-level defects in the published JSONL files at the pinned revision, not local artifacts: - ``interactive_agent`` — heterogeneous schema (``Column(.../member_id/type) changed from string to array``) at JSONL row 4. - ``tool_calling`` — malformed JSON row in a later shard, fails at sample ~885 with ``Missing a closing quotation mark in string``. - ``search`` — streams cleanly (verified to 2500 samples). Commit ``10f3cfd`` corrects ``nemotron-sft-agentic-v2`` to use only ``search``, with an updated comment. The CHANGELOG calls this out as a separate bullet so the behavior change on a previously-released dataset entry from #1498 is discoverable. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Dataset combo support: a single dataset token can expand into multiple registered datasets with even sample splitting; predefined combos added (e.g., cnn_nemotron_v2_mix, nemotron-post-training-v3) and listed as supported. * **Updates** * Default dataset when none specified now uses cnn_nemotron_v2_mix. * Calibration size default increased from 512 to 1024. * **Bug Fixes** * nemotron-sft-agentic-v2 now uses only the deterministic "search" split to avoid streaming JSON errors. * **Tests** * Added coverage for combo expansion, splitting, overlap validation, and rejection behavior. <!-- review_stack_entry_start --> [](https://app.coderabbit.ai/change-stack/NVIDIA/Model-Optimizer/pull/1508?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack) <!-- review_stack_entry_end --> <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chenjie Luo
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
d7733562dd
commit
a5bc6f8123
@@ -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.
|
||||
|
||||
**Bug Fixes**
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -109,10 +109,12 @@ SUPPORTED_DATASET_CONFIG: dict[str, Any] = {
|
||||
"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 @@ SUPPORTED_DATASET_CONFIG: dict[str, Any] = {
|
||||
},
|
||||
}
|
||||
|
||||
# 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
|
||||
|
||||
@@ -610,6 +610,89 @@ class TestGetDatasetDataloaderBlending:
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user