perf: Add num_workers in DPO, GRPO and SFT for loading data (#1314)

Signed-off-by: Kate Cheng <[email protected]>

Author: katec846

examples/configs/dpo.yaml

@@ -158,6 +158,7 @@ policy:
 data:
   max_input_seq_length: ${policy.max_total_sequence_length}
   shuffle: true
+  num_workers: 1
 
   dataset_name: HelpSteer3
   # You can use custom preference datasets for training and validation. For example:

examples/configs/grpo_math_1B.yaml

@@ -219,6 +219,7 @@ data:
   prompt_file: "examples/prompts/cot.txt"
   system_prompt_file: null
   shuffle: true
+  num_workers: 1
 
   dataset_name: "OpenMathInstruct-2"
   # You can use custom response datasets for training and validation. For example:

examples/configs/rm.yaml

@@ -129,6 +129,7 @@ policy:
 data:
   max_input_seq_length: ${policy.max_total_sequence_length}
   shuffle: true
+  num_workers: 1
 
   dataset_name: HelpSteer3
   # You can use custom preference datasets for training and validation. For example:

examples/configs/sft_openmathinstruct2_megatron.yaml

@@ -132,6 +132,7 @@ data:
   add_eos: true
   add_generation_prompt: true
   output_key: 'generated_solution'
+  num_workers: 1
 
 logger:
   log_dir: "logs"  # Base directory for all logs

examples/configs/vlm_grpo_3B.yaml

@@ -203,6 +203,7 @@ data:
   dataset_name: "clevr-cogent"
   split: "trainA"
   shuffle: true
+  num_workers: 1
 
 env:
   clevr-cogent:

examples/configs/vlm_grpo_3B_megatron.yaml

@@ -156,6 +156,7 @@ data:
   dataset_name: clevr-cogent
   split: trainA
   shuffle: true
+  num_workers: 1
 env:
   clevr-cogent:
     num_workers: 8

nemo_rl/algorithms/dpo.py

@@ -176,6 +176,7 @@ def setup(
             add_loss_mask=True,
         ),
         drop_last=True,
+        num_workers=data_config["num_workers"],
     )
 
     if last_checkpoint_path is not None:
@@ -198,6 +199,7 @@ def setup(
                 add_loss_mask=True,
             ),
             drop_last=False,
+            num_workers=data_config["num_workers"],
         )
         for k, v in val_dataset.items()
     }

nemo_rl/algorithms/grpo.py

@@ -207,6 +207,7 @@ def setup(
         shuffle=data_config["shuffle"],
         collate_fn=rl_collate_fn,
         drop_last=True,
+        num_workers=data_config["num_workers"],
     )
     if last_checkpoint_path is not None:
         dataloader_state_dict = torch.load(
@@ -228,6 +229,7 @@ def setup(
             batch_size=grpo_config["val_batch_size"],
             shuffle=False,
             collate_fn=rl_collate_fn,
+            num_workers=data_config["num_workers"],
         )
         print(
             f"  ✓ Validation dataloader loaded with {len(val_dataset)} samples",

nemo_rl/algorithms/rm.py

@@ -151,6 +151,7 @@ def setup(
             add_loss_mask=False,
         ),
         drop_last=True,
+        num_workers=data_config["num_workers"],
     )
 
     if last_checkpoint_path is not None:
@@ -173,6 +174,7 @@ def setup(
                 add_loss_mask=False,
             ),
             drop_last=False,
+            num_workers=data_config["num_workers"],
         )
         for k, v in val_dataset.items()
     }

tests/unit/algorithms/test_grpo.py

@@ -240,7 +240,7 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_single_node():
             "val_period": 0,
             "val_at_start": False,
         },
-        "data": {"shuffle": False},
+        "data": {"shuffle": False, "num_workers": 1},
         "logger": {},  # Config extraction requires this key
         "checkpointing": {},  # Config extraction requires this key
         "cluster": {
@@ -296,7 +296,7 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_multi_node():
             "val_period": 0,
             "val_at_start": False,
         },
-        "data": {"shuffle": False},
+        "data": {"shuffle": False, "num_workers": 1},
         "logger": {},  # Config extraction requires this key
         "checkpointing": {},  # Config extraction requires this key
         "cluster": {