Add unit test for prompt filtering and refactor rl_train to extract get_datasets

Author: igorts-git

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -282,37 +282,11 @@ def get_rollout_kwargs_for_parallelism(sampler_config, num_sampler_devices):
   return rollout_kwargs
 
 
-def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
-  """
-  Run RL training with the provided configuration.
-
-  Args:
-    trainer_config: MaxText configuration for the trainer.
-    sampler_config: MaxText configuration for the sampler.
-    trainer_devices: JAX devices for the trainer.
-    sampler_devices: JAX devices for the sampler.
-  """
-  if not trainer_config.debug.rl:
-    # Apply filter to suppress noisy logs
-    noise_filter = max_logging.NoisyLogFilter()
-    logging.getLogger().addFilter(noise_filter)
-    absl_logging.get_absl_logger().addFilter(noise_filter)
-
-  max_logging.log("Starting RL Training")
-  max_logging.log(f"Ensuring TensorBoard log directory exists: {trainer_config.tensorboard_dir}")
-  if not epath.Path(trainer_config.tensorboard_dir).exists():
-    epath.Path(trainer_config.tensorboard_dir).mkdir(parents=True, exist_ok=True)
-
-  if not epath.Path(trainer_config.checkpoint_dir).exists():
-    epath.Path(trainer_config.checkpoint_dir).mkdir(parents=True)
-
-  # Number of training steps.
-  max_train_steps = int(
-      trainer_config.num_batches
-      * trainer_config.rl.num_iterations
-      * trainer_config.train_fraction
-      * trainer_config.num_epoch
-  )
+def get_datasets(
+    model_tokenizer,
+    trainer_config,
+) -> tuple[grain.IterDataset, grain.IterDataset]:
+  """Handles loading, templating, filtering, and batching of train/test datasets."""
   # ====== Data ======
   # Setup data directories
   home = os.path.expanduser("~") + "/"
@@ -323,9 +297,6 @@ def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
   if not os.path.exists(test_data_dir):
     os.makedirs(test_data_dir)
 
-  # Create model tokenizer
-  model_tokenizer = AutoTokenizer.from_pretrained(trainer_config.tokenizer_path)
-
   # Load datasets
   if trainer_config.dataset_name == "huggingface:nvidia/OpenMathInstruct-2":
     import datasets  # pylint: disable=import-outside-toplevel
