From 5e029bf88f15a57f799808230badfca39f0f5eed Mon Sep 17 00:00:00 2001 From: Siddharth Narayanan Date: Wed, 29 Jan 2025 14:34:19 -0800 Subject: [PATCH 1/5] make sure model is in eval mode before generating --- trl/trainer/grpo_trainer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index 473b6e880ca..26e8b991791 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -445,9 +445,11 @@ def _prepare_inputs(self, inputs: dict[str, Union[torch.Tensor, Any]]) -> dict[s else: # Regular generation path with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model: + unwrapped_model.eval() # Needed to make sure use_cache works with gradient_checkpointing prompt_completion_ids = unwrapped_model.generate( prompt_ids, attention_mask=prompt_mask, generation_config=self.generation_config ) + self.model.train() # Compute prompt length and extract completion ids prompt_length = prompt_ids.size(1) From 7f74ce357d8ca6ecb60aa873fc716bd09e4dab1a Mon Sep 17 00:00:00 2001 From: Siddharth Narayanan Date: Wed, 29 Jan 2025 14:54:01 -0800 Subject: [PATCH 2/5] bit more logging --- trl/trainer/grpo_trainer.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index 26e8b991791..718cb934499 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging import os import textwrap import warnings @@ -58,6 +59,9 @@ if is_wandb_available(): import wandb + +logger = logging.getLogger(__name__) + # What we call a reward function is a callable that takes a list of prompts and completions and returns a list of # rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model. RewardFunc = Union[str, PreTrainedModel, Callable[[list, list], list[float]]] @@ -403,6 +407,7 @@ def _prepare_inputs(self, inputs: dict[str, Union[torch.Tensor, Any]]) -> dict[s prompt_ids = prompt_ids[:, -self.max_prompt_length :] prompt_mask = prompt_mask[:, -self.max_prompt_length :] + logger.debug("Starting generation") # Generate completions using either vLLM or regular generation if self.args.use_vllm: # First, have main process load weights if needed From 61e94cc1d576df42d2a12de2338a82a6b28434b3 Mon Sep 17 00:00:00 2001 From: Siddharth Narayanan Date: Wed, 29 Jan 2025 14:55:37 -0800 Subject: [PATCH 3/5] switch to transformers logging --- trl/trainer/grpo_trainer.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index 718cb934499..368d4d166ca 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import logging import os import textwrap import warnings @@ -40,7 +39,7 @@ is_wandb_available, ) from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled -from transformers.utils import is_peft_available +from transformers.utils import is_peft_available, logging from ..data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template from ..import_utils import is_vllm_available @@ -60,7 +59,7 @@ import wandb -logger = logging.getLogger(__name__) +logger = logging.get_logger("GRPOTrainer") # What we call a reward function is a callable that takes a list of prompts and completions and returns a list of # rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model. @@ -407,7 +406,7 @@ def _prepare_inputs(self, inputs: dict[str, Union[torch.Tensor, Any]]) -> dict[s prompt_ids = prompt_ids[:, -self.max_prompt_length :] prompt_mask = prompt_mask[:, -self.max_prompt_length :] - logger.debug("Starting generation") + logger.info("Starting generation") # Generate completions using either vLLM or regular generation if self.args.use_vllm: # First, have main process load weights if needed From 3bf982ca0204d815fb03f45e64bd2f304ae58e0d Mon Sep 17 00:00:00 2001 From: Siddharth M Narayanan Date: Sun, 2 Feb 2025 16:58:25 -0600 Subject: [PATCH 4/5] add more logs --- trl/trainer/grpo_trainer.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index 368d4d166ca..26e8b991791 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -39,7 +39,7 @@ is_wandb_available, ) from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled -from transformers.utils import is_peft_available, logging +from transformers.utils import is_peft_available from ..data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template from ..import_utils import is_vllm_available @@ -58,9 +58,6 @@ if is_wandb_available(): import wandb - -logger = logging.get_logger("GRPOTrainer") - # What we call a reward function is a callable that takes a list of prompts and completions and returns a list of # rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model. RewardFunc = Union[str, PreTrainedModel, Callable[[list, list], list[float]]] @@ -406,7 +403,6 @@ def _prepare_inputs(self, inputs: dict[str, Union[torch.Tensor, Any]]) -> dict[s prompt_ids = prompt_ids[:, -self.max_prompt_length :] prompt_mask = prompt_mask[:, -self.max_prompt_length :] - logger.info("Starting generation") # Generate completions using either vLLM or regular generation if self.args.use_vllm: # First, have main process load weights if needed From edf5b0ad993b600c467b51b6c2b6008d42e793ac Mon Sep 17 00:00:00 2001 From: Siddharth M Narayanan Date: Thu, 6 Feb 2025 17:14:44 -0600 Subject: [PATCH 5/5] unwrap wrapped model --- trl/trainer/grpo_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index fa55a4bb8e8..0027a14120e 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -500,7 +500,7 @@ def _prepare_inputs(self, inputs: dict[str, Union[torch.Tensor, Any]]) -> dict[s prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) else: # Regular generation path - with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model: + with unwrap_model_for_generation(self.model_wrapped, self.accelerator) as unwrapped_model: prompt_completion_ids = unwrapped_model.generate( prompt_ids, attention_mask=prompt_mask, generation_config=self.generation_config )