Add eval_batch_size knob for faster post-train RL evaluation

[py4]

Post-train RL evaluation is currently batched at trainer_config.batch_size, which for GRPO is intentionally small (e.g. 4 prompts per step × 8 generations = 32 trajectories — set by the GRPO recipe to keep trainer HBM workable for the backward pass). At eval time this is wasteful: eval is greedy decode only (no backward), so the trainer's per-step memory budget doesn't apply, and vLLM rollout has many DP replicas sitting idle when only 4 prompts are dispatched per call.

This PR adds an rl.eval_batch_size knob (default -1 = use batch_size, preserving old behavior) that overrides the batch dimension during dataset preparation for the test split. Setting it to e.g. 128 on a sampler with 8 DP replicas gives a ~32× eval throughput improvement on TPU without affecting training behavior.

Changes (3 files, +21/-2 lines):

• src/maxtext/configs/types.py: new Pydantic field RLDataset.eval_batch_size: int = -1
• src/maxtext/configs/post_train/rl.yml: default eval_batch_size: -1 + comment
• src/maxtext/trainers/post_train/rl/train_rl.py:prepare_datasets: when set and positive, use eval_batch_size for the test split's slice + .batch(...) call

NOTE: total eval examples = num_test_batches * eval_batch_size, so users adjusting eval_batch_size should adjust num_test_batches to keep total eval set size constant.

Backward compatible: default -1 falls back to batch_size (identical to old behavior). No effect on training path.

Checklist

• Tested locally on TPU v6e 8×8: with eval_batch_size=128 (8 DP replicas), eval over 1408 examples completes in ~3 min vs ~30+ min at batch_size=4
• Backward compatible: default -1 preserves existing behavior bit-for-bit
• No effect on training path (only eval-side dataset preparation touched)
• No effect on non-RL paths (only RL Pydantic config + RL trainer touched)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[py4]

Updated description to match the style of sibling RLDataset fields: "Batch size for RL evaluation." PTAL.

[A9isha]

Thanks a lot Pooya!

trainer-side KV cache is a bit confusing, please consider updating the description given that KV cache is for rollout.

[py4]

@A9isha thanks for the catch — you're right, KV cache lives on the rollout (vLLM) side, not the trainer. Updated both the PR body and commit message to say the trainer per-step batch is sized for backward-pass HBM (activations + grads), and noted that eval is greedy decode only so that constraint doesn't apply at eval time.