fix(data): warn when behavioral dataset config keys are silently ignored

nemo_rl/data/datasets/response_datasets/__init__.py

@@ -53,7 +53,10 @@
 from nemo_rl.data.datasets.response_datasets.response_dataset import ResponseDataset
 from nemo_rl.data.datasets.response_datasets.squad import SquadDataset
 from nemo_rl.data.datasets.response_datasets.tulu3 import Tulu3SftMixtureDataset
-from nemo_rl.data.datasets.utils import resolve_external_dataset_class
+from nemo_rl.data.datasets.utils import (
+    resolve_external_dataset_class,
+    warn_on_unsupported_dataset_config_keys,
+)
 
 DATASET_REGISTRY = {
     # built-in datasets
@@ -119,6 +122,10 @@ def load_response_dataset(data_config: ResponseDatasetConfig):
             "(ensure it is installed and importable from PYTHONPATH)."
         )
 
+    # Every dataset class accepts **kwargs, so unsupported config keys are
+    # otherwise swallowed silently (e.g. `subset` on GSM8KDataset).
+    warn_on_unsupported_dataset_config_keys(dataset_class, data_config)
+
     dataset = dataset_class(
         **data_config  # pyrefly: ignore[missing-argument]  `data_path` is required for some classes
     )

nemo_rl/data/datasets/utils.py

@@ -13,11 +13,14 @@
 # limitations under the License.
 
 import base64
+import functools
 import importlib
+import inspect
 import io
 import os
+import warnings
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Mapping, Optional, Union
 
 import numpy as np
 import torch
@@ -182,6 +185,72 @@ def resolve_external_dataset_class(dataset_name: str) -> type:
     return dataset_class
 
 
+# Behavioral keys from ResponseDatasetConfig that change *which data* a run
+# sees. If one of these is set but the resolved dataset class does not accept
+# it, the user silently gets different data than they configured.
+_BEHAVIORAL_DATASET_CONFIG_KEYS = ("seed", "split", "split_validation_size", "subset")
+
+
+def warn_on_unsupported_dataset_config_keys(
+    dataset_class: Any, data_config: Mapping[str, Any]
+) -> None:
+    """Warn when behavioral dataset config keys would be silently ignored.
+
+    The dataset dispatchers instantiate with ``dataset_class(**data_config)``
+    and every built-in dataset accepts ``**kwargs``, so a config key the class
+    does not actually support is swallowed without any feedback — e.g.
+    ``subset`` on ``GSM8KDataset`` (which hardcodes the ``main`` config) or
+    ``split_validation_size`` on datasets that never call
+    ``split_train_validation``. For the keys in
+    ``_BEHAVIORAL_DATASET_CONFIG_KEYS``, warn when the resolved class does not
+    declare the parameter anywhere in its ``__init__`` MRO.
+
+    The check walks the MRO because some datasets consume keys via
+    ``**kwargs`` and forward them to a base-class ``__init__`` that declares
+    them (e.g. the intent datasets). ``functools.partial`` registry entries
+    (the AIME variants) are unwrapped first. ``None`` values are skipped
+    (``subset: null`` is the documented default), as is a falsy
+    ``split_validation_size`` (0 means "no validation split" everywhere).
+    """
+    while isinstance(dataset_class, functools.partial):
+        dataset_class = dataset_class.func
+    if not isinstance(dataset_class, type):
+        return
+
+    accepted: set[str] = set()
+    for klass in dataset_class.__mro__:
+        init = klass.__dict__.get("__init__")
+        if init is None:
+            continue
+        try:
+            signature = inspect.signature(init)
+        except (TypeError, ValueError):
+            # Signature not introspectable (e.g. C extension): don't guess.
+            return
+        accepted.update(
+            param.name
+            for param in signature.parameters.values()
+            if param.kind
+            in (
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                inspect.Parameter.KEYWORD_ONLY,
+            )
+        )
+
+    for key in _BEHAVIORAL_DATASET_CONFIG_KEYS:
+        if key not in data_config or key in accepted:
+            continue
+        value = data_config[key]
+        if value is None or (key == "split_validation_size" and not value):
+            continue
+        warnings.warn(
+            f"Dataset config key {key}={value!r} is not supported by "
+            f"{dataset_class.__name__} and will be ignored. Remove the key, or "
+            "use a dataset class that supports it.",
+            stacklevel=3,
+        )
+
+
 def update_single_dataset_config(data_config: dict, default_data_config: dict) -> None:
     """Fill the single dataset config with default dataset config."""
     for key in default_data_config.keys():

