feat: support multiple datasets for response dataset

[yuki-97]

Related issue: #1049

Usage

data:
  _override_: true # override the data config instead of merging with it
  # other data settings, see `examples/configs/sft.yaml` for more details
  ...
  # dataset settings
  train:
    # train dataset 1
    - dataset_name: OpenMathInstruct-2
      split_validation_size: 0.05 # use 5% of the training data as validation data
      seed: 42  # seed for train/validation split when split_validation_size > 0
    # train dataset 2
    - dataset_name: DeepScaler
  validation:
    # validation dataset 1
    - dataset_name: AIME2024
      repeat: 16
    # validation dataset 2
    - dataset_name: DAPOMathAIME2024
  # default settings for all datasets
  default:
    ...

Summary by CodeRabbit

• New Features

  • Added multi-dataset training support, enabling models to train on multiple datasets simultaneously.
  • Introduced configuration structure for multi-dataset GRPO experimentation with customizable dataset-specific parameters.

• Tests

  • Added functional test coverage for multi-dataset training workflows.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Introduces multi-dataset support for the GRPO framework through a new configuration file, refactored data handling across two modules, an override-aware config merge utility, and corresponding functional tests.

Changes

Cohort / File(s)	Summary
Configuration examples/configs/grpo_multiple_datasets.yaml	New YAML configuration defining multi-dataset GRPO setup with defaults reference, data overrides, dataset-wide settings (max_input_seq_length, shuffle, num_workers), and per-dataset train/validation lists with optional split/repeat parameters.
Config Utilities nemo_rl/utils/config.py	Added merge_with_override() utility function that honors _override_: true markers in child configs to fully override parent config subtrees, improving inheritance control in load_config_with_inheritance().
Data Handling examples/run_sft.py, nemo_rl/data/utils.py	Modified setup_data() and related functions to support multiple training/validation datasets instead of single dataset; collects datasets into lists, optionally applies per-dataset defaults, merges via concatenate_datasets(), and maintains per-task processor/environment bindings across merged datasets.
Testing tests/functional/grpo_multiple_datasets.sh, tests/functional/L1_Functional_Tests_GPU.sh	New functional test script orchestrating GRPO multi-dataset experiment with metrics validation; test invocation added to L1 GPU test suite.

Sequence Diagram(s)

sequenceDiagram
    participant Config as Config Loader
    participant DataList as Multiple Datasets
    participant Merge as Dataset Merger
    participant Processor as Task Processors
    participant Train as Training

    Config->>DataList: Load each dataset config
    DataList->>DataList: Apply per-dataset defaults
    DataList->>Merge: Collect all datasets
    Merge->>Merge: Concatenate datasets
    Merge->>Processor: Build per-task processors<br/>(per dataset)
    Processor->>Processor: Merge processor mappings<br/>(task_name → processor)
    Processor->>Train: Create AllTaskProcessedDataset<br/>with merged data + processors
    Train->>Train: Execute training

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Suggested labels

Run CICD

Suggested reviewers

• terrykong

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR description lacks actual test results, metric values, and validation confirmation for the major multi-dataset feature despite functional tests being added.	Add test execution results showing actual metric values, convergence validation, and confirmation that all code review issues are resolved.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately and concisely describes the main objective: adding support for multiple datasets in response/training data handling. It directly reflects the core changes across the codebase.
Docstring Coverage	✅ Passed	Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (3)

nemo_rl/utils/config.py (1)

1-1: Update the copyright header year to include 2026.

Line 1 still shows 2025; please update to include the current year. As per coding guidelines, headers must include the current year.

📝 Suggested header update

-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2025-2026, NVIDIA CORPORATION.  All rights reserved.

examples/run_sft.py (1)

1-1: Update the copyright header year to include 2026.

Line 1 still shows 2025; please update to include the current year. As per coding guidelines, headers must include the current year.

📝 Suggested header update

-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2025-2026, NVIDIA CORPORATION.  All rights reserved.

nemo_rl/data/utils.py (1)

1-1: Update the copyright header year to include 2026.

Line 1 still shows 2025; please update to include the current year. As per coding guidelines, headers must include the current year.

📝 Suggested header update

-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2025-2026, NVIDIA CORPORATION.  All rights reserved.

🤖 Fix all issues with AI agents

In `@nemo_rl/data/utils.py`:
- Around line 75-77: The normalization logic in utils.py currently only wraps
single-dataset entries when data_config["train"] or data_config["validation"] is
a plain dict, but misses omegaconf.DictConfig; import DictConfig from omegaconf
and update the two isinstance checks (the one that normalizes
data_config["train"] and the one for data_config["validation"]) to check
isinstance(..., (dict, DictConfig)) so single-dataset DictConfig objects are
wrapped into a list before iterating (this fixes errors in
load_response_dataset).

In `@tests/functional/grpo_multiple_datasets.sh`:
- Around line 20-40: The shell command at the end of
tests/functional/grpo_multiple_datasets.sh uses an unquoted $@ which can break
arguments with spaces; update the invocation to use "$@" (i.e., replace $@ with
"$@") so all passed overrides/arguments preserve their boundaries when appended
to the uv run command before the redirection and tee pipeline.

🧹 Nitpick comments (1)

nemo_rl/utils/config.py (1)

30-44: Consider supporting nested _override_ markers for consistency and future-proofing.

Currently, the function only processes top-level _override_ markers. While no nested markers are used today, a recursive approach would ensure consistent behavior if nested overrides are ever needed.

Suggested recursive handling

+def _apply_override_markers(
+    base_cfg: DictConfig, override_cfg: DictConfig
+) -> None:
+    for key, value in list(override_cfg.items()):
+        if isinstance(value, DictConfig):
+            if value.get("_override_", False):
+                value.pop("_override_")
+                base_cfg.pop(key, None)
+            else:
+                child_base = base_cfg.get(key)
+                if isinstance(child_base, DictConfig):
+                    _apply_override_markers(child_base, value)
+                else:
+                    _apply_override_markers(OmegaConf.create({}), value)
+
 def merge_with_override(
     base_config: DictConfig, override_config: DictConfig
 ) -> DictConfig:
     """Merge configs with support for _override_ marker to completely override sections."""
-    for key in list(override_config.keys()):
-        if isinstance(override_config[key], DictConfig):
-            if override_config[key].get("_override_", False):
-                # remove the _override_ marker
-                override_config[key].pop("_override_")
-                # remove the key from base_config so it won't be merged
-                if key in base_config:
-                    base_config.pop(key)
+    _apply_override_markers(base_config, override_config)