From 5293d07bd283136af50a4c0e2ad9ce75b9d7f82a Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 1 Apr 2025 14:58:57 -0400 Subject: [PATCH 01/43] Introduce TP in coloc mode --- trl/extras/vllm_client.py | 38 +++++++++++++++++++++++++++++++++----- trl/trainer/grpo_config.py | 14 ++++++++++++++ 2 files changed, 47 insertions(+), 5 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 0cbb9d2bdf6..09f872d952b 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -354,10 +354,11 @@ class VLLMColocationClient: 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): + def __init__(self, args: GRPOConfig, model, accelerator): self.args: GRPOConfig = args self.model = model - self.vllm_device = vllm_device + self.vllm_device = accelerator.device + self.tp_size = accelerator.num_processes self.llm = LLM( model=self.model.name_or_path, @@ -366,6 +367,7 @@ def __init__(self, args: GRPOConfig, model, vllm_device): dtype=self.args.vllm_dtype, enable_prefix_caching=self.args.vllm_enable_prefix_caching, max_model_len=self.args.vllm_max_model_len, + tensor_parallel_size=self.tp_size if args.vllm_tp else 1, distributed_executor_backend="external_launcher", ) @@ -382,6 +384,20 @@ def update_named_param(self, name: str, weights: torch.Tensor): llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) + 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] + def generate( self, prompts: list[str], @@ -427,8 +443,15 @@ def generate( else: guided_decoding = None + num_gen = 1 + if self.args.vllm_tp: + orig_size = len(prompts) # local size of prompts + prompts = self._gather(prompts) + prompts = prompts[::n] + num_gen = self.args.num_generations + sampling_params = SamplingParams( - n=1, # vLLM on each GPU generates only 1 in vllm_colocation mode + n=num_gen, # vLLM on each GPU generates only 1 in vllm_colocation mode, args.num_generations in vllm_tp mode repetition_penalty=repetition_penalty, temperature=temperature, top_p=top_p, @@ -441,7 +464,12 @@ def generate( 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] + + if self.args.vllm_tp: + completion_ids = self._broadcast_and_slice(completion_ids, orig_size) + return completion_ids def reset_prefix_cache(self): @@ -467,8 +495,8 @@ def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNo 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) + if args.vllm_colocation or args.vllm_tp: + return VLLMColocationClient(args, model, accelerator) elif accelerator.is_main_process: return VLLMClient( args.vllm_server_host, args.vllm_server_port, connection_timeout=args.vllm_server_timeout, diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index 23e13ed5450..9163eda0189 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -94,6 +94,10 @@ class GRPOConfig(TrainingArguments): 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. + vllm_tp (`bool`, *optional*, defaults to `False`): + Flag to enable tensor parallelism with vLLM using the external_launcher backend. + When set to True, vLLM will be initialized on all processes, with each assigned to its own device. + This enables distributed execution across multiple GPUs or nodes, allowing large-scale inference by splitting the model across devices. > Parameters that control the training @@ -262,6 +266,16 @@ class GRPOConfig(TrainingArguments): "multi-GPU inference." }, ) + vllm_tp: Optional[bool] = field( + default=False, + metadata={ + "help": ( + "Enable tensor parallel execution with vLLM using the external launcher backend. " + "When set to `True`, vLLM is initialized on all processes, each bound to its own device. " + "This allows efficient distributed inference across multiple GPUs." + ) + }, + ) # Parameters that control the training learning_rate: float = field( From 7e6245ec6edcea5b0c7d0a71945cb1ec82497d4e Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 1 Apr 2025 16:17:04 -0400 Subject: [PATCH 02/43] Introduce sleep mode in colocated vllms --- trl/extras/vllm_client.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 09f872d952b..6365c297ca6 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -369,7 +369,10 @@ def __init__(self, args: GRPOConfig, model, accelerator): max_model_len=self.args.vllm_max_model_len, tensor_parallel_size=self.tp_size if args.vllm_tp else 1, distributed_executor_backend="external_launcher", + enable_sleep_mode=True ) + + self.llm.sleep(level=2) def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -381,8 +384,10 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ + self.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) + self.llm.sleep(level=2) def _gather(self, prompts): return gather_object(prompts) @@ -437,6 +442,8 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ + + self.llm.wake_up() # Guided decoding, if enabled if guided_decoding_regex is not None: guided_decoding = GuidedDecodingParams(backend="outlines", regex=guided_decoding_regex) @@ -465,6 +472,8 @@ def generate( prompts, sampling_params=sampling_params, use_tqdm=False ) + self.llm.sleep(level=2) + completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] if self.args.vllm_tp: @@ -476,7 +485,9 @@ def reset_prefix_cache(self): """ Resets the prefix cache for the model. """ + self.llm.wake_up() self.llm.reset_prefix_cache() + self.llm.sleep(level=2) def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From 040b6e40aec25bc6f0544d5eb489f37fd18a6ae8 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 1 Apr 2025 16:49:48 -0400 Subject: [PATCH 03/43] Fix process index bug --- trl/extras/vllm_client.py | 1 + 1 file changed, 1 insertion(+) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 6365c297ca6..19e6f69e107 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -359,6 +359,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.model = model self.vllm_device = accelerator.device self.tp_size = accelerator.num_processes + self.process_index = accelerator.process_index self.llm = LLM( model=self.model.name_or_path, From 3aa300a66ad12dbee830aa533986c95bdf823640 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 1 Apr 2025 17:13:27 -0400 Subject: [PATCH 04/43] Ignore prefix cache for now --- trl/extras/vllm_client.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 19e6f69e107..3cf17a525fe 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -473,22 +473,23 @@ def generate( prompts, sampling_params=sampling_params, use_tqdm=False ) - self.llm.sleep(level=2) - completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] if self.args.vllm_tp: completion_ids = self._broadcast_and_slice(completion_ids, orig_size) + self.llm.sleep(level=2) + return completion_ids def reset_prefix_cache(self): """ Resets the prefix cache for the model. """ - self.llm.wake_up() - self.llm.reset_prefix_cache() - self.llm.sleep(level=2) + pass + # self.llm.wake_up() + # self.llm.reset_prefix_cache() + # self.llm.sleep(level=2) def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From 4e90aaa375722cc543ab3f051e4b62eef55af75a Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 1 Apr 2025 19:33:56 -0400 Subject: [PATCH 05/43] Fix sleep issues --- trl/extras/vllm_client.py | 29 +++++++++-------------------- 1 file changed, 9 insertions(+), 20 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 3cf17a525fe..c3e92db9694 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -372,8 +372,6 @@ def __init__(self, args: GRPOConfig, model, accelerator): distributed_executor_backend="external_launcher", enable_sleep_mode=True ) - - self.llm.sleep(level=2) def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -388,21 +386,9 @@ def update_named_param(self, name: str, weights: torch.Tensor): self.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - self.llm.sleep(level=2) 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] def generate( self, @@ -443,7 +429,7 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ - + torch.cuda.empty_cache() self.llm.wake_up() # Guided decoding, if enabled if guided_decoding_regex is not None: @@ -453,13 +439,11 @@ def generate( num_gen = 1 if self.args.vllm_tp: - orig_size = len(prompts) # local size of prompts + orig_size = len(prompts) # size of local prompts (for splitting later) prompts = self._gather(prompts) - prompts = prompts[::n] - num_gen = self.args.num_generations sampling_params = SamplingParams( - n=num_gen, # vLLM on each GPU generates only 1 in vllm_colocation mode, args.num_generations in vllm_tp mode + n=num_gen, # vLLM on each GPU generates only 1 in vllm_colocation mode, args.num_generations (or 1?) in vllm_tp mode repetition_penalty=repetition_penalty, temperature=temperature, top_p=top_p, @@ -476,7 +460,12 @@ def generate( completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] if self.args.vllm_tp: - completion_ids = self._broadcast_and_slice(completion_ids, orig_size) + # just do split - no broadcast! + tp_slice = slice( + self.process_index * orig_size, + (self.process_index) * orig_size + ) + completion_ids = completion_ids[tp_slice] self.llm.sleep(level=2) From 85d930998e728920471ba6ce9289518abc801cc8 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 1 Apr 2025 19:49:33 -0400 Subject: [PATCH 06/43] Fix tp slices --- trl/extras/vllm_client.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index c3e92db9694..40c73d0eabf 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -442,6 +442,8 @@ def generate( orig_size = len(prompts) # size of local prompts (for splitting later) prompts = self._gather(prompts) + print("-----\n orig size", orig_size , " prompt size", len(prompts)) + sampling_params = SamplingParams( n=num_gen, # vLLM on each GPU generates only 1 in vllm_colocation mode, args.num_generations (or 1?) in vllm_tp mode repetition_penalty=repetition_penalty, @@ -458,14 +460,16 @@ def generate( ) completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] + print("-----\n completion id size", len(completion_ids)) if self.args.vllm_tp: # just do split - no broadcast! tp_slice = slice( - self.process_index * orig_size, - (self.process_index) * orig_size + self.process_index * orig_size, + (self.process_index + 1) * orig_size ) completion_ids = completion_ids[tp_slice] + print("-----\n Sliced completion id size", len(completion_ids), " and slice: ", tp_slice) self.llm.sleep(level=2) From d2bccd36ed44b12571e7eca1d5787097e1ae4529 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 1 Apr 2025 19:56:02 -0400 Subject: [PATCH 07/43] Fix typo in wake up --- trl/extras/vllm_client.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 40c73d0eabf..5df08bbc537 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -383,7 +383,7 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ - self.wake_up() + self.llm.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) @@ -442,8 +442,6 @@ def generate( orig_size = len(prompts) # size of local prompts (for splitting later) prompts = self._gather(prompts) - print("-----\n orig size", orig_size , " prompt size", len(prompts)) - sampling_params = SamplingParams( n=num_gen, # vLLM on each GPU generates only 1 in vllm_colocation mode, args.num_generations (or 1?) in vllm_tp mode repetition_penalty=repetition_penalty, @@ -460,7 +458,6 @@ def generate( ) completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] - print("-----\n completion id size", len(completion_ids)) if self.args.vllm_tp: # just do split - no broadcast! @@ -469,7 +466,6 @@ def generate( (self.process_index + 1) * orig_size ) completion_ids = completion_ids[tp_slice] - print("-----\n Sliced completion id size", len(completion_ids), " and slice: ", tp_slice) self.llm.sleep(level=2) From d876fe4d39ee743977dd1fc1dd75c9905d19a1c1 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 2 Apr 2025 08:55:00 -0400 Subject: [PATCH 08/43] Debugging memory --- trl/extras/vllm_client.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 5df08bbc537..f9f9218abf4 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -361,6 +361,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.tp_size = accelerator.num_processes self.process_index = accelerator.process_index + print("\n\n\n\n\n------------- Initializing LLM") self.llm = LLM( model=self.model.name_or_path, device=self.vllm_device, @@ -372,6 +373,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): distributed_executor_backend="external_launcher", enable_sleep_mode=True ) + print("\n\n\n\n\n------------- Initialized LLM") def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -383,9 +385,11 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ + print("\n\n\n\n\n------------- Updating model - waking up") self.llm.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) + print("\n\n\n\n\n------------- Updated the model") def _gather(self, prompts): return gather_object(prompts) @@ -429,8 +433,16 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ + print("\n\n\n\n\n------------- Generation") + mem = torch.cuda.memory_allocated() + print("\n\n\n\n\n------------- Mememory before empty_cache", mem) torch.cuda.empty_cache() + + mem = torch.cuda.memory_allocated() + print("\n\n\n\n\n------------- memory after empty_cache, and waking up", mem) + self.llm.wake_up() + print("\n\n\n\n\n------------- Woke up, will do generation") # Guided decoding, if enabled if guided_decoding_regex is not None: guided_decoding = GuidedDecodingParams(backend="outlines", regex=guided_decoding_regex) @@ -468,6 +480,7 @@ def generate( completion_ids = completion_ids[tp_slice] self.llm.sleep(level=2) + print("\n\n\n\n\n------------- GEnerated, going back to sleep") return completion_ids @@ -475,10 +488,12 @@ def reset_prefix_cache(self): """ Resets the prefix cache for the model. """ - pass - # self.llm.wake_up() - # self.llm.reset_prefix_cache() - # self.llm.sleep(level=2) + # pass + print("\n\n\n\n\n------------- Will reset prefix cache now - waking up") + self.llm.wake_up() + self.llm.reset_prefix_cache() + self.llm.sleep(level=2) + print("\n\n\n\n\n------------- Reset done, going back to sleep") def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From 712e85b1eb3db379edff6c921fc6d67f72319094 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 2 Apr 2025 09:42:20 -0400 Subject: [PATCH 09/43] Measure memory during model update --- trl/extras/vllm_client.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index f9f9218abf4..3e38df5541b 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -373,7 +373,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): distributed_executor_backend="external_launcher", enable_sleep_mode=True ) - print("\n\n\n\n\n------------- Initialized LLM") + print("------------- Initialized LLM") def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -385,11 +385,16 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ - print("\n\n\n\n\n------------- Updating model - waking up") + print(self.process_index, "------------- Updating model - waking up") + mem = torch.cuda.memory_allocated() + print(self.process_index, "------------- Mememory before empty_cache in updating model", mem) + torch.cuda.empty_cache() + mem = torch.cuda.memory_allocated() + print(self.process_index, "------------- memory after empty_cache, and waking up in updating model", mem) self.llm.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - print("\n\n\n\n\n------------- Updated the model") + print(self.process_index, "------------- Updated the model") def _gather(self, prompts): return gather_object(prompts) @@ -433,16 +438,16 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ - print("\n\n\n\n\n------------- Generation") + print(self.process_index, "------------- Generation") mem = torch.cuda.memory_allocated() - print("\n\n\n\n\n------------- Mememory before empty_cache", mem) + print("------------- Mememory before empty_cache", mem) torch.cuda.empty_cache() mem = torch.cuda.memory_allocated() - print("\n\n\n\n\n------------- memory after empty_cache, and waking up", mem) + print(self.process_index, "------------- memory after empty_cache, and waking up", mem) self.llm.wake_up() - print("\n\n\n\n\n------------- Woke up, will do generation") + print(self.process_index, "------------- Woke up, will do generation") # Guided decoding, if enabled if guided_decoding_regex is not None: guided_decoding = GuidedDecodingParams(backend="outlines", regex=guided_decoding_regex) @@ -480,7 +485,7 @@ def generate( completion_ids = completion_ids[tp_slice] self.llm.sleep(level=2) - print("\n\n\n\n\n------------- GEnerated, going back to sleep") + print(self.process_index, "------------- GEnerated, going back to sleep") return completion_ids @@ -489,11 +494,11 @@ def reset_prefix_cache(self): Resets the prefix cache for the model. """ # pass - print("\n\n\n\n\n------------- Will reset prefix cache now - waking up") + print(self.process_index, "------------- Will reset prefix cache now - waking up") self.llm.wake_up() self.llm.reset_prefix_cache() self.llm.sleep(level=2) - print("\n\n\n\n\n------------- Reset done, going back to sleep") + print(self.process_index, "------------- Reset done, going back to sleep") def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From 5127b97fe96d15751672c628ae0304859eb4b42d Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 2 Apr 2025 10:01:40 -0400 Subject: [PATCH 10/43] Debug memory reserved --- trl/extras/vllm_client.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 3e38df5541b..2d4896082f7 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -386,10 +386,10 @@ def update_named_param(self, name: str, weights: torch.Tensor): Tensor containing the updated weights. """ print(self.process_index, "------------- Updating model - waking up") - mem = torch.cuda.memory_allocated() + mem = torch.cuda.memory_reserved() print(self.process_index, "------------- Mememory before empty_cache in updating model", mem) torch.cuda.empty_cache() - mem = torch.cuda.memory_allocated() + mem = torch.cuda.memory_reserved() print(self.process_index, "------------- memory after empty_cache, and waking up in updating model", mem) self.llm.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model @@ -439,11 +439,10 @@ def generate( List of lists of token IDs representing the model-generated completions for each prompt. """ print(self.process_index, "------------- Generation") - mem = torch.cuda.memory_allocated() + mem = torch.cuda.memory_reserved() print("------------- Mememory before empty_cache", mem) torch.cuda.empty_cache() - - mem = torch.cuda.memory_allocated() + mem = torch.cuda.memory_reserved() print(self.process_index, "------------- memory after empty_cache, and waking up", mem) self.llm.wake_up() From a8a6f693f84870ee23303059e099b9fc438dab19 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 2 Apr 2025 12:38:51 -0400 Subject: [PATCH 11/43] Remove prints --- trl/extras/vllm_client.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 2d4896082f7..865ef169c7b 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -361,7 +361,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.tp_size = accelerator.num_processes self.process_index = accelerator.process_index - print("\n\n\n\n\n------------- Initializing LLM") + # print("\n\n\n\n\n------------- Initializing LLM") self.llm = LLM( model=self.model.name_or_path, device=self.vllm_device, @@ -373,7 +373,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): distributed_executor_backend="external_launcher", enable_sleep_mode=True ) - print("------------- Initialized LLM") + # print("------------- Initialized LLM") def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -385,16 +385,16 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ - print(self.process_index, "------------- Updating model - waking up") + # print(self.process_index, "------------- Updating model - waking up") mem = torch.cuda.memory_reserved() - print(self.process_index, "------------- Mememory before empty_cache in updating model", mem) + # print(self.process_index, "------------- Mememory before empty_cache in updating model", mem) torch.cuda.empty_cache() mem = torch.cuda.memory_reserved() - print(self.process_index, "------------- memory after empty_cache, and waking up in updating model", mem) + # print(self.process_index, "------------- memory after empty_cache, and waking up in updating model", mem) self.llm.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - print(self.process_index, "------------- Updated the model") + # print(self.process_index, "------------- Updated the model") def _gather(self, prompts): return gather_object(prompts) @@ -438,15 +438,15 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ - print(self.process_index, "------------- Generation") - mem = torch.cuda.memory_reserved() - print("------------- Mememory before empty_cache", mem) + # print(self.process_index, "------------- Generation") + # mem = torch.cuda.memory_reserved() + # print("------------- Mememory before empty_cache", mem) torch.cuda.empty_cache() - mem = torch.cuda.memory_reserved() - print(self.process_index, "------------- memory after empty_cache, and waking up", mem) + # mem = torch.cuda.memory_reserved() + # print(self.process_index, "------------- memory after empty_cache, and waking up", mem) self.llm.wake_up() - print(self.process_index, "------------- Woke up, will do generation") + # print(self.process_index, "------------- Woke up, will do generation") # Guided decoding, if enabled if guided_decoding_regex is not None: guided_decoding = GuidedDecodingParams(backend="outlines", regex=guided_decoding_regex) @@ -484,7 +484,7 @@ def generate( completion_ids = completion_ids[tp_slice] self.llm.sleep(level=2) - print(self.process_index, "------------- GEnerated, going back to sleep") + # print(self.process_index, "------------- GEnerated, going back to sleep") return completion_ids @@ -493,11 +493,11 @@ def reset_prefix_cache(self): Resets the prefix cache for the model. """ # pass - print(self.process_index, "------------- Will reset prefix cache now - waking up") + # print(self.process_index, "------------- Will reset prefix cache now - waking up") self.llm.wake_up() self.llm.reset_prefix_cache() self.llm.sleep(level=2) - print(self.process_index, "------------- Reset done, going back to sleep") + # print(self.process_index, "------------- Reset done, going back to sleep") def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From 5de5f3cdf77621e90d8a78d523de22bb81ddc8be Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 2 Apr 2025 13:15:51 -0400 Subject: [PATCH 12/43] Validate via experiments --- trl/extras/vllm_client.py | 24 ++---------------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 865ef169c7b..ce6505594da 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -361,7 +361,6 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.tp_size = accelerator.num_processes self.process_index = accelerator.process_index - # print("\n\n\n\n\n------------- Initializing LLM") self.llm = LLM( model=self.model.name_or_path, device=self.vllm_device, @@ -373,7 +372,6 @@ def __init__(self, args: GRPOConfig, model, accelerator): distributed_executor_backend="external_launcher", enable_sleep_mode=True ) - # print("------------- Initialized LLM") def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -385,16 +383,10 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ - # print(self.process_index, "------------- Updating model - waking up") - mem = torch.cuda.memory_reserved() - # print(self.process_index, "------------- Mememory before empty_cache in updating model", mem) torch.cuda.empty_cache() - mem = torch.cuda.memory_reserved() - # print(self.process_index, "------------- memory after empty_cache, and waking up in updating model", mem) self.llm.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - # print(self.process_index, "------------- Updated the model") def _gather(self, prompts): return gather_object(prompts) @@ -438,28 +430,20 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ - # print(self.process_index, "------------- Generation") - # mem = torch.cuda.memory_reserved() - # print("------------- Mememory before empty_cache", mem) torch.cuda.empty_cache() - # mem = torch.cuda.memory_reserved() - # print(self.process_index, "------------- memory after empty_cache, and waking up", mem) - self.llm.wake_up() - # print(self.process_index, "------------- Woke up, will do generation") # Guided decoding, if enabled if guided_decoding_regex is not None: guided_decoding = GuidedDecodingParams(backend="outlines", regex=guided_decoding_regex) else: guided_decoding = None - num_gen = 1 if self.args.vllm_tp: orig_size = len(prompts) # size of local prompts (for splitting later) prompts = self._gather(prompts) sampling_params = SamplingParams( - n=num_gen, # vLLM on each GPU generates only 1 in vllm_colocation mode, args.num_generations (or 1?) in vllm_tp mode + n=1, # vLLM on each GPU generates only 1 in vllm_colocation mode repetition_penalty=repetition_penalty, temperature=temperature, top_p=top_p, @@ -484,20 +468,16 @@ def generate( completion_ids = completion_ids[tp_slice] self.llm.sleep(level=2) - # print(self.process_index, "------------- GEnerated, going back to sleep") - return completion_ids def reset_prefix_cache(self): """ Resets the prefix cache for the model. """ - # pass - # print(self.process_index, "------------- Will reset prefix cache now - waking up") + # ToDo: perhaps we need to just pass self.llm.wake_up() self.llm.reset_prefix_cache() self.llm.sleep(level=2) - # print(self.process_index, "------------- Reset done, going back to sleep") def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From cbb5ef3560828160cf86e5fde3d0dab118455fe8 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 7 Apr 2025 11:01:34 -0400 Subject: [PATCH 13/43] Introduce flexible TP where TP may not be equal to world size --- trl/extras/vllm_client.py | 62 +++++++++++++++++++++++++++----------- trl/trainer/grpo_config.py | 32 +++++++------------- 2 files changed, 56 insertions(+), 38 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index ce6505594da..1d6f791b72a 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -14,6 +14,7 @@ import atexit import logging +import os import time from typing import Optional @@ -355,11 +356,36 @@ class VLLMColocationClient: """ def __init__(self, args: GRPOConfig, model, accelerator): - self.args: GRPOConfig = args + self.args = args self.model = model self.vllm_device = accelerator.device - self.tp_size = accelerator.num_processes + self.world_size = accelerator.num_processes self.process_index = accelerator.process_index + set_seed(42) + print(f"\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") + + if self.args.vllm_colocation_tp: + # Ensure TP value is valid (at least 1) + assert self.args.vllm_colocation_tp >= 1, "vllm_colocation_tp must be greater than 0" + + # Get local world size from environment (https://pytorch.org/docs/stable/elastic/run.html#environment-variables) + self.local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) + + # Make sure TP group size evenly divides the local world size + # This ensures each group has the same number of ranks + assert self.local_world_size % self.args.vllm_colocation_tp == 0, ( + f"TP size of vllm_colocation_tp ({self.args.vllm_colocation_tp}) must divide LOCAL_WORLD_SIZE " + f"({self.local_world_size}) evenly." + ) + + # Create subgroups of ranks for TP, each group with `vllm_colocation_tp` ranks. + # For example, if world_size=8 and vllm_colocation_tp=2 → groups: [0,1], [2,3], [4,5], [6,7] + self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration( + [ + list(range(i*self.args.vllm_colocation_tp, (i+1) * self.args.vllm_colocation_tp)) + for i in range(self.world_size // self.args.vllm_colocation_tp) + ] + ) self.llm = LLM( model=self.model.name_or_path, @@ -368,7 +394,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): dtype=self.args.vllm_dtype, enable_prefix_caching=self.args.vllm_enable_prefix_caching, max_model_len=self.args.vllm_max_model_len, - tensor_parallel_size=self.tp_size if args.vllm_tp else 1, + tensor_parallel_size=args.vllm_colocation_tp, distributed_executor_backend="external_launcher", enable_sleep_mode=True ) @@ -388,9 +414,6 @@ def update_named_param(self, name: str, weights: torch.Tensor): llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - def _gather(self, prompts): - return gather_object(prompts) - def generate( self, prompts: list[str], @@ -438,12 +461,18 @@ def generate( else: guided_decoding = None - if self.args.vllm_tp: - orig_size = len(prompts) # size of local prompts (for splitting later) - prompts = self._gather(prompts) + if self.args.vllm_colocation_tp: + # Gather prompts from all ranks in the TP group and flatten. + # Each rank starts with its own prompts; after gathering, all ranks see the full group set. + orig_size = len(prompts) + gathered_prompts = [None for _ in range(self.args.vllm_colocation_tp)] + torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) + prompts = [p for sublist in gathered_prompts for p in sublist] + + print("\n\n---Rank ", self.process_index, " colocation check prompts, orig_size: ", orig_size, " local group prompts size", len(prompts), " should be equal to ", orig_size*self.args.vllm_colocation_tp ) sampling_params = SamplingParams( - n=1, # vLLM on each GPU generates only 1 in vllm_colocation mode + n=1, # vLLM on each device generates only 1 in vllm_colocation mode repetition_penalty=repetition_penalty, temperature=temperature, top_p=top_p, @@ -459,12 +488,11 @@ def generate( completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] - if self.args.vllm_tp: - # just do split - no broadcast! - tp_slice = slice( - self.process_index * orig_size, - (self.process_index + 1) * orig_size - ) + if self.args.vllm_colocation_tp: + # Slice completions for this rank within its TP group. + # Each rank generates all outputs — we keep only our share. + local_rank_in_group = torch.distributed.get_rank(group=self.tp_group) + tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) completion_ids = completion_ids[tp_slice] self.llm.sleep(level=2) @@ -496,7 +524,7 @@ def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNo 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 or args.vllm_tp: + if args.vllm_colocation_tp: return VLLMColocationClient(args, model, accelerator) elif accelerator.is_main_process: return VLLMClient( diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index 9163eda0189..8928a08f70d 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -90,14 +90,11 @@ 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. - vllm_tp (`bool`, *optional*, defaults to `False`): - Flag to enable tensor parallelism with vLLM using the external_launcher backend. - When set to True, vLLM will be initialized on all processes, with each assigned to its own device. - This enables distributed execution across multiple GPUs or nodes, allowing large-scale inference by splitting the model across devices. + vllm_colocation_tp (`int` or `None`, *optional*, defaults to `None`): + Controls colocated vLLM execution and tensor parallelism via the `external_launcher` backend. + - Set to `None` to disable colocated vLLM entirely. + - Set to `1` to enable colocated vLLM on a single GPU with no tensor parallelism. + - Set to a value >1 to enable colocated vLLM with tensor parallelism across multiple GPUs or nodes. > Parameters that control the training @@ -258,21 +255,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." - }, - ) - vllm_tp: Optional[bool] = field( - default=False, + vllm_colocation_tp: Optional[int] = field( + default=None, metadata={ "help": ( - "Enable tensor parallel execution with vLLM using the external launcher backend. " - "When set to `True`, vLLM is initialized on all processes, each bound to its own device. " - "This allows efficient distributed inference across multiple GPUs." + "Controls colocated vLLM execution and tensor parallelism using the `external_launcher` backend. " + "Set to `None` to disable colocated vLLM. " + "Set to `1` to enable colocated vLLM on a single device (no tensor parallelism). " + "Set to a value >1 to enable colocated vLLM with tensor parallelism across multiple devices." ) }, ) From c6b36e70afc8d02ea084b779eed2057cf6b98c03 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 7 Apr 2025 12:07:01 -0400 Subject: [PATCH 14/43] Remove prints --- trl/extras/vllm_client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 1d6f791b72a..abd87789759 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -362,7 +362,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.world_size = accelerator.num_processes self.process_index = accelerator.process_index set_seed(42) - print(f"\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") + # print(f"\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") if self.args.vllm_colocation_tp: # Ensure TP value is valid (at least 1) @@ -469,7 +469,7 @@ def generate( torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) prompts = [p for sublist in gathered_prompts for p in sublist] - print("\n\n---Rank ", self.process_index, " colocation check prompts, orig_size: ", orig_size, " local group prompts size", len(prompts), " should be equal to ", orig_size*self.args.vllm_colocation_tp ) + # print("\n\n---Rank ", self.process_index, " colocation check prompts, orig_size: ", orig_size, " local group prompts size", len(prompts), " should be equal to ", orig_size*self.args.vllm_colocation_tp ) sampling_params = SamplingParams( n=1, # vLLM on each device generates only 1 in vllm_colocation mode From cc3023c59fe6183fe460d3324482561fe4a58356 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 7 Apr 2025 13:25:06 -0400 Subject: [PATCH 15/43] Dont wakeup in pfix reset --- trl/extras/vllm_client.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index abd87789759..566b99a414a 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -362,7 +362,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.world_size = accelerator.num_processes self.process_index = accelerator.process_index set_seed(42) - # print(f"\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") + print(f"\n\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") if self.args.vllm_colocation_tp: # Ensure TP value is valid (at least 1) @@ -469,7 +469,9 @@ def generate( torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) prompts = [p for sublist in gathered_prompts for p in sublist] - # print("\n\n---Rank ", self.process_index, " colocation check prompts, orig_size: ", orig_size, " local group prompts size", len(prompts), " should be equal to ", orig_size*self.args.vllm_colocation_tp ) + print(f"\n\n---Rank {self.process_index} colocation check prompts, " + f"orig_size: {orig_size}, local group prompts size: {len(prompts)}, " + f"should be equal to: {orig_size * self.args.vllm_colocation_tp}") sampling_params = SamplingParams( n=1, # vLLM on each device generates only 1 in vllm_colocation mode @@ -503,9 +505,8 @@ def reset_prefix_cache(self): Resets the prefix cache for the model. """ # ToDo: perhaps we need to just pass - self.llm.wake_up() + # no need to wake up already awake) self.llm.reset_prefix_cache() - self.llm.sleep(level=2) def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From 8ae9f01a4b7c0ddbe3559ad858b285e64d9f217d Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 7 Apr 2025 18:47:08 -0400 Subject: [PATCH 16/43] Tp size should divide global world size evenly --- trl/extras/vllm_client.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 566b99a414a..3411df8eb06 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -368,14 +368,11 @@ def __init__(self, args: GRPOConfig, model, accelerator): # Ensure TP value is valid (at least 1) assert self.args.vllm_colocation_tp >= 1, "vllm_colocation_tp must be greater than 0" - # Get local world size from environment (https://pytorch.org/docs/stable/elastic/run.html#environment-variables) - self.local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) - - # Make sure TP group size evenly divides the local world size + # Make sure TP group size evenly divides the world size # This ensures each group has the same number of ranks - assert self.local_world_size % self.args.vllm_colocation_tp == 0, ( - f"TP size of vllm_colocation_tp ({self.args.vllm_colocation_tp}) must divide LOCAL_WORLD_SIZE " - f"({self.local_world_size}) evenly." + assert self.world_size % self.args.vllm_colocation_tp == 0, ( + f"TP size of vllm_colocation_tp ({self.args.vllm_colocation_tp}) must divide world size " + f"({self.world_size}) evenly." ) # Create subgroups of ranks for TP, each group with `vllm_colocation_tp` ranks. From 2de090f99a2408260646a0bbcb6d69ad3d58a91d Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 7 Apr 2025 20:20:41 -0400 Subject: [PATCH 17/43] Add max num seq --- trl/extras/vllm_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 3411df8eb06..2694ceb8d51 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -393,7 +393,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): max_model_len=self.args.vllm_max_model_len, tensor_parallel_size=args.vllm_colocation_tp, distributed_executor_backend="external_launcher", - enable_sleep_mode=True + max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation_tp ) def update_named_param(self, name: str, weights: torch.Tensor): From 8151e1168d74f064efe8aa19f76a26a8a5d04058 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 7 Apr 2025 21:14:19 -0400 Subject: [PATCH 18/43] Bring back sleep --- trl/extras/vllm_client.py | 1 + 1 file changed, 1 insertion(+) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 2694ceb8d51..5ac99e0ba62 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -393,6 +393,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): max_model_len=self.args.vllm_max_model_len, tensor_parallel_size=args.vllm_colocation_tp, distributed_executor_backend="external_launcher", + enable_sleep_mode=True, max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation_tp ) From 946d49d573681162c3d0bbd0cceafa2fed04ae9d Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 8 Apr 2025 11:43:14 -0400 Subject: [PATCH 19/43] Fix sleep bug for grad accumulations --- trl/extras/vllm_client.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 5ac99e0ba62..5b9de685327 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -361,6 +361,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.vllm_device = accelerator.device self.world_size = accelerator.num_processes self.process_index = accelerator.process_index + self._step = 0 set_seed(42) print(f"\n\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") @@ -452,7 +453,6 @@ def generate( List of lists of token IDs representing the model-generated completions for each prompt. """ torch.cuda.empty_cache() - self.llm.wake_up() # Guided decoding, if enabled if guided_decoding_regex is not None: guided_decoding = GuidedDecodingParams(backend="outlines", regex=guided_decoding_regex) @@ -467,9 +467,9 @@ def generate( torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) prompts = [p for sublist in gathered_prompts for p in sublist] - print(f"\n\n---Rank {self.process_index} colocation check prompts, " + print(f"\n\n---Rank {self.process_index} colocation, step {self._step}, grad accumulation {self.args.gradient_accumulation_steps}, check prompts, " f"orig_size: {orig_size}, local group prompts size: {len(prompts)}, " - f"should be equal to: {orig_size * self.args.vllm_colocation_tp}") + f"should be equal to: {orig_size * self.args.vllm_colocation_tp}") if self.process_index == 0 else None sampling_params = SamplingParams( n=1, # vLLM on each device generates only 1 in vllm_colocation mode @@ -495,7 +495,13 @@ def generate( tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) completion_ids = completion_ids[tp_slice] - self.llm.sleep(level=2) + # Only sleep after the last mini-step + if self.args.gradient_accumulation_steps == 1 or \ + ((self._step + 1) % self.args.gradient_accumulation_steps == 0): + print(f"\n\n---Rank {self.process_index} sleeping now") if self.process_index == 0 else None + self.llm.sleep(level=2) + + self._step += 1 return completion_ids def reset_prefix_cache(self): @@ -503,7 +509,6 @@ def reset_prefix_cache(self): Resets the prefix cache for the model. """ # ToDo: perhaps we need to just pass - # no need to wake up already awake) self.llm.reset_prefix_cache() def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: From fe8f68493c64d456984c919f9c393b2459d0109a Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 8 Apr 2025 12:18:24 -0400 Subject: [PATCH 20/43] Reload model during grad accumulation --- trl/extras/vllm_client.py | 34 +++++++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 5b9de685327..49ef32aa629 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -37,6 +37,10 @@ from accelerate import Accelerator from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed +from contextlib import nullcontext +from ..import_utils import is_deepspeed_available +if is_deepspeed_available(): + import deepspeed logger = logging.getLogger(__name__) @@ -358,10 +362,11 @@ class VLLMColocationClient: def __init__(self, args: GRPOConfig, model, accelerator): self.args = args self.model = model + self.accelerator = accelerator self.vllm_device = accelerator.device self.world_size = accelerator.num_processes self.process_index = accelerator.process_index - self._step = 0 + self._model_needs_loading = True set_seed(42) print(f"\n\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") @@ -412,6 +417,19 @@ def update_named_param(self, name: str, weights: torch.Tensor): self.llm.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) + self._model_needs_loading = False + + def load_model_during_grad_accumulation(self): + # Only load model during the grad accumulation steps - otherwise model was just loaded + if self._model_needs_loading: + print(f"\n\n---Rank {self.process_index} updating the model now") if self.process_index == 0 else None + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 + gather_if_zero3 = deepspeed.zero.GatheredParameters if zero_stage_3 else nullcontext + # ToDo: For now, focused on non-PEFT models, simply gather and update each parameter individually. + for name, param in self.model.named_parameters(): + with gather_if_zero3([param]): + self.update_named_param(name, param.data) def generate( self, @@ -453,6 +471,9 @@ def generate( List of lists of token IDs representing the model-generated completions for each prompt. """ torch.cuda.empty_cache() + # Load model during grad accumulation steps (we lost model weights during previous sleep for memory) + self.load_model_during_grad_accumulation() + # Guided decoding, if enabled if guided_decoding_regex is not None: guided_decoding = GuidedDecodingParams(backend="outlines", regex=guided_decoding_regex) @@ -467,7 +488,7 @@ def generate( torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) prompts = [p for sublist in gathered_prompts for p in sublist] - print(f"\n\n---Rank {self.process_index} colocation, step {self._step}, grad accumulation {self.args.gradient_accumulation_steps}, check prompts, " + print(f"\n\n---Rank {self.process_index} generation... check prompts, " f"orig_size: {orig_size}, local group prompts size: {len(prompts)}, " f"should be equal to: {orig_size * self.args.vllm_colocation_tp}") if self.process_index == 0 else None @@ -495,13 +516,8 @@ def generate( tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) completion_ids = completion_ids[tp_slice] - # Only sleep after the last mini-step - if self.args.gradient_accumulation_steps == 1 or \ - ((self._step + 1) % self.args.gradient_accumulation_steps == 0): - print(f"\n\n---Rank {self.process_index} sleeping now") if self.process_index == 0 else None - self.llm.sleep(level=2) - - self._step += 1 + self.llm.sleep(level=2) + self._model_needs_loading = True # we lose the weights after sleep - so ensure to reload it before next generation return completion_ids def reset_prefix_cache(self): From c3509de4616866645770f37a3d019d98690a15cd Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 8 Apr 2025 13:13:29 -0400 Subject: [PATCH 21/43] Switch to sleep level 1 --- trl/extras/vllm_client.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 49ef32aa629..093456db1a2 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -472,7 +472,8 @@ def generate( """ torch.cuda.empty_cache() # Load model during grad accumulation steps (we lost model weights during previous sleep for memory) - self.load_model_during_grad_accumulation() + # self.load_model_during_grad_accumulation() + self.llm.wake_up() # Guided decoding, if enabled if guided_decoding_regex is not None: @@ -516,7 +517,7 @@ def generate( tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) completion_ids = completion_ids[tp_slice] - self.llm.sleep(level=2) + self.llm.sleep(level=1) self._model_needs_loading = True # we lose the weights after sleep - so ensure to reload it before next generation return completion_ids From aca6242465ce1ed7c41289431520eb300b90a92f Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 8 Apr 2025 13:50:27 -0400 Subject: [PATCH 22/43] Sleep 1 during acc steps and levl 2 otherwise --- trl/extras/vllm_client.py | 34 ++++++++++++++++------------------ 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 093456db1a2..ee89e73ba70 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -366,7 +366,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.vllm_device = accelerator.device self.world_size = accelerator.num_processes self.process_index = accelerator.process_index - self._model_needs_loading = True + self._accumulation_step = 0 set_seed(42) print(f"\n\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") @@ -417,19 +417,9 @@ def update_named_param(self, name: str, weights: torch.Tensor): self.llm.wake_up() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - self._model_needs_loading = False - - def load_model_during_grad_accumulation(self): - # Only load model during the grad accumulation steps - otherwise model was just loaded - if self._model_needs_loading: - print(f"\n\n---Rank {self.process_index} updating the model now") if self.process_index == 0 else None - deepspeed_plugin = self.accelerator.state.deepspeed_plugin - zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 - gather_if_zero3 = deepspeed.zero.GatheredParameters if zero_stage_3 else nullcontext - # ToDo: For now, focused on non-PEFT models, simply gather and update each parameter individually. - for name, param in self.model.named_parameters(): - with gather_if_zero3([param]): - self.update_named_param(name, param.data) + self._accumulation_step = 0 + if self.process_index == 0: + print(f"[RANK 0] update_named_param done: {name}, accumulation_step reset to 0.") def generate( self, @@ -471,8 +461,6 @@ def generate( List of lists of token IDs representing the model-generated completions for each prompt. """ torch.cuda.empty_cache() - # Load model during grad accumulation steps (we lost model weights during previous sleep for memory) - # self.load_model_during_grad_accumulation() self.llm.wake_up() # Guided decoding, if enabled @@ -517,8 +505,16 @@ def generate( tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) completion_ids = completion_ids[tp_slice] - self.llm.sleep(level=1) - self._model_needs_loading = True # we lose the weights after sleep - so ensure to reload it before next generation + # during grad accumulation, sleep 1, otherwise sleep 2 as we can safely forget weights - new model will be loaded + is_grad_accum = ((self._accumulation_step + 1) % self.args.gradient_accumulation_steps) != 0 + if is_grad_accum: + self.llm.sleep(level=1) + else: + self.llm.sleep(level=2) + self._accumulation_step += 1 + + if self.process_index == 0: + print(f"[RANK 0] generate done. Updated accumulation_step to {self._accumulation_step} and is_grad_accum was {is_grad_accum}") return completion_ids def reset_prefix_cache(self): @@ -527,6 +523,8 @@ def reset_prefix_cache(self): """ # ToDo: perhaps we need to just pass self.llm.reset_prefix_cache() + if self.process_index == 0: + print(f"[RANK 0] reset_prefix_cache done.") def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From 110cbce4abbdd878a24233fad16ab5f83ea341d2 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 9 Apr 2025 08:51:08 -0400 Subject: [PATCH 23/43] Fix config dfefinition --- trl/trainer/grpo_config.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index 8928a08f70d..2761906ee8a 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -93,8 +93,8 @@ class GRPOConfig(TrainingArguments): vllm_colocation_tp (`int` or `None`, *optional*, defaults to `None`): Controls colocated vLLM execution and tensor parallelism via the `external_launcher` backend. - Set to `None` to disable colocated vLLM entirely. - - Set to `1` to enable colocated vLLM on a single GPU with no tensor parallelism. - - Set to a value >1 to enable colocated vLLM with tensor parallelism across multiple GPUs or nodes. + - Set to `1` to enable colocated vLLM on each GPU with no tensor parallelism. + - Set to a value >1 to enable colocated vLLM with tensor parallelism across multiple GPUs. > Parameters that control the training @@ -261,7 +261,7 @@ class GRPOConfig(TrainingArguments): "help": ( "Controls colocated vLLM execution and tensor parallelism using the `external_launcher` backend. " "Set to `None` to disable colocated vLLM. " - "Set to `1` to enable colocated vLLM on a single device (no tensor parallelism). " + "Set to `1` to enable colocated vLLM on each device (no tensor parallelism). " "Set to a value >1 to enable colocated vLLM with tensor parallelism across multiple devices." ) }, From 9e5128a54647997a94a8fd69ab684628150691bb Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 9 Apr 2025 09:11:05 -0400 Subject: [PATCH 24/43] Debug generations --- trl/extras/vllm_client.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index ee89e73ba70..3c59e780711 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -496,6 +496,15 @@ def generate( prompts, sampling_params=sampling_params, use_tqdm=False ) + if self.process_index == 0: + print("\n\n==== Rank 0 Prompt/Generation Output ====\n") + for i, (prompt, outputs) in enumerate(zip(prompts, all_outputs)): + print(f"--- Prompt {i+1} ---") + print(prompt) + print(f"--- Generation {i+1} ---") + print(outputs.outputs[0].text.strip()) + print("=" * 40) + completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] if self.args.vllm_colocation_tp: @@ -514,14 +523,13 @@ def generate( self._accumulation_step += 1 if self.process_index == 0: - print(f"[RANK 0] generate done. Updated accumulation_step to {self._accumulation_step} and is_grad_accum was {is_grad_accum}") + print(f"[RANK 0] generate done. Updated accumulation_step to {self._accumulation_step} and sleep level was 1 during grad accumulation {is_grad_accum}") return completion_ids def reset_prefix_cache(self): """ Resets the prefix cache for the model. """ - # ToDo: perhaps we need to just pass self.llm.reset_prefix_cache() if self.process_index == 0: print(f"[RANK 0] reset_prefix_cache done.") From 675a1ed2c231b9827e73a2733336ec6c2164665c Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 9 Apr 2025 10:00:08 -0400 Subject: [PATCH 25/43] Revert to sleep level 1 - as level 2 generates randomly --- trl/extras/vllm_client.py | 39 ++++++++++++++++++--------------------- 1 file changed, 18 insertions(+), 21 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 3c59e780711..d08635c5289 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -36,11 +36,7 @@ from vllm.sampling_params import GuidedDecodingParams from accelerate import Accelerator -from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed -from contextlib import nullcontext -from ..import_utils import is_deepspeed_available -if is_deepspeed_available(): - import deepspeed +from accelerate.utils import broadcast_object_list, gather_object, set_seed logger = logging.getLogger(__name__) @@ -366,7 +362,7 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.vllm_device = accelerator.device self.world_size = accelerator.num_processes self.process_index = accelerator.process_index - self._accumulation_step = 0 + self._is_sleeping = False set_seed(42) print(f"\n\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") @@ -402,6 +398,16 @@ def __init__(self, args: GRPOConfig, model, accelerator): enable_sleep_mode=True, max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation_tp ) + + def wake_up_vllm(self): + torch.cuda.empty_cache() + if self._is_sleeping: + self.llm.wake_up() + self._is_sleeping = False + + def sleep_vllm(self): + self.llm.sleep(level=1) + self._is_sleeping = True def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -413,13 +419,12 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ - torch.cuda.empty_cache() - self.llm.wake_up() + + self.wake_up_vllm() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - self._accumulation_step = 0 if self.process_index == 0: - print(f"[RANK 0] update_named_param done: {name}, accumulation_step reset to 0.") + print(f"---[RANK 0] update_named_param done: {name}.") def generate( self, @@ -460,8 +465,7 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ - torch.cuda.empty_cache() - self.llm.wake_up() + self.wake_up_vllm() # Guided decoding, if enabled if guided_decoding_regex is not None: @@ -514,16 +518,9 @@ def generate( tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) completion_ids = completion_ids[tp_slice] - # during grad accumulation, sleep 1, otherwise sleep 2 as we can safely forget weights - new model will be loaded - is_grad_accum = ((self._accumulation_step + 1) % self.args.gradient_accumulation_steps) != 0 - if is_grad_accum: - self.llm.sleep(level=1) - else: - self.llm.sleep(level=2) - self._accumulation_step += 1 - + self.sleep_vllm() if self.process_index == 0: - print(f"[RANK 0] generate done. Updated accumulation_step to {self._accumulation_step} and sleep level was 1 during grad accumulation {is_grad_accum}") + print(f"[RANK 0] generate done.") return completion_ids def reset_prefix_cache(self): From 91d7e7225b11d26a611920bb54f4af5f3d145ba6 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 9 Apr 2025 10:27:32 -0400 Subject: [PATCH 26/43] Conduct 72b experiment --- trl/extras/vllm_client.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index d08635c5289..d577ebc701b 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -419,12 +419,9 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ - self.wake_up_vllm() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - if self.process_index == 0: - print(f"---[RANK 0] update_named_param done: {name}.") def generate( self, @@ -481,8 +478,8 @@ def generate( torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) prompts = [p for sublist in gathered_prompts for p in sublist] - print(f"\n\n---Rank {self.process_index} generation... check prompts, " - f"orig_size: {orig_size}, local group prompts size: {len(prompts)}, " + print(f"\n\n------Rank {self.process_index} generation... check prompts, " + f"orig_size: {orig_size} * tp_size: {self.args.vllm_colocation_tp} = local group prompts size: {len(prompts)}, " f"should be equal to: {orig_size * self.args.vllm_colocation_tp}") if self.process_index == 0 else None sampling_params = SamplingParams( @@ -520,7 +517,7 @@ def generate( self.sleep_vllm() if self.process_index == 0: - print(f"[RANK 0] generate done.") + print(f"------[RANK 0] generate done.") return completion_ids def reset_prefix_cache(self): @@ -529,7 +526,7 @@ def reset_prefix_cache(self): """ self.llm.reset_prefix_cache() if self.process_index == 0: - print(f"[RANK 0] reset_prefix_cache done.") + print(f"----[RANK 0] reset_prefix_cache done.") def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From 65666eb782066861d0c0ad96ebc07d3f66ec9e96 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 9 Apr 2025 14:50:58 -0400 Subject: [PATCH 27/43] Remove prints --- trl/extras/vllm_client.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index d577ebc701b..da3e6a7809b 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -364,7 +364,6 @@ def __init__(self, args: GRPOConfig, model, accelerator): self.process_index = accelerator.process_index self._is_sleeping = False set_seed(42) - print(f"\n\n------ device {self.vllm_device}, tp size: {self.args.vllm_colocation_tp}, process index: {self.process_index}, process length {self.world_size}") if self.args.vllm_colocation_tp: # Ensure TP value is valid (at least 1) @@ -478,10 +477,6 @@ def generate( torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) prompts = [p for sublist in gathered_prompts for p in sublist] - print(f"\n\n------Rank {self.process_index} generation... check prompts, " - f"orig_size: {orig_size} * tp_size: {self.args.vllm_colocation_tp} = local group prompts size: {len(prompts)}, " - f"should be equal to: {orig_size * self.args.vllm_colocation_tp}") if self.process_index == 0 else None - sampling_params = SamplingParams( n=1, # vLLM on each device generates only 1 in vllm_colocation mode repetition_penalty=repetition_penalty, @@ -497,15 +492,6 @@ def generate( prompts, sampling_params=sampling_params, use_tqdm=False ) - if self.process_index == 0: - print("\n\n==== Rank 0 Prompt/Generation Output ====\n") - for i, (prompt, outputs) in enumerate(zip(prompts, all_outputs)): - print(f"--- Prompt {i+1} ---") - print(prompt) - print(f"--- Generation {i+1} ---") - print(outputs.outputs[0].text.strip()) - print("=" * 40) - completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] if self.args.vllm_colocation_tp: @@ -516,8 +502,6 @@ def generate( completion_ids = completion_ids[tp_slice] self.sleep_vllm() - if self.process_index == 0: - print(f"------[RANK 0] generate done.") return completion_ids def reset_prefix_cache(self): @@ -525,8 +509,6 @@ def reset_prefix_cache(self): Resets the prefix cache for the model. """ self.llm.reset_prefix_cache() - if self.process_index == 0: - print(f"----[RANK 0] reset_prefix_cache done.") def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNoOpClient: """ From 710da692ae7eeb7fd267846cea9bdd776b4ec18b Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 16 Apr 2025 11:08:37 -0400 Subject: [PATCH 28/43] Make sleep optional --- trl/extras/vllm_client.py | 7 ++++--- trl/trainer/grpo_config.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index da3e6a7809b..fd2c2f2ebf1 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -400,13 +400,14 @@ def __init__(self, args: GRPOConfig, model, accelerator): def wake_up_vllm(self): torch.cuda.empty_cache() - if self._is_sleeping: + if self.args.vllm_sleep_enabled and self._is_sleeping: self.llm.wake_up() self._is_sleeping = False def sleep_vllm(self): - self.llm.sleep(level=1) - self._is_sleeping = True + if self.args.vllm_sleep_enabled: + self.llm.sleep(level=1) + self._is_sleeping = True def update_named_param(self, name: str, weights: torch.Tensor): """ diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index 2761906ee8a..20d4fdd795e 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -95,6 +95,9 @@ class GRPOConfig(TrainingArguments): - Set to `None` to disable colocated vLLM entirely. - Set to `1` to enable colocated vLLM on each GPU with no tensor parallelism. - Set to a value >1 to enable colocated vLLM with tensor parallelism across multiple GPUs. + vllm_sleep_enabled (`bool`, *optional*, defaults to `False`): + Indicates whether to enable the sleep operation for vLLM during training. + If set to `True`, vLLM will remain in sleep mode throughout the training stage. > Parameters that control the training @@ -266,6 +269,16 @@ class GRPOConfig(TrainingArguments): ) }, ) + vllm_sleep_enabled: Optional[bool] = field( + default=False, + metadata={ + "help": ( + "Enables sleep mode for colocated vLLM during training. " + "Set to `True` to keep vLLM in sleep state during training steps, helping reduce memory usage. " + "Set to `False` to disable this behavior." + ) + }, + ) # Parameters that control the training learning_rate: float = field( From a426f59d61b68e492266ba26318f90285ccef2db Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Wed, 16 Apr 2025 16:24:23 -0400 Subject: [PATCH 29/43] Incorporate feedback --- trl/extras/vllm_client.py | 97 +++++++++++++++++++++----------------- trl/trainer/grpo_config.py | 4 +- 2 files changed, 56 insertions(+), 45 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index fd2c2f2ebf1..88131da4b4c 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -14,7 +14,6 @@ import atexit import logging -import os import time from typing import Optional @@ -343,48 +342,48 @@ def close_communicator(self): class VLLMColocationClient: """ - A client class to interact with vLLM processes colocated with the training process. + A client class for interacting with vLLM models 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. + This client eliminates remote communication overhead by directly interfacing with the in-process vLLM engine. + It supports weight updates and text generation, and is optimized for tensor-parallel setups where multiple + ranks share a single vLLM engine per node or process group. 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"). + args (GRPOConfig): Configuration object with vLLM-specific parameters. + model (transformers.PreTrainedModel): The model used for generation and weight updates. + accelerator_device (str): Device where the model is loaded. + accelerator_num_processes (int): Total number of distributed processes (world size). + accelerator_process_index (int): Index of the current process in the distributed setup. """ - - def __init__(self, args: GRPOConfig, model, accelerator): + def __init__(self, args: GRPOConfig, model, accelerator_device, accelerator_num_processes, accelerator_process_index): self.args = args self.model = model - self.accelerator = accelerator - self.vllm_device = accelerator.device - self.world_size = accelerator.num_processes - self.process_index = accelerator.process_index + self.vllm_device = accelerator_device + self.world_size = accelerator_num_processes + self.process_index = accelerator_process_index self._is_sleeping = False set_seed(42) - if self.args.vllm_colocation_tp: - # Ensure TP value is valid (at least 1) - assert self.args.vllm_colocation_tp >= 1, "vllm_colocation_tp must be greater than 0" - - # Make sure TP group size evenly divides the world size - # This ensures each group has the same number of ranks - assert self.world_size % self.args.vllm_colocation_tp == 0, ( - f"TP size of vllm_colocation_tp ({self.args.vllm_colocation_tp}) must divide world size " - f"({self.world_size}) evenly." - ) + # Ensure TP value is valid (at least 1) + assert self.args.vllm_colocation >= 1, "vllm_colocation must be greater than 0" - # Create subgroups of ranks for TP, each group with `vllm_colocation_tp` ranks. - # For example, if world_size=8 and vllm_colocation_tp=2 → groups: [0,1], [2,3], [4,5], [6,7] - self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration( - [ - list(range(i*self.args.vllm_colocation_tp, (i+1) * self.args.vllm_colocation_tp)) - for i in range(self.world_size // self.args.vllm_colocation_tp) - ] + # Make sure TP group size evenly divides the world size + # This ensures each group has the same number of ranks + assert self.world_size % self.args.vllm_colocation == 0, ( + f"TP size of vllm_colocation ({self.args.vllm_colocation}) must divide world size " + f"({self.world_size}) evenly." ) + if self.args.vllm_colocation > 1: # if model is sharded, create subgroups + # Create subgroups of ranks for TP, each group with `vllm_colocation` ranks. + # For example, if world_size=8 and vllm_colocation=2 → groups: [0,1], [2,3], [4,5], [6,7] + self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration( + [ + list(range(i*self.args.vllm_colocation, (i+1) * self.args.vllm_colocation)) + for i in range(self.world_size // self.args.vllm_colocation) + ] + ) + self.llm = LLM( model=self.model.name_or_path, device=self.vllm_device, @@ -392,19 +391,31 @@ def __init__(self, args: GRPOConfig, model, accelerator): dtype=self.args.vllm_dtype, enable_prefix_caching=self.args.vllm_enable_prefix_caching, max_model_len=self.args.vllm_max_model_len, - tensor_parallel_size=args.vllm_colocation_tp, + tensor_parallel_size=args.vllm_colocation, distributed_executor_backend="external_launcher", enable_sleep_mode=True, - max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation_tp + max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation ) - def wake_up_vllm(self): + def maybe_wake_up_vllm(self): + """ + Wakes up the vLLM engine if it is currently in sleep mode. + + This is useful before any generation or weight update calls to ensure the model is active. + It also calls `torch.cuda.empty_cache()` to free unused memory, helping avoid OOM errors. + """ torch.cuda.empty_cache() if self.args.vllm_sleep_enabled and self._is_sleeping: self.llm.wake_up() self._is_sleeping = False - def sleep_vllm(self): + def maybe_sleep_vllm(self): + """ + Puts the vLLM engine into sleep mode after generation is complete. + + This helps conserve memory by offloading cached resources. + The sleep only happens if `vllm_sleep_enabled` is set to True in the config. + """ if self.args.vllm_sleep_enabled: self.llm.sleep(level=1) self._is_sleeping = True @@ -419,7 +430,7 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ - self.wake_up_vllm() + self.maybe_wake_up_vllm() llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) @@ -462,7 +473,7 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ - self.wake_up_vllm() + self.maybe_wake_up_vllm() # Guided decoding, if enabled if guided_decoding_regex is not None: @@ -470,11 +481,11 @@ def generate( else: guided_decoding = None - if self.args.vllm_colocation_tp: + if self.args.vllm_colocation > 1: # Gather prompts from all ranks in the TP group and flatten. # Each rank starts with its own prompts; after gathering, all ranks see the full group set. orig_size = len(prompts) - gathered_prompts = [None for _ in range(self.args.vllm_colocation_tp)] + gathered_prompts = [None for _ in range(self.args.vllm_colocation)] torch.distributed.all_gather_object(gathered_prompts, prompts, group=self.tp_group) prompts = [p for sublist in gathered_prompts for p in sublist] @@ -495,14 +506,14 @@ def generate( completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] - if self.args.vllm_colocation_tp: + if self.args.vllm_colocation > 1: # Slice completions for this rank within its TP group. # Each rank generates all outputs — we keep only our share. local_rank_in_group = torch.distributed.get_rank(group=self.tp_group) tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) completion_ids = completion_ids[tp_slice] - self.sleep_vllm() + self.maybe_sleep_vllm() return completion_ids def reset_prefix_cache(self): @@ -528,8 +539,8 @@ def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNo 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_tp: - return VLLMColocationClient(args, model, accelerator) + if args.vllm_colocation: + return VLLMColocationClient(args, model, accelerator.device, accelerator.num_processes, accelerator.process_index) elif accelerator.is_main_process: return VLLMClient( args.vllm_server_host, args.vllm_server_port, connection_timeout=args.vllm_server_timeout, diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index 20d4fdd795e..c200696accd 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -90,7 +90,7 @@ 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_tp (`int` or `None`, *optional*, defaults to `None`): + vllm_colocation (`int` or `None`, *optional*, defaults to `None`): Controls colocated vLLM execution and tensor parallelism via the `external_launcher` backend. - Set to `None` to disable colocated vLLM entirely. - Set to `1` to enable colocated vLLM on each GPU with no tensor parallelism. @@ -258,7 +258,7 @@ class GRPOConfig(TrainingArguments): default=None, metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."}, ) - vllm_colocation_tp: Optional[int] = field( + vllm_colocation: Optional[int] = field( default=None, metadata={ "help": ( From a1dd8e4ad1010e6b8377db9d4aa0e30fe5d454b4 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Fri, 18 Apr 2025 10:10:14 -0400 Subject: [PATCH 30/43] Incorporate Fabians comments --- trl/extras/vllm_client.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 3d59ae75049..b380e00173c 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -359,16 +359,16 @@ class VLLMColocationClient: Args: args (GRPOConfig): Configuration object with vLLM-specific parameters. model (transformers.PreTrainedModel): The model used for generation and weight updates. - accelerator_device (str): Device where the model is loaded. - accelerator_num_processes (int): Total number of distributed processes (world size). - accelerator_process_index (int): Index of the current process in the distributed setup. + device (str): Device where the model is loaded. + num_processes (int): Total number of distributed processes (world size). + process_index (int): Index of the current process in the distributed setup. """ - def __init__(self, args: GRPOConfig, model, accelerator_device, accelerator_num_processes, accelerator_process_index): + def __init__(self, args: GRPOConfig, model, device, num_processes, process_index): self.args = args self.model = model - self.vllm_device = accelerator_device - self.world_size = accelerator_num_processes - self.process_index = accelerator_process_index + self.vllm_device = device + self.world_size = num_processes + self.process_index = process_index self._is_sleeping = False set_seed(42) @@ -401,7 +401,7 @@ def __init__(self, args: GRPOConfig, model, accelerator_device, accelerator_num_ max_model_len=self.args.vllm_max_model_len, tensor_parallel_size=args.vllm_colocation, distributed_executor_backend="external_launcher", - enable_sleep_mode=True, + enable_sleep_mode=self.args.vllm_sleep_enabled, max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation ) From 2f95c00d010fa77139e1b7fd916100f49008d463 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 11:52:54 -0400 Subject: [PATCH 31/43] Revert to sleep 2 and reload model during grad accumulation --- trl/extras/vllm_client.py | 35 ++++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index b380e00173c..abab1a7737f 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -39,6 +39,11 @@ from accelerate import Accelerator from accelerate.utils import broadcast_object_list, gather_object, set_seed +if is_deepspeed_available(): + import deepspeed + +from contextlib import nullcontext + logger = logging.getLogger(__name__) class VLLMNoOpClient: @@ -370,6 +375,7 @@ def __init__(self, args: GRPOConfig, model, device, num_processes, process_index self.world_size = num_processes self.process_index = process_index self._is_sleeping = False + self._grad_accumulation = True set_seed(42) # Ensure TP value is valid (at least 1) @@ -405,6 +411,17 @@ def __init__(self, args: GRPOConfig, model, device, num_processes, process_index max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation ) + def load_model_during_grad_accumulation(self): + # Only load model during the grad accumulation steps - otherwise model was just loaded + print(f"\n\n---Rank {self.process_index} updating the model now during grad accumulation") if self.process_index == 0 else None + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 + gather_if_zero3 = deepspeed.zero.GatheredParameters if zero_stage_3 else nullcontext + # ToDo: For now, focused on non-PEFT models, simply gather and update each parameter individually. + for name, param in self.model.named_parameters(): + with gather_if_zero3([param]): + self.update_named_param(name, param.data) + def maybe_wake_up_vllm(self): """ Wakes up the vLLM engine if it is currently in sleep mode. @@ -415,6 +432,9 @@ def maybe_wake_up_vllm(self): torch.cuda.empty_cache() if self.args.vllm_sleep_enabled and self._is_sleeping: self.llm.wake_up() + print(f"\n\n---Rank {self.process_index} woke up now") if self.process_index == 0 else None + if self._grad_accumulation: + self.load_model_during_grad_accumulation() self._is_sleeping = False def maybe_sleep_vllm(self): @@ -425,7 +445,8 @@ def maybe_sleep_vllm(self): The sleep only happens if `vllm_sleep_enabled` is set to True in the config. """ if self.args.vllm_sleep_enabled: - self.llm.sleep(level=1) + print(f"\n\n---Rank {self.process_index} sleeping now with level 2") if self.process_index == 0 else None + self.llm.sleep(level=2) self._is_sleeping = True def update_named_param(self, name: str, weights: torch.Tensor): @@ -438,9 +459,12 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ + self._grad_accumulation = False # updating model weights - not grad accumulation self.maybe_wake_up_vllm() + print(f"\n\n---Rank {self.process_index} updating model at big step") if self.process_index == 0 else None llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) + self._grad_accumulation = True def generate( self, @@ -512,6 +536,15 @@ def generate( prompts, sampling_params=sampling_params, use_tqdm=False ) + if self.process_index == 0: + print("\n\n==== Rank 0 Prompt/Generation Output ====\n") + for i, (prompt, outputs) in enumerate(zip(prompts, all_outputs)): + print(f"--- Prompt {i+1} ---") + print(prompt) + print(f"--- Generation {i+1} ---") + print(outputs.outputs[0].text.strip()) + print("=" * 40) + completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] if self.args.vllm_colocation > 1: From 9c1504465d3617e2cc7e974cd3e4d4678da82ad8 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 12:00:52 -0400 Subject: [PATCH 32/43] Include accelerator in vllm client to access deepspeed --- trl/extras/vllm_client.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index abab1a7737f..1b30d48541b 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -368,12 +368,13 @@ class VLLMColocationClient: num_processes (int): Total number of distributed processes (world size). process_index (int): Index of the current process in the distributed setup. """ - def __init__(self, args: GRPOConfig, model, device, num_processes, process_index): + def __init__(self, args: GRPOConfig, model, device, num_processes, process_index, accelerator): self.args = args self.model = model self.vllm_device = device self.world_size = num_processes self.process_index = process_index + self.accelerator = accelerator self._is_sleeping = False self._grad_accumulation = True set_seed(42) @@ -581,7 +582,7 @@ def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNo accelerator (`Accelerator`): Hugging Face `Accelerator` object that helps with multi-GPU training. """ if args.vllm_colocation: - return VLLMColocationClient(args, model, accelerator.device, accelerator.num_processes, accelerator.process_index) + return VLLMColocationClient(args, model, accelerator.device, accelerator.num_processes, accelerator.process_index, accelerator) elif accelerator.is_main_process: return VLLMClient( args.vllm_server_host, args.vllm_server_port, connection_timeout=args.vllm_server_timeout, From 315182829c653a70ffb56a48175fb8d6639b989e Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 12:03:40 -0400 Subject: [PATCH 33/43] Import deepspeed avaialble --- trl/extras/vllm_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 1b30d48541b..746ed3b7e8b 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -20,7 +20,7 @@ import torch from torch import nn -from ..import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available +from ..import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available, is_deepspeed_available from ..trainer.grpo_config import GRPOConfig From f44b0fe37da90c897b933b6d9e0ad08b3b00d728 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 12:08:14 -0400 Subject: [PATCH 34/43] Set seed in llm init --- trl/extras/vllm_client.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 746ed3b7e8b..5a3886edd3b 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -14,6 +14,7 @@ import atexit import logging +import os import time from typing import Optional @@ -406,10 +407,11 @@ def __init__(self, args: GRPOConfig, model, device, num_processes, process_index dtype=self.args.vllm_dtype, enable_prefix_caching=self.args.vllm_enable_prefix_caching, max_model_len=self.args.vllm_max_model_len, - tensor_parallel_size=args.vllm_colocation, + tensor_parallel_size=self.args.vllm_colocation, distributed_executor_backend="external_launcher", enable_sleep_mode=self.args.vllm_sleep_enabled, - max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation + max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation, + seed=int(os.getenv("RANK", "0")) // self.args.vllm_colocation, ) def load_model_during_grad_accumulation(self): From 0bb81023f47d0c6a1c795559fdaf1b896bf8a346 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 13:48:19 -0400 Subject: [PATCH 35/43] Parametrize sleep level 2 and compare --- trl/extras/vllm_client.py | 4 ++-- trl/trainer/grpo_config.py | 9 +++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 5a3886edd3b..d3c5927952d 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -436,7 +436,7 @@ def maybe_wake_up_vllm(self): if self.args.vllm_sleep_enabled and self._is_sleeping: self.llm.wake_up() print(f"\n\n---Rank {self.process_index} woke up now") if self.process_index == 0 else None - if self._grad_accumulation: + if self._grad_accumulation and self.args.vllm_sleep_level2: self.load_model_during_grad_accumulation() self._is_sleeping = False @@ -449,7 +449,7 @@ def maybe_sleep_vllm(self): """ if self.args.vllm_sleep_enabled: print(f"\n\n---Rank {self.process_index} sleeping now with level 2") if self.process_index == 0 else None - self.llm.sleep(level=2) + self.llm.sleep(level=2) if self.args.vllm_sleep_level2 else self.llm.sleep(level=1) self._is_sleeping = True def update_named_param(self, name: str, weights: torch.Tensor): diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index a85c9c5cbc6..3263f7ccbd4 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -308,6 +308,7 @@ class GRPOConfig(TrainingArguments): ) }, ) + vllm_sleep_enabled: Optional[bool] = field( default=False, metadata={ @@ -318,6 +319,14 @@ class GRPOConfig(TrainingArguments): ) }, ) + vllm_sleep_level2: Optional[bool] = field( + default=False, + metadata={ + "help": ( + "Sleep level 2 enabled - otherwise sleep level 1" + ) + }, + ) # Parameters that control the training learning_rate: float = field( From 589ffc97b080c68e40cb262089b4ddbb20f2fa57 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 15:21:45 -0400 Subject: [PATCH 36/43] Debug level 1 and level 2 --- trl/extras/vllm_client.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index d3c5927952d..50aab20b0da 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -448,8 +448,12 @@ def maybe_sleep_vllm(self): The sleep only happens if `vllm_sleep_enabled` is set to True in the config. """ if self.args.vllm_sleep_enabled: - print(f"\n\n---Rank {self.process_index} sleeping now with level 2") if self.process_index == 0 else None - self.llm.sleep(level=2) if self.args.vllm_sleep_level2 else self.llm.sleep(level=1) + if self.args.vllm_sleep_level2: + print(f"\n\n---Rank {self.process_index} sleeping now with level 2") if self.process_index == 0 else None + self.llm.sleep(level=2) + else: + print(f"\n\n---Rank {self.process_index} sleeping now with level 1") if self.process_index == 0 else None + self.llm.sleep(level=1) self._is_sleeping = True def update_named_param(self, name: str, weights: torch.Tensor): From 2eb185548dcdf6c96218fb2a08fc8ca8f672c413 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 15:35:18 -0400 Subject: [PATCH 37/43] Fix grad accumulation for sleep 2 --- trl/extras/vllm_client.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 50aab20b0da..e9f1000e74f 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -433,9 +433,10 @@ def maybe_wake_up_vllm(self): It also calls `torch.cuda.empty_cache()` to free unused memory, helping avoid OOM errors. """ torch.cuda.empty_cache() - if self.args.vllm_sleep_enabled and self._is_sleeping: - self.llm.wake_up() - print(f"\n\n---Rank {self.process_index} woke up now") if self.process_index == 0 else None + if self.args.vllm_sleep_enabled: + if self._is_sleeping: + self.llm.wake_up() + print(f"\n\n---Rank {self.process_index} woke up now") if self.process_index == 0 else None if self._grad_accumulation and self.args.vllm_sleep_level2: self.load_model_during_grad_accumulation() self._is_sleeping = False @@ -455,6 +456,7 @@ def maybe_sleep_vllm(self): print(f"\n\n---Rank {self.process_index} sleeping now with level 1") if self.process_index == 0 else None self.llm.sleep(level=1) self._is_sleeping = True + self._grad_accumulation = True # done with current step, so next is grad accumulation unless it is set in update_named_param def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -471,7 +473,7 @@ def update_named_param(self, name: str, weights: torch.Tensor): print(f"\n\n---Rank {self.process_index} updating model at big step") if self.process_index == 0 else None llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - self._grad_accumulation = True + # self._grad_accumulation = True def generate( self, From 82d8d9424533b1b97ed17284f2ed8b9a0842db5f Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 15:54:41 -0400 Subject: [PATCH 38/43] Fix grad accumulation for sleep 2 --- trl/extras/vllm_client.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index e9f1000e74f..af888b8099c 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -416,14 +416,17 @@ def __init__(self, args: GRPOConfig, model, device, num_processes, process_index def load_model_during_grad_accumulation(self): # Only load model during the grad accumulation steps - otherwise model was just loaded - print(f"\n\n---Rank {self.process_index} updating the model now during grad accumulation") if self.process_index == 0 else None - deepspeed_plugin = self.accelerator.state.deepspeed_plugin - zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 - gather_if_zero3 = deepspeed.zero.GatheredParameters if zero_stage_3 else nullcontext - # ToDo: For now, focused on non-PEFT models, simply gather and update each parameter individually. - for name, param in self.model.named_parameters(): - with gather_if_zero3([param]): - self.update_named_param(name, param.data) + print(f"\n---Rank {self.process_index} load_model_during_grad_accumulation checking") if self.process_index == 0 else None + if self.args.vllm_sleep_level2: + print(f"\n---Rank {self.process_index} updating the model now during grad accumulation") if self.process_index == 0 else None + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 + gather_if_zero3 = deepspeed.zero.GatheredParameters if zero_stage_3 else nullcontext + # ToDo: For now, focused on non-PEFT models, simply gather and update each parameter individually. + for name, param in self.model.named_parameters(): + with gather_if_zero3([param]): + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name,param.data)]) def maybe_wake_up_vllm(self): """ @@ -434,10 +437,12 @@ def maybe_wake_up_vllm(self): """ torch.cuda.empty_cache() if self.args.vllm_sleep_enabled: + print(f"\n\n---Rank {self.process_index} vllm_sleep_enabled - check to wake up") if self.process_index == 0 else None if self._is_sleeping: self.llm.wake_up() print(f"\n\n---Rank {self.process_index} woke up now") if self.process_index == 0 else None - if self._grad_accumulation and self.args.vllm_sleep_level2: + if self._grad_accumulation: + print("\n\n grad accumulation - check and load the model") self.load_model_during_grad_accumulation() self._is_sleeping = False @@ -470,7 +475,7 @@ def update_named_param(self, name: str, weights: torch.Tensor): """ self._grad_accumulation = False # updating model weights - not grad accumulation self.maybe_wake_up_vllm() - print(f"\n\n---Rank {self.process_index} updating model at big step") if self.process_index == 0 else None + print(f"---Rank {self.process_index} updating model at big step") if self.process_index == 0 else None llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) # self._grad_accumulation = True From e4fafee04c385346b37f12bfbe0a573a021e2db7 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 16:23:47 -0400 Subject: [PATCH 39/43] Fix grad accumulation for sleep 2 --- trl/extras/vllm_client.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index af888b8099c..8b7416f84ad 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -377,7 +377,7 @@ def __init__(self, args: GRPOConfig, model, device, num_processes, process_index self.process_index = process_index self.accelerator = accelerator self._is_sleeping = False - self._grad_accumulation = True + self._grad_accumulation = False set_seed(42) # Ensure TP value is valid (at least 1) @@ -413,12 +413,14 @@ def __init__(self, args: GRPOConfig, model, device, num_processes, process_index max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation, seed=int(os.getenv("RANK", "0")) // self.args.vllm_colocation, ) + + print(f"---Rank {self.process_index} initialized llm for the first time") if self.process_index == 0 else None def load_model_during_grad_accumulation(self): # Only load model during the grad accumulation steps - otherwise model was just loaded - print(f"\n---Rank {self.process_index} load_model_during_grad_accumulation checking") if self.process_index == 0 else None + print(f"---Rank {self.process_index} load_model_during_grad_accumulation checking") if self.process_index == 0 else None if self.args.vllm_sleep_level2: - print(f"\n---Rank {self.process_index} updating the model now during grad accumulation") if self.process_index == 0 else None + print(f"---Rank {self.process_index} updating the model now during grad accumulation") if self.process_index == 0 else None deepspeed_plugin = self.accelerator.state.deepspeed_plugin zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 gather_if_zero3 = deepspeed.zero.GatheredParameters if zero_stage_3 else nullcontext @@ -436,8 +438,9 @@ def maybe_wake_up_vllm(self): It also calls `torch.cuda.empty_cache()` to free unused memory, helping avoid OOM errors. """ torch.cuda.empty_cache() + print(f'----[wake_check] sleeping: {self._is_sleeping}, grad_accumulation: {self._grad_accumulation}, vllm_sleep_enabled: {self.args.vllm_sleep_enabled}, vllm_sleep_level2: {self.args.vllm_sleep_level2}') if self.process_index == 0 else None if self.args.vllm_sleep_enabled: - print(f"\n\n---Rank {self.process_index} vllm_sleep_enabled - check to wake up") if self.process_index == 0 else None + # print(f"\n\n---Rank {self.process_index} vllm_sleep_enabled - check to wake up") if self.process_index == 0 else None if self._is_sleeping: self.llm.wake_up() print(f"\n\n---Rank {self.process_index} woke up now") if self.process_index == 0 else None @@ -454,6 +457,7 @@ def maybe_sleep_vllm(self): The sleep only happens if `vllm_sleep_enabled` is set to True in the config. """ if self.args.vllm_sleep_enabled: + print(f'-----[sleep_check] sleeping: {self._is_sleeping}, grad_accumulation: {self._grad_accumulation}, vllm_sleep_enabled: {self.args.vllm_sleep_enabled}, vllm_sleep_level2: {self.args.vllm_sleep_level2}') if self.process_index == 0 else None if self.args.vllm_sleep_level2: print(f"\n\n---Rank {self.process_index} sleeping now with level 2") if self.process_index == 0 else None self.llm.sleep(level=2) @@ -462,6 +466,7 @@ def maybe_sleep_vllm(self): self.llm.sleep(level=1) self._is_sleeping = True self._grad_accumulation = True # done with current step, so next is grad accumulation unless it is set in update_named_param + print(f'----[sleep_check_end] sleeping: {self._is_sleeping}, grad_accumulation: {self._grad_accumulation}, vllm_sleep_enabled: {self.args.vllm_sleep_enabled}, vllm_sleep_level2: {self.args.vllm_sleep_level2}') if self.process_index == 0 else None def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -478,7 +483,6 @@ def update_named_param(self, name: str, weights: torch.Tensor): print(f"---Rank {self.process_index} updating model at big step") if self.process_index == 0 else None llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) - # self._grad_accumulation = True def generate( self, @@ -519,6 +523,8 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ + + print(f"---Rank {self.process_index} start generation") if self.process_index == 0 else None self.maybe_wake_up_vllm() # Guided decoding, if enabled From 6191b8926761adb20a0802809ebdb43a9670a38d Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 17:22:28 -0400 Subject: [PATCH 40/43] Revert tpcoloc branch to commit a1dd8e4 without rewriting history --- trl/extras/vllm_client.py | 69 +++++--------------------------------- trl/trainer/grpo_config.py | 9 ----- 2 files changed, 8 insertions(+), 70 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 8b7416f84ad..b380e00173c 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -14,14 +14,13 @@ import atexit import logging -import os import time from typing import Optional import torch from torch import nn -from ..import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available, is_deepspeed_available +from ..import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available from ..trainer.grpo_config import GRPOConfig @@ -40,11 +39,6 @@ from accelerate import Accelerator from accelerate.utils import broadcast_object_list, gather_object, set_seed -if is_deepspeed_available(): - import deepspeed - -from contextlib import nullcontext - logger = logging.getLogger(__name__) class VLLMNoOpClient: @@ -369,15 +363,13 @@ class VLLMColocationClient: num_processes (int): Total number of distributed processes (world size). process_index (int): Index of the current process in the distributed setup. """ - def __init__(self, args: GRPOConfig, model, device, num_processes, process_index, accelerator): + def __init__(self, args: GRPOConfig, model, device, num_processes, process_index): self.args = args self.model = model self.vllm_device = device self.world_size = num_processes self.process_index = process_index - self.accelerator = accelerator self._is_sleeping = False - self._grad_accumulation = False set_seed(42) # Ensure TP value is valid (at least 1) @@ -407,29 +399,12 @@ def __init__(self, args: GRPOConfig, model, device, num_processes, process_index dtype=self.args.vllm_dtype, enable_prefix_caching=self.args.vllm_enable_prefix_caching, max_model_len=self.args.vllm_max_model_len, - tensor_parallel_size=self.args.vllm_colocation, + tensor_parallel_size=args.vllm_colocation, distributed_executor_backend="external_launcher", enable_sleep_mode=self.args.vllm_sleep_enabled, - max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation, - seed=int(os.getenv("RANK", "0")) // self.args.vllm_colocation, + max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation ) - - print(f"---Rank {self.process_index} initialized llm for the first time") if self.process_index == 0 else None - def load_model_during_grad_accumulation(self): - # Only load model during the grad accumulation steps - otherwise model was just loaded - print(f"---Rank {self.process_index} load_model_during_grad_accumulation checking") if self.process_index == 0 else None - if self.args.vllm_sleep_level2: - print(f"---Rank {self.process_index} updating the model now during grad accumulation") if self.process_index == 0 else None - deepspeed_plugin = self.accelerator.state.deepspeed_plugin - zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 - gather_if_zero3 = deepspeed.zero.GatheredParameters if zero_stage_3 else nullcontext - # ToDo: For now, focused on non-PEFT models, simply gather and update each parameter individually. - for name, param in self.model.named_parameters(): - with gather_if_zero3([param]): - llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model - llm_model.load_weights([(name,param.data)]) - def maybe_wake_up_vllm(self): """ Wakes up the vLLM engine if it is currently in sleep mode. @@ -438,15 +413,8 @@ def maybe_wake_up_vllm(self): It also calls `torch.cuda.empty_cache()` to free unused memory, helping avoid OOM errors. """ torch.cuda.empty_cache() - print(f'----[wake_check] sleeping: {self._is_sleeping}, grad_accumulation: {self._grad_accumulation}, vllm_sleep_enabled: {self.args.vllm_sleep_enabled}, vllm_sleep_level2: {self.args.vllm_sleep_level2}') if self.process_index == 0 else None - if self.args.vllm_sleep_enabled: - # print(f"\n\n---Rank {self.process_index} vllm_sleep_enabled - check to wake up") if self.process_index == 0 else None - if self._is_sleeping: - self.llm.wake_up() - print(f"\n\n---Rank {self.process_index} woke up now") if self.process_index == 0 else None - if self._grad_accumulation: - print("\n\n grad accumulation - check and load the model") - self.load_model_during_grad_accumulation() + if self.args.vllm_sleep_enabled and self._is_sleeping: + self.llm.wake_up() self._is_sleeping = False def maybe_sleep_vllm(self): @@ -457,16 +425,8 @@ def maybe_sleep_vllm(self): The sleep only happens if `vllm_sleep_enabled` is set to True in the config. """ if self.args.vllm_sleep_enabled: - print(f'-----[sleep_check] sleeping: {self._is_sleeping}, grad_accumulation: {self._grad_accumulation}, vllm_sleep_enabled: {self.args.vllm_sleep_enabled}, vllm_sleep_level2: {self.args.vllm_sleep_level2}') if self.process_index == 0 else None - if self.args.vllm_sleep_level2: - print(f"\n\n---Rank {self.process_index} sleeping now with level 2") if self.process_index == 0 else None - self.llm.sleep(level=2) - else: - print(f"\n\n---Rank {self.process_index} sleeping now with level 1") if self.process_index == 0 else None - self.llm.sleep(level=1) + self.llm.sleep(level=1) self._is_sleeping = True - self._grad_accumulation = True # done with current step, so next is grad accumulation unless it is set in update_named_param - print(f'----[sleep_check_end] sleeping: {self._is_sleeping}, grad_accumulation: {self._grad_accumulation}, vllm_sleep_enabled: {self.args.vllm_sleep_enabled}, vllm_sleep_level2: {self.args.vllm_sleep_level2}') if self.process_index == 0 else None def update_named_param(self, name: str, weights: torch.Tensor): """ @@ -478,9 +438,7 @@ def update_named_param(self, name: str, weights: torch.Tensor): weights (`torch.Tensor`): Tensor containing the updated weights. """ - self._grad_accumulation = False # updating model weights - not grad accumulation self.maybe_wake_up_vllm() - print(f"---Rank {self.process_index} updating model at big step") if self.process_index == 0 else None llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model llm_model.load_weights([(name,weights)]) @@ -523,8 +481,6 @@ def generate( `list[list[int]]`: List of lists of token IDs representing the model-generated completions for each prompt. """ - - print(f"---Rank {self.process_index} start generation") if self.process_index == 0 else None self.maybe_wake_up_vllm() # Guided decoding, if enabled @@ -556,15 +512,6 @@ def generate( prompts, sampling_params=sampling_params, use_tqdm=False ) - if self.process_index == 0: - print("\n\n==== Rank 0 Prompt/Generation Output ====\n") - for i, (prompt, outputs) in enumerate(zip(prompts, all_outputs)): - print(f"--- Prompt {i+1} ---") - print(prompt) - print(f"--- Generation {i+1} ---") - print(outputs.outputs[0].text.strip()) - print("=" * 40) - completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] if self.args.vllm_colocation > 1: @@ -601,7 +548,7 @@ def get_vllm_client(args: GRPOConfig, model, accelerator: Accelerator) -> VLLMNo accelerator (`Accelerator`): Hugging Face `Accelerator` object that helps with multi-GPU training. """ if args.vllm_colocation: - return VLLMColocationClient(args, model, accelerator.device, accelerator.num_processes, accelerator.process_index, accelerator) + return VLLMColocationClient(args, model, accelerator.device, accelerator.num_processes, accelerator.process_index) elif accelerator.is_main_process: return VLLMClient( args.vllm_server_host, args.vllm_server_port, connection_timeout=args.vllm_server_timeout, diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index 3263f7ccbd4..a85c9c5cbc6 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -308,7 +308,6 @@ class GRPOConfig(TrainingArguments): ) }, ) - vllm_sleep_enabled: Optional[bool] = field( default=False, metadata={ @@ -319,14 +318,6 @@ class GRPOConfig(TrainingArguments): ) }, ) - vllm_sleep_level2: Optional[bool] = field( - default=False, - metadata={ - "help": ( - "Sleep level 2 enabled - otherwise sleep level 1" - ) - }, - ) # Parameters that control the training learning_rate: float = field( From 717c8b4d8e8494f56a56522e0f6688c1660f4b1e Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Mon, 21 Apr 2025 17:25:21 -0400 Subject: [PATCH 41/43] Revert to sleep 2 after grad acc optimization fixing the model load error --- trl/extras/vllm_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index b380e00173c..3a6b760595d 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -425,7 +425,7 @@ def maybe_sleep_vllm(self): The sleep only happens if `vllm_sleep_enabled` is set to True in the config. """ if self.args.vllm_sleep_enabled: - self.llm.sleep(level=1) + self.llm.sleep(level=2) self._is_sleeping = True def update_named_param(self, name: str, weights: torch.Tensor): From 4badf5bc517d36784526a96458fd0b824e60ef6f Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Tue, 22 Apr 2025 08:54:29 -0400 Subject: [PATCH 42/43] Comparison of sleep levels --- trl/extras/vllm_client.py | 9 +++++++-- trl/trainer/grpo_config.py | 8 ++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index 3a6b760595d..bb1757617cf 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -14,6 +14,7 @@ import atexit import logging +import os import time from typing import Optional @@ -402,7 +403,8 @@ def __init__(self, args: GRPOConfig, model, device, num_processes, process_index tensor_parallel_size=args.vllm_colocation, distributed_executor_backend="external_launcher", enable_sleep_mode=self.args.vllm_sleep_enabled, - max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation + max_num_seqs=self.args.per_device_train_batch_size * self.args.vllm_colocation, + seed=int(os.getenv("RANK", "0")) // self.args.vllm_colocation, # feed identical seed for tp groups ) def maybe_wake_up_vllm(self): @@ -425,7 +427,10 @@ def maybe_sleep_vllm(self): The sleep only happens if `vllm_sleep_enabled` is set to True in the config. """ if self.args.vllm_sleep_enabled: - self.llm.sleep(level=2) + if self.args.vllm_sleep_level1: + self.llm.sleep(level=1) + else: + self.llm.sleep(level=2) self._is_sleeping = True def update_named_param(self, name: str, weights: torch.Tensor): diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py index a85c9c5cbc6..4d9751e2878 100644 --- a/trl/trainer/grpo_config.py +++ b/trl/trainer/grpo_config.py @@ -318,6 +318,14 @@ class GRPOConfig(TrainingArguments): ) }, ) + vllm_sleep_level1: Optional[bool] = field( + default=False, + metadata={ + "help": ( + "Sleep level 1 enabled - otherwise sleep level 2 default" + ) + }, + ) # Parameters that control the training learning_rate: float = field( From a0c677a651e86ec6000b2d51ac15ce6da42668f1 Mon Sep 17 00:00:00 2001 From: Mert Toslali Date: Sat, 10 May 2025 09:57:48 -0400 Subject: [PATCH 43/43] Add max_num_batched_tokens needed for v1 profiling --- trl/extras/vllm_client.py | 1 + 1 file changed, 1 insertion(+) diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py index bb1757617cf..0b5e364eaa2 100644 --- a/trl/extras/vllm_client.py +++ b/trl/extras/vllm_client.py @@ -400,6 +400,7 @@ def __init__(self, args: GRPOConfig, model, device, num_processes, process_index dtype=self.args.vllm_dtype, enable_prefix_caching=self.args.vllm_enable_prefix_caching, max_model_len=self.args.vllm_max_model_len, + max_num_batched_tokens=self.args.vllm_max_model_len, tensor_parallel_size=args.vllm_colocation, distributed_executor_backend="external_launcher", enable_sleep_mode=self.args.vllm_sleep_enabled,