[yuki-97]

What does this PR do ?

1. Refactor and add load_response_dataset, load_preference_dataset APIs.
2. Support load dataset from local or Hugging Face:
   1. ResponseDataset for SFT and RL.
   2. BinaryPreferenceDataset and PreferenceDataset for DPO and RM.
3. Refactor structure to:

   ├── __init__.py
   ├── processed_dataset.py
   ├── utils.py
   ├── eval_datasets
   │   ├── __init__.py
   │   ├── local_math_dataset.py
   │   └── xxx_dataset.py
   ├── preference_datasets
   │   ├── __init__.py
   │   ├── binary_preference_dataset.py
   │   ├── preference_dataset.py
   │   └── xxx_dataset.py
   └── response_datasets
       ├── __init__.py
       ├── response_dataset.py
       └── xxx_dataset.py
   

4. Remove data_cls which is duplicated with dataset_name. The logic will become below after this PR:
   1. If dataset_name is a built-in dataset, we'll use that dataset class as before.
   2. If dataset_name is in ["ResponseDataset", "BinaryPreferenceDataset", "PreferenceDataset"], we'll load data from local or Hugging Face through train_data_path and val_data_path.

Issues

Related #909.

Usage

1. For datasets supported by default (e.g. squad, open_assistant, HelpSteer3, etc.), the usage is the same as before:

   data:
     dataset_name: "squad"
     ...

2. For prompt-response data for SFT and RL, we'll use ResponseDataset:

   data:
     dataset_name: ResponseDataset
     train_data_path: <PathToTrainingDataset>
     val_data_path: <PathToValidationDataset>
     input_key: <QuestionKey>, default is "input"
     output_key: <AnswerKey>, default is "output"
     train_split: <TrainSplit>, used for HuggingFace datasets, default is None
     val_split: <ValSplit>, used for HuggingFace datasets, default is None
     ...

3. For preference data for DPO and RM, we'll use PreferenceDataset:

   data:
     dataset_name: PreferenceDataset
     train_data_path: <PathToTrainingDataset>
     val_data_paths:
       <NameOfValidationDataset1>: <PathToValidationDataset1>
       ...
     train_split: <TrainSplit>, used for HuggingFace datasets, default is None
     val_split: <ValSplit>, used for HuggingFace datasets, default is None
     ...

4. For binary preference data for DPO and RM, we'll use BinaryPreferenceDataset:

   data:
     dataset_name: BinaryPreferenceDataset
     train_data_path: <PathToTrainingDataset>
     val_data_path: <PathToValidationDataset>
     prompt_key: <PromptKey>, default is "prompt"
     chosen_key: <ChosenKey>, default is "chosen"
     rejected_key: <RejectedKey>, default is "rejected"
     train_split: <TrainSplit>, used for HuggingFace datasets, default is None
     val_split: <ValSplit>, used for HuggingFace datasets, default is None
     ...

Test Result

Run well on below configs and some other configs.

"examples/configs/recipes/llm/sft-llama3.2-1b-1n8g-fsdp2tp1.v3.yaml"
"examples/configs/recipes/llm/dpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v2.yaml"
"examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K.yaml"
"examples/configs/recipes/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n2g-dtensor2tp1.v1.yaml"

convergence test with llama8b

sft-llama3.2-1b-1n8g-fsdp2tp1.v3

dpo-llama3.2-1b-instruct-1n8g-fsdp2tp1.v2

grpo-deepscaler-1.5b-8K

vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n2g-dtensor2tp1.v1

Summary by CodeRabbit

• New Features

  • Unified dataset loaders (response, preference, eval); new ResponseDataset and BinaryPreferenceDataset; HelpSteer3 and Tulu3Preference added as built-in defaults with on‑the‑fly HuggingFace downloads and multi‑validation support.

• Documentation

  • Expanded guides and example configs with concrete dataset blocks, keys (paths/input/output/chosen/rejected) and split options; remote CSV eval guidance.

• Refactor

  • Reorganized data package, centralized dataset utilities/exports, consolidated AIME variants into a single AIMEDataset; moved dataset helpers to a shared utils module.

• Tests

  • Added loader tests; removed/skipped flaky remote-download tests.

[parthchadha]

LGTM! @terrykong can you review as well