@@ -334,7 +305,6 @@ def prepare_openinstructmath2_dataset(
         split: str = "train_1M",
         seed: int = 42,
         test_size: float = 0.05,
-        output_key: str = "expected_answer",
     ):
       """Load and split the OpenMathInstruct-2 dataset into train and validation sets using HF's train_test_split."""
       max_logging.log(
@@ -422,16 +392,54 @@ def _filter_long_prompts(x):
 
   if trainer_config.debug.rl:
     # Let's see how one batch of the dataset looks like!
-    if trainer_config.debug.rl:
-      for i, ele in enumerate(train_dataset):
-        if i >= 5:
-          break
-        pprint(ele)
-    if trainer_config.debug.rl:
-      for i, ele in enumerate(test_dataset):
-        if i >= 5:
-          break
-        pprint(ele)
+    for i, ele in enumerate(train_dataset):
+      if i >= 5:
+        break
+      pprint(ele)
+    for i, ele in enumerate(test_dataset):
+      if i >= 5:
+        break
+      pprint(ele)
+
+  return train_dataset, test_dataset
+
+
+def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
+  """
+  Run RL training with the provided configuration.
+
+  Args:
+    trainer_config: MaxText configuration for the trainer.
+    sampler_config: MaxText configuration for the sampler.
+    trainer_devices: JAX devices for the trainer.
+    sampler_devices: JAX devices for the sampler.
+  """
+  if not trainer_config.debug.rl:
+    # Apply filter to suppress noisy logs
+    noise_filter = max_logging.NoisyLogFilter()
+    logging.getLogger().addFilter(noise_filter)
+    absl_logging.get_absl_logger().addFilter(noise_filter)
+
+  max_logging.log("Starting RL Training")
+  max_logging.log(f"Ensuring TensorBoard log directory exists: {trainer_config.tensorboard_dir}")
+  if not epath.Path(trainer_config.tensorboard_dir).exists():
+    epath.Path(trainer_config.tensorboard_dir).mkdir(parents=True, exist_ok=True)
+
+  if not epath.Path(trainer_config.checkpoint_dir).exists():
+    epath.Path(trainer_config.checkpoint_dir).mkdir(parents=True)
+
+  # Number of training steps.
+  max_train_steps = int(
+      trainer_config.num_batches
+      * trainer_config.rl.num_iterations
+      * trainer_config.train_fraction
+      * trainer_config.num_epoch
+  )
+  # ====== Data ======
+  # Create model tokenizer
+  model_tokenizer = AutoTokenizer.from_pretrained(trainer_config.tokenizer_path)
+
+  train_dataset, test_dataset = get_datasets(model_tokenizer, trainer_config)
 
   # Load reference model
   max_logging.log("Creating reference model and also meshes for reference and rollout")

tests/unit/train_rl_test.py

@@ -16,6 +16,7 @@
 
 import unittest
 from unittest import mock
+import grain
 import pytest
 from types import SimpleNamespace
 import jax
@@ -203,6 +204,80 @@ def test_get_rollout_kwargs_errors(self):
     with self.assertRaisesRegex(ValueError, r"!= len\(sampler_devices\)"):
       train_rl.get_rollout_kwargs_for_parallelism(sampler_config, 8)
 
+  @pytest.mark.cpu_only
+  def test_prompt_filtering(self):
+    """Test that prompts longer than max_prefill_predict_length are filtered out."""
+    # Setup mocks
+    mock_tokenizer = mock.MagicMock()
+
+    # Define tokenizer side effect
+    def tokenize_side_effect(text):
+      if text == "short":
+        return [0] * 5
+      else:
+        return [0] * 15
+
+    mock_tokenizer.tokenize.side_effect = tokenize_side_effect
+
+    # Define dataset mock data
+    train_data = [{"prompts": "short"}, {"prompts": "long"}, {"prompts": "short"}, {"prompts": "long"}]
+    test_data = [{"prompts": "short"}, {"prompts": "long"}]
+    train_map_ds = grain.MapDataset.source(train_data)
+    test_map_ds = grain.MapDataset.source(test_data)
+
+    def get_dataset_side_effect(model_tokenizer, config, data_dir, split, data_files=None, dataset_name=None):
+      if split == "train":
+        return train_map_ds
+      else:
+        return test_map_ds
+
+    # Configs
+    trainer_config = SimpleNamespace(
+        debug=SimpleNamespace(rl=False),
+        tokenizer_path="dummy_path",
+        dataset_name="dummy_dataset",
+        train_split="train",
+        eval_split="eval",
+        hf_train_files=None,
+        hf_eval_files=None,
+        max_prefill_predict_length=10,
+        batch_size=2,
+        num_batches=2,
+        train_fraction=1.0,
+        num_epoch=1,
+        num_test_batches=1,
+    )
+
+    # Patch everything!
+    with (
+        mock.patch("maxtext.trainers.post_train.rl.train_rl.get_dataset", side_effect=get_dataset_side_effect),
+        mock.patch("maxtext.trainers.post_train.rl.train_rl.os.makedirs"),
+        mock.patch("maxtext.trainers.post_train.rl.train_rl.os.path.exists", return_value=True),
+    ):
+      train_dataset, test_dataset = train_rl.get_datasets(mock_tokenizer, trainer_config)
+
+      # Check filtered train dataset
+      elements = list(train_dataset)
+      # dataset_size = 4. Indices [0,1,2,3] are [short, long, short, long].
+      # Filtered results: [short, short].
+      # batch(2) will return 1 batch of 2 elements.
+      self.assertEqual(len(elements), 1)
+      batch = elements[0]
+      self.assertEqual(len(batch["prompts"]), 2)
+      for prompt in batch["prompts"]:
+        self.assertEqual(prompt, "short")
+
+      # Check filtered test dataset
+      test_elements = list(test_dataset)
+      # test_data indices [0,1] are [short, long].
+      # num_test_batches=1, batch_size=2 -> test dataset_size = 2.
+      # Filtering results: [short].
+      # batch(2) will return 1 batch of 1 element.
+      self.assertEqual(len(test_elements), 1)
+      test_batch = test_elements[0]
+      self.assertEqual(len(test_batch["prompts"]), 1)
+      self.assertEqual(test_batch["prompts"][0], "short")
+
 
 if __name__ == "__main__":
   unittest.main()