🫔 [GRPO] Pass wrapped model to unwrap_model_for_generation for DeepSp…

…eed Stage-3 compatibility (#2871) Co-authored-by: Quentin Gallouédec <[email protected]>
huggingface · Feb 28, 2025 · ad6a35b · ad6a35b
1 parent 7bc9858
commit ad6a35b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
@@ -655,7 +655,7 @@ def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep)
     @profiling_decorator
     def _move_model_to_vllm(self):
         with unwrap_model_for_generation(
-            self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation
+            self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation
         ) as unwrapped_model:
             if is_compiled_module(unwrapped_model):
                 unwrapped_model = unwrapped_model._orig_mod
@@ -754,7 +754,7 @@ def _generate_and_score_completions(
             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
                 )