Skip to content

Add 7 nvidia/Nemotron-* calibration datasets to SUPPORTED_DATASET_CONFIG #1498

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cjluo-nv merged 2 commits into from
May 16, 2026
Merged

Add 7 nvidia/Nemotron-* calibration datasets to SUPPORTED_DATASET_CONFIG #1498

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bade82e
Add 7 nvidia/Nemotron-* calibration datasets to SUPPORTED_DATASET_CONFIG
Fridah-nv May 15, 2026
f9dd1bc
address reviews
Fridah-nv May 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 91 additions & 2 deletions modelopt/torch/utils/dataset_utils.py
View file
Open in desktop

@kevalmorabia97 kevalmorabia97 May 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Have you tried directly trying these datasets (e.g. get_dataset_dataloader("nvidia/Nemotron-SFT-Instruction-Following-Chat-v2", ...)) without adding to this dictionary? We have a fallback logic that auto-detects dataset type and parses it if not present in the dictionary

@Fridah-nv Fridah-nv May 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tried it on main for nvidia/Nemotron-SFT-Instruction-Following-Chat-v2 — the fallback hits three issues:

  1. get_dataset_samples(path, num_samples=2) → ValueError: Bad split: train. Available splits: ['reasoning_off', 'reasoning_on'].
    Auto-detect defaults to split=["train"] (dataset_utils.py:415); none of the 7 Nemotron datasets has a train split.
  2. With the correct split but no tokenizer → ValueError: Dataset '…' has a 'messages' column but no tokenizer with apply_chat_template was provided (dataset_utils.py:256-260).
  3. With a tokenizer but apply_chat_template=False → the fallback still calls apply_chat_template unconditionally, which fails on tokenizers without a chat template (e.g. gpt2).

We can think about more of improving the fallback pass.
As a background, we are evaluating if these dataset should be the default calibration dataset instead. So adding them in this PR

Original file line number Diff line number Diff line change
Expand Up @@ -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] = {
Expand Down Expand Up @@ -61,15 +66,99 @@
"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": {
"config": {
"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": {
# Skips ``reasoning_on`` split: heterogeneous messages schema fails streaming cast.
"config": {
"path": "nvidia/Nemotron-SFT-Instruction-Following-Chat-v2",
"split": ["reasoning_off"],
},
"preprocess": _join_messages_content,
"chat_key": "messages",
},
"nemotron-science-v1": {
"config": {
"path": "nvidia/Nemotron-Science-v1",
"split": ["MCQ", "RQA"],
},
"preprocess": _join_messages_content,
"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": _join_messages_content,
"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": _join_messages_content,
"chat_key": "messages",
},
"nemotron-math-v2": {
"config": {
"path": "nvidia/Nemotron-Math-v2",
"split": ["high_part00", "high_part01", "high_part02", "medium", "low"],
},
"preprocess": _join_messages_content,
"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": _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",
],
},
"preprocess": _join_messages_content,
"chat_key": "messages",
},
"magpie": {
Expand Down
55 changes: 55 additions & 0 deletions tests/unit/torch/utils/test_dataset_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

import pytest
import torch
from huggingface_hub import get_token

@coderabbitai coderabbitai Bot May 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Avoid module-level import of optional test dependency.

Importing huggingface_hub at file import time can fail test collection in environments missing that package. Move get_token import inside test_get_dataset_samples_new_nemotron() (or use pytest.importorskip("huggingface_hub")) so unrelated unit tests still collect/run.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/torch/utils/test_dataset_utils.py` at line 20, The module-level
import "from huggingface_hub import get_token" causes test collection failures
if huggingface_hub is not installed; move this import into the test function
test_get_dataset_samples_new_nemotron() (or replace it with
pytest.importorskip("huggingface_hub") at the top of that test) so the optional
dependency is only required when that test runs and unrelated tests can still be
collected.

from torch.utils.data import DataLoader

from modelopt.torch.utils.dataset_utils import (
Expand Down Expand Up @@ -689,3 +690,57 @@ 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


_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 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")
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)

assert isinstance(samples, list)
assert len(samples) == 2
assert all(isinstance(s, str) and len(s) > 0 for s in samples)
Loading