refactor: split train and val dataset in response dataset

[yuki-97]

Related issue: #1050

1. Split train and val in build-in dataset, so that we could unblock multiple dataset support.
2. Unify the built-in datasets under nemo_rl/data/datasets/response_datasets/ into a similar format.
3. Remove duplicated dataset name: clevr_cogent and openmathinstruct2.

New Param
Add a new param split_validation_size to handle the case that one dataset is used for both training and validation. (e.g., OpenMathInstruct-2 in examples/configs/grpo_math_1B.yaml)

1. If data.train.split_validation_size > 0 and data.validation is None, will use part of the training dataset as validation dataset.
2. If data.train.split_validation_size > 0 and data.validation is not None, will use both "part of the training dataset" and "provided validation dataset" as validation dataset.

Usage

train:
  dataset_name: ResponseDataset
  data_path: <PathToTrainingDataset>  # e.g., /path/to/local/dataset.jsonl or hf_org/hf_dataset_name (HuggingFace)
  input_key: <QuestionKey>, default is "input"
  output_key: <AnswerKey>, default is "output"
  split: <TrainSplit>, default is None  # used for HuggingFace datasets
  split_validation_size: 0.05 # use 5% of the training data as validation data
validation:
  dataset_name: ResponseDataset
  data_path: <PathToValidationDataset>
  input_key: <QuestionKey>, default is "input"
  output_key: <AnswerKey>, default is "output"
  split: <ValidationSplit>, default is None  # used for HuggingFace datasets

Test Result

algo	result
sft
sft-vlm
grpo
grpo-vlm
distillation

Summary by CodeRabbit

Release Notes

• New Features

  • Added support for separate training and validation dataset configuration with new train and validation blocks in data settings
  • Introduced new datasets: AIME2024, DAPOMath variants with automatic validation split capability
  • Enhanced dataset framework with improved flexibility for processor selection and environment configuration

• Documentation

  • Updated guides with new data configuration structure and examples for train/validation dataset setup
  • Clarified supported dataset listings and configuration format for multi-dataset training scenarios

• Bug Fixes & Improvements

  • Improved dataset loading workflow with better support for shared datasets and per-task processing
  • Streamlined configuration migration from flat to nested dataset structure across all example configs

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

[terrykong]

some initial thoughts

since it's a big PR @ashors1 could you help as a second review?

[terrykong]

i kind of feel we shouldn't split on the fly, it makes reproducing potentially problematic. i think it's better that each dataset is static at the time of running

[yuki-97]

I think it's reproducible since it'll use seed when using train_test_split. actually we also used this before, just not expose the split_validation_size param.

RL/nemo_rl/data/datasets/response_datasets/openmathinstruct2.py

Line 59 in 48dbb37

split_ds = original_ds.train_test_split(test_size=test_size, seed=seed)

[yuki-97]

btw for the seed, which one do you think is better?

1. remove seed in the data config, pass it thru load_response_dataset using config["grpo"]["seed"].
2. keep seed in the data config, inherit from ${grpo.seed}.

[terrykong]

i see now that there are two, should we just remove this outer one to avoid dealing with surprising precedence issues if someone forgot to set one

[yuki-97]

let's discuss it here #1649 (comment).

and even if we have a default like I said in that conversation, we still need to keep it for now, because this PR only refactor the response dataset, the preference dataset will still need to use it.

[terrykong]

wdyt about just expecting users to populate the train config? then we don't have dup keys

[yuki-97]

I think we should have a default value especially when we support multiple datasets in next PR, otherwise people need to write the same things for every dataset, then the data config will be a bit redundant.

and I'm thinking if it's better to provide a default like train and validation, it seems more directly than just put it outside. wdyt?

# now
data:
    train:
        # this dataset will override prompt_key and use the default values for other vars
        - data_path: /path/to/local/train_dataset_1.jsonl
          prompt_key: question
        # this dataset will use all the default values
        - data_path: /path/to/local/train_dataset_2.jsonl
    validation:
        - data_path: /path/to/local/val_dataset.jsonl
    # will use below vars as default values if dataset doesn't specify it
    dataset_name: BinaryPreferenceDataset
    prompt_key: prompt
    chosen_key: chosen
    rejected_key: rejected
    prompt_file: null
    system_prompt_file: null
    env_name: math

# add `default`
data:
    train:
        # this dataset will override prompt_key and use the default values for other vars
        - data_path: /path/to/local/train_dataset_1.jsonl
          prompt_key: question
        # this dataset will use all the default values
        - data_path: /path/to/local/train_dataset_2.jsonl
    validation:
        - data_path: /path/to/local/val_dataset.jsonl
    default:
        # will use below vars as default values if dataset doesn't specify it
        dataset_name: BinaryPreferenceDataset
        prompt_key: prompt
        chosen_key: chosen
        rejected_key: rejected
        prompt_file: null
        system_prompt_file: null
        env_name: math

[terrykong]

do you think we need to keep this? it kind of seems like we could do without it

[yuki-97]

some datasets are associated with processor (e.g., helpsteer3), so we need to keep it for now.
I think we don't need to keep this eventually, as designed I'll make the processor associated with algorithm instead of dataset in later PR.

[yuki-97]

track it here #1658.

[terrykong]

what was this comment referring to? not sure i follow from the changes you made

[yuki-97]

for these two commits: 487d354, d9344a6.
we have openmathinstruct2 and OpenMathInstruct-2 both use OpenMathInstruct2Dataset, clevr_cogent and clevr-cogent both use CLEVRCoGenTDataset before.

[terrykong]

this looks a lot cleaner :)

[yuki-97]

@terrykong shall we enable this?