[RL] Fix shape mismatch on tail batch in GRPO training

src/maxtext/examples/chat_templates/gpt_oss_rl.json

@@ -0,0 +1,4 @@
+{
+  "SYSTEM_PROMPT": "You are given a problem. Think about the problem and provide your reasoning. Place it between {reasoning_start_token} and {reasoning_end_token}. Then, provide the final answer (i.e., just one numerical value) between {solution_start_token} and {solution_end_token}.",
+  "TEMPLATE": "<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\nKnowledge cutoff: 2024-06\nCurrent date: 2026-06-19\n\nReasoning: medium\n\n# Valid channels: analysis, commentary, final. Channel must be included for every message.\nCalls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>user<|message|><start_of_turn>user\n{system_prompt}\n\n{question}<end_of_turn>\n<start_of_turn>model<|end|><|start|>assistant"
+}

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

@@ -226,9 +226,9 @@ def math_verify_pool(
 
   cpu_count = multiprocessing.cpu_count()
   if num_procs is None:
-    num_procs = min(_DEFAULT_MAX_PROCS, len(items), cpu_count)
+    num_procs = min(_DEFAULT_MAX_PROCS, cpu_count)
   else:
-    num_procs = max(1, min(num_procs, len(items), cpu_count))
+    num_procs = max(1, min(num_procs, cpu_count))
 
   cnt = 0
   pool = _get_pool(num_procs)

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

@@ -339,7 +339,7 @@ def _use_raw_prompt(x):
   dataset_size = int(trainer_config.num_batches * trainer_config.batch_size * trainer_config.train_fraction)
   train_dataset = train_dataset[:dataset_size]
   train_dataset = train_dataset.repeat(trainer_config.num_epoch)
-  train_dataset = train_dataset.to_iter_dataset().batch(trainer_config.batch_size)
+  train_dataset = train_dataset.to_iter_dataset().batch(trainer_config.batch_size, drop_remainder=True)
 
   if trainer_config.num_test_batches > 0:
     # eval_batch_size = -1 (default) → use trainer_config.batch_size (legacy
@@ -358,7 +358,7 @@ def _use_raw_prompt(x):
     test_dataset = test_dataset[
         trainer_config.test_batch_start_index : trainer_config.num_test_batches * eval_batch_size_for_eval
     ]
-    test_dataset = test_dataset.to_iter_dataset().batch(eval_batch_size_for_eval)
+    test_dataset = test_dataset.to_iter_dataset().batch(eval_batch_size_for_eval, drop_remainder=True)
 
   return train_dataset, test_dataset
 

tests/post_training/unit/train_rl_test.py

@@ -359,6 +359,7 @@ def get_filtered_data_side_effect(dataset_name, model_tokenizer, template_config
         data_shuffle_seed=42,
         max_prefill_predict_length=10,
         batch_size=2,
+        eval_batch_size=1,
         num_batches=2,
         train_fraction=1.0,
         num_epoch=1,
@@ -422,7 +423,8 @@ def test_prepare_datasets_with_split(self, mock_load):
         data_template_path="maxtext/examples/chat_templates/gsm8k_rl.json",
         data_shuffle_seed=42,
         num_batches=1,
-        batch_size=5,
+        batch_size=2,
+        eval_batch_size=1,
         train_fraction=1.0,
         num_epoch=1,
         num_test_batches=1,