Merge pull request #3580 from AI-Hypercomputer:anisha-microbatch

PiperOrigin-RevId: 895488431

Author: Google-ML-Automation

src/maxtext/configs/post_train/rl.yml

@@ -111,7 +111,8 @@ batch_size: 1
 num_batches: 4
 # A batch can be split into multiple micro batches for memory management
 # and/or async sampling and training.
-micro_batch_size: -1
+train_micro_batch_size: -1
+rollout_micro_batch_size: -1
 # Keep `num_test_batches` low so that evaluation runs quickly. It can be
 # increased to a max. of 330 (if batch size is 4).
 num_test_batches: 5  # 200

src/maxtext/configs/types.py

@@ -1733,7 +1733,8 @@ class RLDataset(BaseModel):
   num_test_batches: int = Field(5, description="Number of batches for RL evaluation.")
   test_batch_start_index: int = Field(0, description="Start index for the test dataset")
   train_fraction: float = Field(1.0, description="Fraction of the dataset to be used for training.")
-  micro_batch_size: int = Field(-1, description="Micro batch size for rollout and training.")
+  train_micro_batch_size: int = Field(-1, description="Micro batch size for training.")
+  rollout_micro_batch_size: int = Field(-1, description="Micro batch size for rollout.")
 
 
 class RLEvaluation(BaseModel):

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

@@ -461,7 +461,10 @@ def create_rl_components(
     checkpoint_dir = None
 
   # Set up micro batching
-  micro_batch_size = None if trainer_config.micro_batch_size == -1 else trainer_config.micro_batch_size
+  train_micro_batch_size = None if trainer_config.train_micro_batch_size == -1 else trainer_config.train_micro_batch_size
+  rollout_micro_batch_size = (
+      None if trainer_config.rollout_micro_batch_size == -1 else trainer_config.rollout_micro_batch_size
+  )
 
   # Setup metrics logging
   metrics_logging_options = metrics_logger.MetricsLoggerOptions(
@@ -511,8 +514,8 @@ def create_rl_components(
           eval_every_n_steps=trainer_config.eval_interval,
           max_steps=max_train_steps,
           mini_batch_size=trainer_config.batch_size,
-          train_micro_batch_size=micro_batch_size,
-          rollout_micro_batch_size=micro_batch_size,
+          train_micro_batch_size=train_micro_batch_size,
+          rollout_micro_batch_size=rollout_micro_batch_size,
           metrics_logging_options=metrics_logging_options,
           profiler_options=profiler_options,
           checkpoint_root_directory=checkpoint_dir,

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

@@ -15,8 +15,9 @@
 # pylint: disable=bare-except, consider-using-generator, chained-comparison, broad-exception-caught
 """RL Utils Module."""
 import re
+import uuid
+from etils import epath
 import optax
-from maxtext.utils import max_logging
 import numpy as np
 
 
@@ -433,13 +434,6 @@ def check_numbers(prompts, completions, answer, tmvp_config, **kargs):
       extracted_responses.append(fallback_matches[-1].strip() if fallback_matches else None)
 
   scores = []
-  if tmvp_config.debug.rl:
-    max_logging.log("START ============================")
-    max_logging.log(f"Question: {question[0]}")
-    max_logging.log(f"Answer: {answer[0]}")
-    max_logging.log(f"Response: {completions[0]}")
-    max_logging.log(f"Extracted: {extracted_responses[0]}")
-    max_logging.log("END ==============================")
 
   for guess, true_answer in zip(extracted_responses, answer):
     if guess is None:
@@ -469,6 +463,20 @@ def check_numbers(prompts, completions, answer, tmvp_config, **kargs):
         scores.append(tmvp_config.reward_exact_answer if guess_val == true_val else 0.0)
       except:
         scores.append(0)
+  if tmvp_config.debug.rl:
+    debug_log_path = epath.Path(tmvp_config.base_output_directory) / tmvp_config.run_name / "debug_rl_logs"
+    debug_log_path.mkdir(parents=True, exist_ok=True)
+    log_file = debug_log_path / f"check_numbers_{uuid.uuid4().hex}.txt"
+    log_content = (
+        "START ============================\n"
+        f"Question: {question[0]}\n"
+        f"Answer: {answer[0]}\n"
+        f"Response: {completions[0]}\n"
+        f"Extracted: {extracted_responses[0]}\n"
+        f"Reward Score: {scores[0]}\n"
+        "END ==============================\n"
+    )
+    log_file.write_text(log_content)
 
   return scores