Merge pull request #3020 from AI-Hypercomputer:nicogrande/update-train-rl-args

PiperOrigin-RevId: 862978391

Author: Google-ML-Automation

src/MaxText/configs/rl.yml

@@ -144,6 +144,11 @@ swap_space_vllm_gb: 2
 decode_sampling_temperature: 0.9
 decode_sampling_top_k: 50
 decode_sampling_nucleus_p: 1.0
+# Optional sharding configuration for samplers
+enable_dp_attention: False
+# Performance tuning for samplers
+max_num_batched_tokens: null
+max_num_seqs: null
 
 # ====== Checkpoint Configuration ======
 enable_checkpointing: True

src/MaxText/configs/types.py

@@ -1485,6 +1485,9 @@ class VLLM(BaseModel):
   kv_cache_buffer: int = Field(256, description="Buffer for KV cache.")
   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.")
+  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.")
   vllm_additional_config: dict[str, Any] = Field(default_factory=dict, description="Additional vLLM config options.")
   vllm_hf_config_path: str = Field("", description="Path to HuggingFace model config for MaxText model.")
 

src/MaxText/rl/train_rl.py

@@ -437,6 +437,9 @@ def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
           rollout_vllm_hf_config_path=trainer_config.vllm_hf_config_path,
           rollout_vllm_additional_config=rollout_additional_config,
           rollout_vllm_init_with_random_weights=True,
+          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,
           **get_rollout_kwargs_for_data_parallelism(sampler_config, len(sampler_devices)),
       ),
   )