tests/unit/data/datasets/test_response_dataset_registry.py

@@ -22,8 +22,10 @@
 
 from __future__ import annotations
 
+import functools
 import sys
 import types
+import warnings
 
 import pytest
 
@@ -39,7 +41,10 @@
 from nemo_rl.data.datasets.response_datasets import (
     load_response_dataset,
 )
-from nemo_rl.data.datasets.utils import resolve_external_dataset_class
+from nemo_rl.data.datasets.utils import (
+    resolve_external_dataset_class,
+    warn_on_unsupported_dataset_config_keys,
+)
 
 
 class _StubResponseDataset:
@@ -185,3 +190,83 @@ def test_load_preference_dataset_bad_dotted_path_errors():
     config = {"dataset_name": "nemo_rl_missing_module.MyDataset"}
     with pytest.raises(ValueError, match="Could not import module"):
         load_preference_dataset(config)
+
+
+# ---------------------------------------------------------------------------
+# warn_on_unsupported_dataset_config_keys (dispatcher guard, issue #3270)
+# ---------------------------------------------------------------------------
+
+
+class _SplitOnlyDataset:
+    """Stub declaring only ``split``; other behavioral keys are unsupported."""
+
+    def __init__(self, split="train", **kwargs):
+        pass
+
+
+class _KwargsForwardingDataset(_SplitOnlyDataset):
+    """Consumes keys via ``**kwargs`` and forwards to a base that declares
+    them (mirrors the intent datasets)."""
+
+    def __init__(self, **kwargs):
+        kwargs.setdefault("split", "train")
+        super().__init__(**kwargs)
+
+
+def test_warn_on_unsupported_subset():
+    with pytest.warns(UserWarning, match="subset='socratic'.*_SplitOnlyDataset"):
+        warn_on_unsupported_dataset_config_keys(
+            _SplitOnlyDataset, {"subset": "socratic"}
+        )
+
+
+def test_warn_on_unsupported_split_validation_size_and_seed():
+    with pytest.warns(UserWarning) as record:
+        warn_on_unsupported_dataset_config_keys(
+            _SplitOnlyDataset, {"split_validation_size": 0.05, "seed": 42}
+        )
+    messages = [str(w.message) for w in record]
+    assert any("split_validation_size=0.05" in m for m in messages)
+    assert any("seed=42" in m for m in messages)
+
+
+def test_no_warning_for_supported_key():
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        warn_on_unsupported_dataset_config_keys(_SplitOnlyDataset, {"split": "test"})
+
+
+def test_no_warning_for_none_or_zero_values():
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        warn_on_unsupported_dataset_config_keys(
+            _SplitOnlyDataset, {"subset": None, "split_validation_size": 0.0}
+        )
+
+
+def test_no_warning_when_base_class_declares_key():
+    """MRO-aware: subclasses forwarding **kwargs to a declaring base must
+    not be flagged."""
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        warn_on_unsupported_dataset_config_keys(
+            _KwargsForwardingDataset, {"split": "validation"}
+        )
+
+
+def test_warns_through_functools_partial():
+    """Registry entries like the AIME variants are ``functools.partial``."""
+    wrapped = functools.partial(_SplitOnlyDataset, split="train")
+    with pytest.warns(UserWarning, match="_SplitOnlyDataset"):
+        warn_on_unsupported_dataset_config_keys(wrapped, {"subset": "all"})
+
+
+def test_load_response_dataset_warns_on_swallowed_key(monkeypatch, stub_module):
+    """End-to-end: the dispatcher warns before instantiating a class that
+    would silently swallow a behavioral key."""
+    config = {
+        "dataset_name": f"{stub_module}.StubResponseDataset",
+        "subset": "socratic",
+    }
+    with pytest.warns(UserWarning, match="subset='socratic'"):
+        load_response_dataset(config)