diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 761a79f0bf4..c5483feec41 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -22,24 +22,81 @@ from ..import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available +from ..trainer.grpo_config import GRPOConfig if is_requests_available(): import requests from requests import ConnectionError - if is_vllm_available(): from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.utils import StatelessProcessGroup - + from vllm import SamplingParams, LLM + from vllm.sampling_params import GuidedDecodingParams if is_vllm_ascend_available(): from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator +from accelerate import Accelerator +from accelerate.utils import broadcast_object_list, gather_object + logger = logging.getLogger(__name__) +class VLLMNoOpClient: + """ + A no-op vLLM client used in distributed training when the process is neither the main process + nor running in vLLM colocation mode. -class VLLMClient: + This stub client ensures compatibility in distributed setups without performing actual + inference or model updates. + + Methods like `generate` and `update_named_param` are implemented as no-ops or return default + values to maintain consistent interfaces across processes. + + This class should only be used internally by `get_vllm_client`. + """ + + def __init__(self, process_index: int): + self.process_index = process_index + + def generate( + self, + prompts: list[str], + n: int = 1, + repetition_penalty: float = 1.0, + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + max_tokens: int = 16, + guided_decoding_regex: Optional[str] = None, + ) -> list[list[str]]: + orig_size = len(prompts) + prompts = gather_object(prompts) + completion_ids = [None] * len(prompts) + return self._broadcast_and_slice(completion_ids, orig_size) + + def update_named_param(self, name: str, weights: torch.Tensor): + pass + + def reset_prefix_cache(self): + pass + + def _gather(self, prompts): + return gather_object(prompts) + + def _broadcast_and_slice(self, completion_ids: list, slice_size: int): + # Broadcast the completions from the main process to all processes, ensuring each process receives its + # corresponding slice + + completion_ids = broadcast_object_list(completion_ids, from_process=0) + process_slice = slice( + self.process_index * slice_size, + (self.process_index + 1) * slice_size, + ) + return completion_ids[process_slice] + +class VLLMClient(VLLMNoOpClient): """ A client class to interact with a vLLM server. @@ -84,13 +141,17 @@ class VLLMClient: """ def __init__( - self, host: str = "0.0.0.0", server_port: int = 8000, group_port: int = 51216, connection_timeout: float = 0.0 + self, host: str = "0.0.0.0", server_port: int = 8000, group_port: int = 51216, connection_timeout: float = 0.0, + distributed: bool = False ): + super().__init__(process_index=0) + if not is_requests_available(): raise ImportError("requests is not installed. Please install it with `pip install requests`.") if not is_vllm_available(): raise ImportError("vLLM is not installed. Please install it with `pip install vllm`.") + self.distributed = distributed self.session = requests.Session() self.host = host self.server_port = server_port @@ -170,6 +231,16 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ + + if self.distributed: + orig_size = len(prompts) + prompts = self._gather(prompts) + + # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate + # num_generations outputs for each one. This is faster than generating outputs for each duplicate + # prompt individually + prompts = prompts[::n] + url = f"http://{self.host}:{self.server_port}/generate/" response = self.session.post( url, @@ -186,10 +257,15 @@ def generate( }, ) if response.status_code == 200: - return response.json()["completion_ids"] + completion_ids = response.json()["completion_ids"] else: raise Exception(f"Request failed: {response.status_code}, {response.text}") + if self.distributed: + completion_ids = self._broadcast_and_slice(completion_ids, orig_size) + + return completion_ids + def init_communicator(self): """ Initializes the weight update group in a distributed setup for model synchronization. @@ -280,8 +356,145 @@ def close_communicator(self): if response.status_code != 200: raise Exception(f"Request failed: {response.status_code}, {response.text}") +class VLLMColocationClient: + """ + A client class to interact with vLLM processes colocated with the training process. + + This client bypasses remote communication and directly interacts with the in-process vLLM engine. + It supports weight updates and text generation functionalities similar to `VLLMClient`, but is optimized + for scenarios where vLLM is running in the same process or node as training. + + Args: + args (`GRPOConfig`): Configuration object containing vLLM parameters. + model (`transformers.PreTrainedModel`): The model being used. + vllm_device (`torch.device` or `str`): Device on which the model is loaded (e.g., "cuda:0"). + """ + + def __init__(self, args: GRPOConfig, model, vllm_device): + self.args: GRPOConfig = args + self.model = model + self.vllm_device = vllm_device + + self.llm = LLM( + model=self.model.name_or_path, + device=self.vllm_device, + gpu_memory_utilization=self.args.vllm_gpu_memory_utilization, + dtype=self.args.vllm_dtype, + max_model_len=self.args.vllm_max_model_len, + distributed_executor_backend="external_launcher", + seed=0 + ) + + def update_named_param(self, name: str, weights: torch.Tensor): + """ + Updates a specific named parameter in the model. + + Args: + name (`str`): + Name of the layer whose weights are being updated. + weights (`torch.Tensor`): + Tensor containing the updated weights. + """ + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name,weights)]) + + def generate( + self, + prompts: list[str], + n: int = 1, + repetition_penalty: float = 1.0, + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + max_tokens: int = 16, + guided_decoding_regex: Optional[str] = None, + ) -> list[list[str]]: + """ + Generates model completions for the provided prompts. + + Args: + prompts (`list[str]`): + List of text prompts for which the model will generate completions. + n (`int`, *optional*, defaults to `1`): + Number of completions to generate for each prompt. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Parameter for repetition penalty. 1.0 means no penalty. + temperature (`float`, *optional*, defaults to `1.0`): + Temperature parameter for sampling. Higher values increase diversity. + top_p (`float`, *optional*, defaults to `1.0`): + Top-p sampling parameter.`1.0` means no truncation. + top_k (`int`, *optional*, defaults to `-1`): + Top-k sampling parameter. `-1` means no truncation. + min_p (`float`, *optional*, defaults to `0.0`): + Minimum probability for sampling. + max_tokens (`int`, *optional*, defaults to `16`): + Maximum number of tokens to generate for each prompt. + guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`): + Regular expression to guide the decoding process. + + Returns: + `list[list[int]]`: + List of lists of token IDs representing the model-generated completions for each prompt. + """ + # Guided decoding, if enabled + if guided_decoding_regex is not None: + guided_decoding = GuidedDecodingParams(backend="outlines", regex=guided_decoding_regex) + else: + guided_decoding = None + + sampling_params = SamplingParams( + n=1, # vLLM on each GPU generates only 1 in vllm_colocation mode + repetition_penalty=repetition_penalty, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + max_tokens=max_tokens, + guided_decoding=guided_decoding, + ) + + all_outputs = self.llm.generate( + prompts, sampling_params=sampling_params, use_tqdm=False + ) + completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] + return completion_ids -# Example usage + def reset_prefix_cache(self): + """ + Resets the prefix cache for the model. + """ + self.llm.reset_prefix_cache() + +def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: + """ + Returns the appropriate vLLM client based on the current configuration. + + This function acts as a proxy to initialize and return the correct vLLM client type: + - If colocation is enabled, it returns `VLLMColocationClient`, which interacts directly with + the colocated vLLM process for faster integration. + - If running in the main process (non-colocated mode), it returns `VLLMClient`, which communicates + with an external vLLM server. + - If not the main process and colocation is disabled, it returns a base client (`VLLMNoOpClient`) + for compatibility in distributed settings. + + Args: + args (`GRPOConfig`): Configuration object containing flags for colocation, server host, port, etc. + model (`transformers.PreTrainedModel`): The model to use, passed only for the colocated client. + accelerator (`Accelerator`): Hugging Face `Accelerator` object that helps with multi-GPU training. + """ + if args.vllm_colocation: + return VLLMColocationClient(args, model, accelerator.device) + elif accelerator.is_main_process: + vllm_client = VLLMClient( + args.vllm_server_host, args.vllm_server_port, connection_timeout=args.vllm_server_timeout, + distributed=accelerator.num_processes > 1, + ) + vllm_client.init_communicator() + return vllm_client + return VLLMNoOpClient(accelerator.process_index) + +# Example usage for VLLMCLient if __name__ == "__main__": from vllm import SamplingParams @@ -297,3 +510,4 @@ def close_communicator(self): model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B").to("cuda") client.update_model_params(model) + diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index bc605f0ce3f..21912daee47 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -96,6 +96,10 @@ class GRPOConfig(TrainingArguments): timeout, a `ConnectionError` is raised. vllm_guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`): Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled. + vllm_colocation (`bool`, *optional*, defaults to `False`): + Whether to use colocated vLLM execution via external launcher. If set to `True`, vLLM will be + initialized in **all processes**, each assigned to its respective device. This allows multi-GPU + or multi-node execution with vLLM's external launcher, enabling improved large-scale inference. > Parameters that control the training @@ -295,6 +299,14 @@ class GRPOConfig(TrainingArguments): default=None, metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."}, ) + vllm_colocation: Optional[bool] = field( + default=False, + metadata={ + "help": "Whether to use colocated vLLM execution via external launcher. If set to `True`, vLLM will be " + "initialized in all processes, each assigned to its respective device. This enables optimized " + "multi-GPU inference." + }, + ) # Parameters that control the training learning_rate: float = field( @@ -412,3 +424,26 @@ class GRPOConfig(TrainingArguments): "all prompts are logged." }, ) + vllm_gpu_memory_utilization: Optional[float] = field( + default=0.3, + metadata={ + "help": "This parameter is deprecated and will be removed in version 0.18.0. To control the GPU memory " + "utilization for vLLM, you should now use the `gpu_memory_utilization` parameter in the vLLM server " + "configuration." + }, + ) + vllm_dtype: Optional[str] = field( + default="auto", + metadata={ + "help": "This parameter is deprecated and will be removed in version 0.18.0. To control the data type for " + "vLLM generation, you should now use the `dtype` parameter in the vLLM server configuration." + }, + ) + vllm_max_model_len: Optional[int] = field( + default=None, + metadata={ + "help": "This parameter is deprecated and will be removed in version 0.18.0. To control the " + "`max_model_len` for vLLM, you should now use the `max_model_len` parameter in the vLLM server " + "configuration." + }, + ) \ No newline at end of file diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index f349b80a59a..01d0c06928e 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -46,7 +46,8 @@ from ..data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template from ..extras.profiling import profiling_context, profiling_decorator -from ..extras.vllm_client import VLLMClient + +from ..extras.vllm_client import get_vllm_client from ..import_utils import is_liger_kernel_available, is_rich_available, is_vllm_available from ..models import create_reference_model, prepare_deepspeed, unwrap_model_for_generation from .callbacks import SyncRefModelCallback @@ -619,11 +620,9 @@ def data_collator(features): # No data collation is needed in GRPO "`pip install vllm` to use it." ) - if self.accelerator.is_main_process: - self.vllm_client = VLLMClient( - args.vllm_server_host, args.vllm_server_port, connection_timeout=args.vllm_server_timeout - ) - self.vllm_client.init_communicator() + self.vllm_client = get_vllm_client( + self.args, model, self.accelerator, + ) # vLLM specific sampling arguments self.guided_decoding_regex = args.vllm_guided_decoding_regex @@ -851,8 +850,7 @@ def _move_model_to_vllm(self): continue name = name.replace("modules_to_save.default.", "") - if self.accelerator.is_main_process: - self.vllm_client.update_named_param(name, param.data) + self.vllm_client.update_named_param(name, param.data) # Unmerge adapters while parameters are still gathered self.model.unmerge_adapter() @@ -861,12 +859,10 @@ def _move_model_to_vllm(self): # For non-PEFT models, simply gather and update each parameter individually. for name, param in self.model.named_parameters(): with gather_if_zero3([param]): - if self.accelerator.is_main_process: - self.vllm_client.update_named_param(name, param.data) + self.vllm_client.update_named_param(name, param.data) - # Reset cache on main process - if self.accelerator.is_main_process: - self.vllm_client.reset_prefix_cache() + # Reset cache on main process (if colocated, reset cache on all vllms) + self.vllm_client.reset_prefix_cache() @profiling_decorator def _prepare_inputs( @@ -925,35 +921,18 @@ def _generate_and_score_completions( self._move_model_to_vllm() self._last_loaded_step = self.state.global_step - # Generate completions using vLLM: gather all prompts and use them in a single call in the main process - all_prompts_text = gather_object(prompts_text) - if self.accelerator.is_main_process: - # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate - # num_generations outputs for each one. This is faster than generating outputs for each duplicate - # prompt individually. - ordered_set_of_prompts = all_prompts_text[:: self.num_generations] - with profiling_context(self, "vLLM.generate"): - completion_ids = self.vllm_client.generate( - prompts=ordered_set_of_prompts, - n=self.num_generations, - repetition_penalty=self.repetition_penalty, - temperature=self.temperature, - top_p=self.top_p, - top_k=-1 if self.top_k is None else self.top_k, - min_p=0.0 if self.min_p is None else self.min_p, - max_tokens=self.max_completion_length, - guided_decoding_regex=self.guided_decoding_regex, - ) - else: - completion_ids = [None] * len(all_prompts_text) - # Broadcast the completions from the main process to all processes, ensuring each process receives its - # corresponding slice. - completion_ids = broadcast_object_list(completion_ids, from_process=0) - process_slice = slice( - self.accelerator.process_index * len(prompts), - (self.accelerator.process_index + 1) * len(prompts), + # Generate completions using vLLM: gather all prompts and use them in a single call in the main process (if colocated, work on your own batch) + completion_ids = self.vllm_client.generate( + prompts=prompts_text, + n=self.num_generations, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + top_p=self.top_p, + top_k=-1 if self.top_k is None else self.top_k, + min_p=0.0 if self.min_p is None else self.min_p, + max_tokens=self.max_completion_length, + guided_decoding_regex=self.guided_decoding_regex, ) - completion_ids = completion_ids[process_slice] # Pad the completions, and concatenate them with the prompts completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids]