[Tunix] Enable micro_batch_size for rollout and reference models in PPO learner.

PiperOrigin-RevId: 827589436

Author: lc5211

tunix/rl/ppo/ppo_learner.py

@@ -238,6 +238,7 @@ def _generate_and_compute_advantage(
     # "experiences".
     completion_output = self.rl_cluster.generate(
         prompts=training_input["prompts"],
+        micro_batch_size=self._rollout_micro_batch_size,
     )
     completion_ids = completion_output.tokens
     prompt_ids = completion_output.left_padded_prompt_tokens
@@ -261,6 +262,7 @@ def _generate_and_compute_advantage(
           completion_tokens=completion_ids,
           pad_id=pad_value,
           eos_id=eos_value,
+          micro_batch_size=self._compute_logps_micro_batch_size,
       )
     else:
       ref_per_token_logps = None
@@ -272,6 +274,7 @@ def _generate_and_compute_advantage(
     old_per_token_logps = self.rl_cluster.get_old_per_token_logps(
         prompt_tokens=prompt_ids,
         completion_tokens=completion_ids,
+        micro_batch_size=self._compute_logps_micro_batch_size,
     )
 
     # ===== Value computation ======