perf: Add a field in SFT data config to modify num_workers for loading data (#1143)

Signed-off-by: Kate Cheng <yunhsuanc@nvidia.com>

Author: katec846

examples/configs/sft.yaml

@@ -136,6 +136,7 @@ data:
   add_eos: true
   add_generation_prompt: false
   shuffle: true
+  num_workers: 1
 
   dataset_name: "squad"
   # You can use custom response datasets for training and validation. For example:

nemo_rl/algorithms/sft.py

@@ -138,6 +138,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:
@@ -152,6 +153,7 @@ def setup(
         shuffle=False,
         collate_fn=rl_collate_fn,
         drop_last=False,
+        num_workers=data_config["num_workers"],
     )
 
     # ==========================

nemo_rl/data/__init__.py

@@ -33,6 +33,11 @@ class DataConfig(TypedDict):
     download_dir: NotRequired[str]
     train_data_path: NotRequired[str]
     val_data_paths: NotRequired[dict[str, str]]
+    # Number of data loader workers.
+    # Set to 8 or 10 for large batches to improve loading speed.
+    # This saturates CPU threads without consuming too much memory
+    # However, setting it too high might cause memory issues for long seqlens.
+    num_workers: NotRequired[int]
 
 
 class MathDataConfig(DataConfig):