From d2cd8676ab2b1cbb762bbf5cb1fe7030532d2389 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 a915eb4dab7..3f6d9b62109 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -419,9 +419,11 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N else: # Regular generation path with unwrap_model_for_generation(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_inputs, generation_config=self.generation_config ) + model.train() # Compute prompt length and extract completion ids prompt_length = prompt_inputs["input_ids"].size(1) From 23cbc145a4289b8f46d33717c3e4c9013a47587f 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 3f6d9b62109..a41aaf48bd1 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 @@ -55,6 +56,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]]] @@ -383,6 +387,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -self.max_prompt_length :] prompt_inputs["attention_mask"] = prompt_inputs["attention_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 11f598ad4f54c8ade38a545badf49f250653047d 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 a41aaf48bd1..cb49a3d8b2e 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 @@ -38,7 +37,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 @@ -57,7 +56,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. @@ -387,7 +386,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -self.max_prompt_length :] prompt_inputs["attention_mask"] = prompt_inputs["attention_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 dc172c3d1222f6d11339f0e994fc196b288cee7b 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 | 3 +++ 1 file changed, 3 insertions(+) diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py index cb49a3d8b2e..7dbdfa875cc 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -428,6 +428,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N **prompt_inputs, generation_config=self.generation_config ) model.train() + logger.info("Finishing generation") # Compute prompt length and extract completion ids prompt_length = prompt_inputs["input_ids"].size(1) @@ -547,6 +548,8 @@ def get_per_token_logps(model, input_ids, attention_mask, logits_to_keep): mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) + logger.info("Finishing loss") + return loss def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: Optional[list[str]] = None): From 28cc83e0bb4e161d738d5187c2a18abe35d1f332 Mon Sep 17 00:00:00 2001 From: Andrew White Date: Sun, 2 Feb 2025 22:43:48 -0800 Subject: [PATCH 5/5] pulling in diff from trl-2747 --- 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 7dbdfa875cc..17163363158 100644 --- a/trl/trainer/grpo_trainer.py +++ b/trl/trainer/grpo_trainer.py @@ -442,7 +442,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() # Concatenate prompt_mask with completion_mask for logit computation - prompt_mask_repeated = prompt_inputs["attention_mask"].repeat_interleave(self.num_generations, dim=0) + prompt_mask_repeated = prompt_inputs["attention_mask"].to(device).repeat_interleave(self.num_generations, dim=0) attention_mask = torch.cat([prompt_mask_repeated, completion_mask], dim=1) # (B*G, P+C) # Get the per-token log probabilities for the completions for the model and the reference model