[data, trainer] feat: add support for limiting samples from dataset

e.g.: For RLHFDataset, `filter_overlong_prompts` can be very
expensive and it will be good to add support to limit the sample
size before we do this when the dataset is very large.

Also add support for other kinds of datasets for unification.

Author: HollowMan6

docs/examples/config.rst

@@ -17,6 +17,8 @@ Data
      tokenizer: null
      train_files: ~/data/rlhf/gsm8k/train.parquet
      val_files: ~/data/rlhf/gsm8k/test.parquet
+     train_max_samples: -1  # set to -1 to use full dataset
+     val_max_samples: -1  # set to -1 to use full dataset
      prompt_key: prompt
      max_prompt_length: 512
      max_response_length: 512
@@ -41,6 +43,10 @@ Data
   HDFS path to local path.
 - ``data.val_files``: Validation parquet. Can be a list or a single
   file.
+- ``data.train_max_samples``: Maximum number of samples to use from the
+  training dataset. Set to -1 to use the full dataset.
+- ``data.val_max_samples``: Maximum number of samples to use from the
+  validation dataset. Set to -1 to use the full dataset.
 - ``data.prompt_key``: The field in the dataset where the prompt is
   located. Default is 'prompt'.
 - ``data.max_prompt_length``: Maximum prompt length. All prompts will be

examples/split_placement/config/ppo_trainer_split.yaml

@@ -12,6 +12,8 @@ data:
   tokenizer: null
   train_files: ~/data/rlhf/gsm8k/train.parquet
   val_files: ~/data/rlhf/gsm8k/test.parquet
+  train_max_samples: -1  # set to -1 to use full dataset
+  val_max_samples: -1  # set to -1 to use full dataset
   prompt_key: prompt
   max_prompt_length: 512
   max_response_length: 512

recipe/entropy/main_entropy.py

@@ -162,8 +162,16 @@ def run(self, config):
 
         from verl.utils.dataset.rl_dataset import collate_fn
 
