feat: support multiple dataloader for grpo

[yuki-97]

Closes #1603.

Usage

uv run examples/run_grpo.py \
    --config examples/configs/grpo_multiple_datasets.yaml \
    grpo.num_prompts_per_step=32 \
    data.use_multiple_dataloader=true \
    data.num_prompts_per_dataloader=16 \
    data.custom_dataloader=examples.custom_dataloader.custom_dataloader.example_custom_dataloader

More details at https://github.com/NVIDIA-NeMo/RL/blob/yukih/multiple-dataloader/docs/guides/grpo.md#multiple-dataloaders.

Test Result

1. num_prompts_per_dataloader=16 will get 16 prompts from each dataset in every step. The data distribution of the two train datasets will change comparing to using single dataloader, so the train reward curve will have some difference, this is expected.

     use_multiple_dataloader: true
     num_prompts_per_dataloader: 16
     custom_dataloader: examples.custom_dataloader.custom_dataloader.example_custom_dataloader
     train:
       - dataset_name: OpenMathInstruct-2
         split_validation_size: 0.05 # use 5% of the training data as validation data
         seed: ${grpo.seed} # seed for train/validation split when split_validation_size > 0
       - dataset_name: DeepScaler
     validation:
       - dataset_name: AIME2024
         repeat: 16
       - dataset_name: DAPOMathAIME2024

   
2. This test is using two same OpenMathInstruct-2 dataset, so the data distribution won't change and the train reward curve could match.

     use_multiple_dataloader: true
     num_prompts_per_dataloader: 16
     custom_dataloader: examples.custom_dataloader.custom_dataloader.example_custom_dataloader
     train:
       - dataset_name: OpenMathInstruct-2
       - dataset_name: ResponseDataset
         data_path: nvidia/OpenMathInstruct-2
         input_key: problem
         output_key: expected_answer
         split: train_1M
     validation:
       - dataset_name: AIME2024
         repeat: 16
       - dataset_name: DAPOMathAIME2024

   

Summary by CodeRabbit

• New Features

  • Added support for multiple dataloaders in GRPO training with customizable dataloader configuration.
  • Introduced configuration options: use_multiple_dataloader, num_prompts_per_dataloader, and custom_dataloader.

• Documentation

  • Added comprehensive guide section documenting multiple dataloader setup, configuration parameters, and usage examples.

• Known Limitations

  • Multiple dataloaders not yet supported with async GRPO.

[RayenTian]

Thanks for your hard work!