Merge pull request #2713 from AI-Hypercomputer:xfgu-dp-rl

PiperOrigin-RevId: 842472785

Author: Google-ML-Automation

src/MaxText/configs/rl.yml

@@ -23,6 +23,10 @@ sampler_devices_fraction: 0.5
 chips_per_vm: 4  # depends on hardware, for v5p this is 4
 num_trainer_slices: -1
 num_samplers_slices: -1
+# Only specify rollout_data_parallelism when you would like to use more than one model
+# replicas in rollout. If not specified, rollout_tensor_parallelism will be auto-determined.
+rollout_data_parallelism: -1
+rollout_tensor_parallelism: -1
 
 # ====== Reproducibility ======
 data_shuffle_seed: 42

src/MaxText/configs/types.py

@@ -1333,6 +1333,14 @@ class RLHardware(BaseModel):
   use_pathways: bool = Field(True, description="Whether to use Pathways for multihost orchestration.")
   num_trainer_slices: int = Field(-1, description="Number of slices for the trainer.")
   num_samplers_slices: int = Field(-1, description="Number of slices for the samplers.")
+  rollout_data_parallelism: int = Field(
+      -1,
+      description="Total model replicas for rollout. It should only be specified when you would like to use more "
+      "than one model replica in rollout.",
+  )
+  rollout_tensor_parallelism: int = Field(
+      -1, description="Tensor parallelism per replica for rollout. If not specified, it will be auto-determined."
+  )
 
 
 class VLLM(BaseModel):

src/MaxText/rl/train_rl.py

@@ -206,6 +206,36 @@ def setup_configs_and_devices(argv: Sequence[str]):
   return trainer_config, sampler_config, trainer_devices, sampler_devices
 
 
+def get_rollout_kwargs_for_data_parallelism(sampler_config, num_sampler_devices):
+  """Get rollout kwargs for vLLM rollout when using data parallelism."""
+  dp = sampler_config.rollout_data_parallelism
+  if dp == -1:
+    return {}
+
+  rollout_kwargs = {}
+  tp = sampler_config.rollout_tensor_parallelism
+
+  if tp == -1:
+    if num_sampler_devices % dp != 0:
+      raise ValueError(
+          f"num_sampler_devices({num_sampler_devices}) must be divisible by "
+          f"rollout_data_parallelism({dp}) "
+          f"when rollout_tensor_parallelism is -1."
+      )
+    tp = num_sampler_devices // dp
+  elif tp * dp != num_sampler_devices:
+    raise ValueError(
+        f"rollout_tensor_parallelism({tp}) * "
+        f"rollout_data_parallelism({dp}) "
+        f"!= len(sampler_devices)({num_sampler_devices})"
+    )
+  rollout_kwargs["tensor_parallel_size"] = tp
+  rollout_kwargs["data_parallel_size"] = dp
+  rollout_kwargs["rollout_vllm_async_scheduling"] = True
+
+  return rollout_kwargs
+
+
 def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
   """
   Run RL training with the provided configuration.
@@ -360,6 +390,7 @@ def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
           rollout_vllm_hbm_utilization=trainer_config.hbm_utilization_vllm,
           rollout_vllm_tpu_backend_type="jax",
           rollout_vllm_swap_space_size_gb=trainer_config.swap_space_vllm_gb,
+          **get_rollout_kwargs_for_data_parallelism(sampler_config, len(sampler_devices)),
       ),
   )
   grpo_config = GrpoConfig(