-        train_dataset = create_rl_dataset(config.data.train_files, config.data, tokenizer, processor)
-        val_dataset = create_rl_dataset(config.data.val_files, config.data, tokenizer, processor)
+        train_dataset = create_rl_dataset(
+            config.data.train_files,
+            config.data,
+            tokenizer,
+            processor,
+            max_samples=config.data.get("train_max_samples", -1),
+        )
+        val_dataset = create_rl_dataset(
+            config.data.val_files, config.data, tokenizer, processor, max_samples=config.data.get("val_max_samples", -1)
+        )
         train_sampler = create_rl_sampler(config.data, train_dataset)
         trainer = RayEntropyTrainer(
             config=config,
@@ -183,7 +191,7 @@ def run(self, config):
         trainer.fit()
 
 
-def create_rl_dataset(data_paths, data_config, tokenizer, processor):
+def create_rl_dataset(data_paths, data_config, tokenizer, processor, max_samples: int = -1):
     """Create a dataset.
 
     Arguments:
@@ -216,6 +224,7 @@ def create_rl_dataset(data_paths, data_config, tokenizer, processor):
         tokenizer=tokenizer,
         processor=processor,
         config=data_config,
+        max_samples=max_samples,
     )
 
     return dataset

recipe/one_step_off_policy/main_ppo.py

@@ -212,8 +212,16 @@ def run(self, config):
         from verl.utils.dataset.rl_dataset import collate_fn
 
         # Create training and validation datasets.
-        train_dataset = create_rl_dataset(config.data.train_files, config.data, tokenizer, processor)
-        val_dataset = create_rl_dataset(config.data.val_files, config.data, tokenizer, processor)
+        train_dataset = create_rl_dataset(
+            config.data.train_files,
+            config.data,
+            tokenizer,
+            processor,
+            max_samples=config.data.get("train_max_samples", -1),
+        )
+        val_dataset = create_rl_dataset(
+            config.data.val_files, config.data, tokenizer, processor, max_samples=config.data.get("val_max_samples", -1)
+        )
         train_sampler = create_rl_sampler(config.data, train_dataset)
 
         # Initialize the PPO trainer.

recipe/spin/spin_trainer.py

@@ -393,11 +393,19 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl
 
         if train_dataset is None:
             train_dataset = create_rl_dataset(
-                self.config.data.train_files, self.config.data, self.tokenizer, self.processor
+                self.config.data.train_files,
+                self.config.data,
+                self.tokenizer,
+                self.processor,
+                max_samples=self.config.data.get("train_max_samples", -1),
             )
         if val_dataset is None:
             val_dataset = create_rl_dataset(
-                self.config.data.val_files, self.config.data, self.tokenizer, self.processor
+                self.config.data.val_files,
+                self.config.data,
+                self.tokenizer,
+                self.processor,
+                max_samples=self.config.data.get("val_max_samples", -1),
             )
         self.train_dataset, self.val_dataset = train_dataset, val_dataset
 

tests/special_e2e/sft/test_sp_loss_match.py

@@ -112,8 +112,12 @@ def create_trainer(config):
 
     local_model_path = copy_to_local(src=config.model.partial_pretrain, verbose=True)
     tokenizer = hf_tokenizer(local_model_path, trust_remote_code=config.model.trust_remote_code)
-    train_dataset = create_sft_dataset(config.data.train_files, config.data, tokenizer)
-    val_dataset = create_sft_dataset(config.data.val_files, config.data, tokenizer)
+    train_dataset = create_sft_dataset(
+        config.data.train_files, config.data, tokenizer, max_samples=config.data.get("train_max_samples", -1)
+    )
+    val_dataset = create_sft_dataset(
+        config.data.val_files, config.data, tokenizer, max_samples=config.data.get("val_max_samples", -1)
+    )
 
     return FSDPSFTTrainer(
         config=config,

tests/trainer/config/legacy_ppo_megatron_trainer.yaml

@@ -2,6 +2,8 @@ data:
   tokenizer: null
   train_files: ~/data/rlhf/gsm8k/train.parquet
   val_files: ~/data/rlhf/gsm8k/test.parquet
+  train_max_samples: -1  # set to -1 to use full dataset
+  val_max_samples: -1  # set to -1 to use full dataset
   prompt_key: prompt
   reward_fn_key: data_source
   max_prompt_length: 512

tests/trainer/config/legacy_ppo_trainer.yaml

@@ -22,6 +22,16 @@ data:
   # Validation parquet. Can be a list or a single file.
   val_files: ~/data/rlhf/gsm8k/test.parquet
 
+  # Maximum sample length to be used.
+  # Set to -1 to use full dataset, otherwise, randomly
+  # select the specified number of samples from train dataset
+  train_max_samples: -1
+
+  # Maximum sample length to be used.
+  # Set to -1 to use full dataset, otherwise, randomly
+  # select the specified number of samples from val dataset
+  val_max_samples: -1
+
   # The field in the dataset where the prompt is located. Default is 'prompt'.
   prompt_key: prompt
 

tests/utils/dataset/test_rl_dataset_on_cpu.py

@@ -66,6 +66,25 @@ def test_rl_dataset():
     print(f"\n\noutput: {output}")
 
 
+def test_rl_dataset_with_max_samples():
+    from verl.utils import hf_tokenizer
+    from verl.utils.dataset.rl_dataset import RLHFDataset
+
+    tokenizer = hf_tokenizer("deepseek-ai/deepseek-coder-1.3b-instruct")
+    local_path = get_gsm8k_data()
+    config = OmegaConf.create(
+        {
+            "prompt_key": "prompt",
+            "max_prompt_length": 256,
+            "filter_overlong_prompts": True,
+            "filter_overlong_prompts_workers": 2,
+            "max_samples": 5,
+        }
+    )
+    dataset = RLHFDataset(data_files=local_path, tokenizer=tokenizer, config=config, max_samples=5)
+    assert len(dataset) == 5
+
+
 def test_image_rl_data():
     from verl.utils import hf_processor, hf_tokenizer
     from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn

tests/utils/dataset/test_sft_dataset_on_cpu.py

@@ -72,3 +72,26 @@ def test_sft_dataset():
     output = tokenizer.batch_decode([data])[0]
     assert len(output) > 1
     assert isinstance(output, str)
+
+
+def test_sft_dataset_with_max_samples():
+    tokenizer = hf_tokenizer("deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct")
+    local_path = get_gsm8k_data()
+    from omegaconf import OmegaConf
+
+    dataset = SFTDataset(
+        parquet_files=local_path,
+        tokenizer=tokenizer,
+        config=OmegaConf.create(
+            {
+                "prompt_key": "extra_info",
+                "prompt_dict_keys": ["question"],
+                "response_key": "extra_info",
+                "response_dict_keys": ["answer"],
+                "max_length": 512,
+            }
+        ),
+        max_samples=5,
+    )
+
+    assert len(dataset) == 5