From bade82eccad8eff20d2b1997c3755e97f079437f Mon Sep 17 00:00:00 2001 From: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 14 May 2026 23:06:39 -0700 Subject: [PATCH 1/2] Add 7 nvidia/Nemotron-* calibration datasets to SUPPORTED_DATASET_CONFIG Registers nemotron-{sft-instruction-following-chat-v2, science-v1, competitive-programming-v1, sft-agentic-v2, math-v2, sft-swe-v2, sft-multilingual-v1} so hf_ptq.py's --dataset flag (which enumerates get_supported_datasets() automatically) can select these for PTQ calibration. Splits with heterogeneous parquet schemas that crash streaming CastError mid-iteration are excluded per inline comments. Adds a parametrized smoke test that skips when HF_TOKEN is unset since the nvidia/Nemotron-* datasets are gated. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/utils/dataset_utils.py | 69 ++++++++++++++++++++ tests/unit/torch/utils/test_dataset_utils.py | 31 +++++++++ 2 files changed, 100 insertions(+) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 129052e9ae0..f50e677eae6 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -72,6 +72,75 @@ "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), "chat_key": "messages", }, + "nemotron-sft-instruction-following-chat-v2": { + # Skips ``reasoning_on`` split: heterogeneous messages schema fails streaming cast. + "config": { + "path": "nvidia/Nemotron-SFT-Instruction-Following-Chat-v2", + "split": ["reasoning_off"], + }, + "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "chat_key": "messages", + }, + "nemotron-science-v1": { + "config": { + "path": "nvidia/Nemotron-Science-v1", + "split": ["MCQ", "RQA"], + }, + "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "chat_key": "messages", + }, + "nemotron-competitive-programming-v1": { + # Skips ``infinibyte_part0[0|1]``: heterogeneous schema fails streaming cast. + "config": { + "path": "nvidia/Nemotron-Competitive-Programming-v1", + "split": [ + "competitive_coding_cpp_part00", + "competitive_coding_cpp_part01", + "competitive_coding_python_part00", + "competitive_coding_python_part01", + ], + }, + "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "chat_key": "messages", + }, + "nemotron-sft-agentic-v2": { + # Skips ``search`` split: heterogeneous messages schema fails streaming cast. + "config": { + "path": "nvidia/Nemotron-SFT-Agentic-v2", + "split": ["interactive_agent", "tool_calling"], + }, + "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "chat_key": "messages", + }, + "nemotron-math-v2": { + "config": { + "path": "nvidia/Nemotron-Math-v2", + "split": ["high_part00", "high_part01", "high_part02", "medium", "low"], + }, + "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "chat_key": "messages", + }, + "nemotron-sft-swe-v2": { + # Skips ``openhands_swe`` split: heterogeneous schema fails streaming cast. + "config": { + "path": "nvidia/Nemotron-SFT-SWE-v2", + "split": ["agentless"], + }, + "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "chat_key": "messages", + }, + "nemotron-sft-multilingual-v1": { + "config": { + "path": "nvidia/Nemotron-SFT-Multilingual-v1", + "split": [ + "code_de", "code_es", "code_fr", "code_it", "code_ja", "code_zh", + "math_de", "math_es", "math_fr", "math_it", "math_ja", "math_zh", + "stem_de", "stem_es", "stem_fr", "stem_it", "stem_ja", "stem_zh", + ], + }, + "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "chat_key": "messages", + }, "magpie": { "config": { "path": "Magpie-Align/Magpie-Pro-MT-300K-v0.1", diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index f89663d89b5..83a2da190ea 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os from unittest.mock import Mock, patch import pytest @@ -689,3 +690,33 @@ def test_dataloader_mixing_hf_and_local_jsonl(self, tmp_path, pad_tokenizer): ) batches = list(loader) assert sum(b["input_ids"].shape[0] for b in batches) == 5 + + +@pytest.mark.parametrize( + "dataset_key", + [ + "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 test_get_dataset_samples_new_nemotron(dataset_key): + """Smoke-test the 7 newly registered nvidia/Nemotron-* calibration datasets. + + Skipped when ``HF_TOKEN`` is unset because these datasets live behind the HF Hub + and the CI runner may not have credentials. + """ + pytest.importorskip("datasets") + pytest.importorskip("huggingface_hub") + if not os.environ.get("HF_TOKEN"): + pytest.skip("HF_TOKEN not set; skipping nvidia/Nemotron-* dataset smoke test") + + samples = get_dataset_samples(dataset_key, num_samples=2) + + assert isinstance(samples, list) + assert len(samples) == 2 + assert all(isinstance(s, str) and len(s) > 0 for s in samples) From f9dd1bc5d4ed0218dd274c943e5a2ea3daa17599 Mon Sep 17 00:00:00 2001 From: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> Date: Fri, 15 May 2026 22:09:08 +0000 Subject: [PATCH 2/2] address reviews Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> --- modelopt/torch/utils/dataset_utils.py | 44 ++++++++++---- tests/unit/torch/utils/test_dataset_utils.py | 60 ++++++++++++++------ 2 files changed, 74 insertions(+), 30 deletions(-) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index f50e677eae6..4329f15488f 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -32,6 +32,11 @@ if TYPE_CHECKING: from transformers import PreTrainedTokenizerBase + +def _join_messages_content(sample: dict) -> str: + return "\n".join(turn["content"] for turn in sample["messages"]) + + # Use dict to store the config for each dataset. # If we want to export more options to user like target languages, we need more standardized approach like dataclass. SUPPORTED_DATASET_CONFIG: dict[str, Any] = { @@ -61,7 +66,7 @@ "path": "nvidia/Nemotron-Post-Training-Dataset-v2", "split": ["stem", "chat", "math", "code"], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "nemotron-post-training-dataset-v1": { @@ -69,7 +74,7 @@ "path": "nvidia/Nemotron-Post-Training-Dataset-v1", "split": ["stem", "chat", "math", "code", "tool_calling"], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "nemotron-sft-instruction-following-chat-v2": { @@ -78,7 +83,7 @@ "path": "nvidia/Nemotron-SFT-Instruction-Following-Chat-v2", "split": ["reasoning_off"], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "nemotron-science-v1": { @@ -86,7 +91,7 @@ "path": "nvidia/Nemotron-Science-v1", "split": ["MCQ", "RQA"], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "nemotron-competitive-programming-v1": { @@ -100,7 +105,7 @@ "competitive_coding_python_part01", ], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "nemotron-sft-agentic-v2": { @@ -109,7 +114,7 @@ "path": "nvidia/Nemotron-SFT-Agentic-v2", "split": ["interactive_agent", "tool_calling"], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "nemotron-math-v2": { @@ -117,7 +122,7 @@ "path": "nvidia/Nemotron-Math-v2", "split": ["high_part00", "high_part01", "high_part02", "medium", "low"], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "nemotron-sft-swe-v2": { @@ -126,19 +131,34 @@ "path": "nvidia/Nemotron-SFT-SWE-v2", "split": ["agentless"], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "nemotron-sft-multilingual-v1": { "config": { "path": "nvidia/Nemotron-SFT-Multilingual-v1", "split": [ - "code_de", "code_es", "code_fr", "code_it", "code_ja", "code_zh", - "math_de", "math_es", "math_fr", "math_it", "math_ja", "math_zh", - "stem_de", "stem_es", "stem_fr", "stem_it", "stem_ja", "stem_zh", + "code_de", + "code_es", + "code_fr", + "code_it", + "code_ja", + "code_zh", + "math_de", + "math_es", + "math_fr", + "math_it", + "math_ja", + "math_zh", + "stem_de", + "stem_es", + "stem_fr", + "stem_it", + "stem_ja", + "stem_zh", ], }, - "preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["messages"]), + "preprocess": _join_messages_content, "chat_key": "messages", }, "magpie": { diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 83a2da190ea..7566088a1ad 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -13,11 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os from unittest.mock import Mock, patch import pytest import torch +from huggingface_hub import get_token from torch.utils.data import DataLoader from modelopt.torch.utils.dataset_utils import ( @@ -692,28 +692,52 @@ def test_dataloader_mixing_hf_and_local_jsonl(self, tmp_path, pad_tokenizer): assert sum(b["input_ids"].shape[0] for b in batches) == 5 -@pytest.mark.parametrize( - "dataset_key", - [ - "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", - ], -) +_NEW_NEMOTRON_KEYS = [ + "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", +] + + +@pytest.mark.parametrize("dataset_key", _NEW_NEMOTRON_KEYS) +def test_new_nemotron_registry_shape(dataset_key): + """Always-on shape check on the 7 newly registered nvidia/Nemotron-* entries. + + Complements the gated smoke test below — catches typos in dataset paths or + split names even when the runner has no HF credentials. + """ + from modelopt.torch.utils.dataset_utils import SUPPORTED_DATASET_CONFIG + + assert dataset_key in SUPPORTED_DATASET_CONFIG + entry = SUPPORTED_DATASET_CONFIG[dataset_key] + config = entry["config"] + assert config["path"].startswith("nvidia/Nemotron-") + splits = config["split"] + assert isinstance(splits, list) and splits + assert all(isinstance(s, str) and s for s in splits) + assert len(set(splits)) == len(splits) + assert callable(entry["preprocess"]) + assert entry["chat_key"] == "messages" + + +@pytest.mark.integration +@pytest.mark.parametrize("dataset_key", _NEW_NEMOTRON_KEYS) def test_get_dataset_samples_new_nemotron(dataset_key): """Smoke-test the 7 newly registered nvidia/Nemotron-* calibration datasets. - Skipped when ``HF_TOKEN`` is unset because these datasets live behind the HF Hub - and the CI runner may not have credentials. + Skipped when no HF token is available because these datasets live behind the HF Hub. + ``huggingface_hub.get_token()`` covers both the ``HF_TOKEN`` env var and tokens + cached by ``hf auth login``. """ pytest.importorskip("datasets") - pytest.importorskip("huggingface_hub") - if not os.environ.get("HF_TOKEN"): - pytest.skip("HF_TOKEN not set; skipping nvidia/Nemotron-* dataset smoke test") + if not get_token(): + pytest.skip( + "No HF token (env HF_TOKEN or `hf auth login`); skipping gated Nemotron smoke test" + ) samples = get_dataset_samples(dataset_key, num_samples=2)