Merge pull request #3263 from AI-Hypercomputer:nicogrande/add-stop-strings

PiperOrigin-RevId: 879152699

Author: Google-ML-Automation

src/maxtext/configs/post_train/rl.yml

@@ -149,6 +149,10 @@ enable_dp_attention: False
 # Performance tuning for samplers
 max_num_batched_tokens: null
 max_num_seqs: null
+# If True, enables asynchronous scheduling in vLLM for faster generation
+async_scheduling: True
+# stop generation when any of these strings is generated
+stop_strings: [</answer>]
 
 # ====== Checkpoint Configuration ======
 enable_checkpointing: True

src/maxtext/configs/types.py

@@ -1594,8 +1594,10 @@ class VLLM(BaseModel):
   hbm_utilization_vllm: float = Field(0.72, description="Target HBM utilization for vLLM.")
   swap_space_vllm_gb: int = Field(2, description="Swap space in GB for vLLM.")
   enable_dp_attention: bool = Field(False, description="Enable the attn_dp mesh axis in vLLM.")
+  async_scheduling: bool = Field(False, description="Enable asynchronous scheduling in vLLM.")
   max_num_batched_tokens: Optional[int] = Field(None, description="Max number of batched tokens in vLLM.")
   max_num_seqs: Optional[int] = Field(None, description="Max number of sequences in vLLM.")
+  stop_strings: Optional[list[str]] = Field(None, description="List of stop strings for vLLM decoding.")
   vllm_additional_config: dict[str, Any] = Field(default_factory=dict, description="Additional vLLM config options.")
   vllm_hf_overrides: dict[str, Any] = Field(
       default_factory=dict,

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

@@ -253,7 +253,6 @@ def get_rollout_kwargs_for_data_parallelism(sampler_config, 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
 
@@ -542,9 +541,15 @@ def _filter_long_prompts(x):
           rollout_vllm_enable_dp_attention=trainer_config.enable_dp_attention,
           rollout_vllm_max_num_batched_tokens=trainer_config.max_num_batched_tokens,
           rollout_vllm_max_num_seqs=trainer_config.max_num_seqs,
+          rollout_vllm_async_scheduling=trainer_config.async_scheduling,
           rollout_vllm_kwargs={
               "hf_overrides": trainer_config.vllm_hf_overrides,
           },
+          rollout_vllm_sampling_kwargs={
+              "stop": trainer_config.stop_strings,
+              "detokenize": trainer_config.stop_strings is not None,
+              "include_stop_str_in_output": trainer_config.stop_strings is not None,
+          },
           **get_rollout_kwargs_for_data_parallelism(sampler_config, len(sampler_devices)),
       ),
   )