From c926b9956be160d8a76f42a2458e297c18981c1e Mon Sep 17 00:00:00 2001 From: Sanbao Su Date: Tue, 30 Jun 2026 18:18:19 +0000 Subject: [PATCH] Fix shape mismatch on tail batch in GRPO training and host CPU OOM - Fixes shape mismatch issues on tail batches in GRPO training by enabling drop_remainder=True on training and test datasets. - Updates unit tests (train_rl_test.py) to mock configurations compatible with drop_remainder=True and specific batch sizes. - Prevents host CPU OOM by keeping the pool size of math_verify_pool constant. - Includes a template gpt_oss_rl.json for RL chat. --- src/maxtext/examples/chat_templates/gpt_oss_rl.json | 4 ++++ src/maxtext/trainers/post_train/rl/math_verify_pool.py | 4 ++-- src/maxtext/trainers/post_train/rl/train_rl.py | 4 ++-- tests/post_training/unit/train_rl_test.py | 4 +++- 4 files changed, 11 insertions(+), 5 deletions(-) create mode 100644 src/maxtext/examples/chat_templates/gpt_oss_rl.json diff --git a/src/maxtext/examples/chat_templates/gpt_oss_rl.json b/src/maxtext/examples/chat_templates/gpt_oss_rl.json new file mode 100644 index 0000000000..9f6d522c43 --- /dev/null +++ b/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|>user\n{system_prompt}\n\n{question}\nmodel<|end|><|start|>assistant" +} diff --git a/src/maxtext/trainers/post_train/rl/math_verify_pool.py b/src/maxtext/trainers/post_train/rl/math_verify_pool.py index 2eeadf1747..c0968f184e 100644 --- a/src/maxtext/trainers/post_train/rl/math_verify_pool.py +++ b/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) diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index 9270fb5033..bd60747aeb 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/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 diff --git a/tests/post_training/unit/train_rl_test.py b/tests/post_training/unit/train_rl_test.py index 0396a1f944..1de8850999 100644 --- a/tests/post_training/unit/train_rl_test.py +++ b/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,