10x generate in ppo w/ zero3

huggingface/trl#1483
lightDev0405 · May 28, 2024 · d0d39d5 · d0d39d5
1 parent 85d0185
commit d0d39d5
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/src/llamafactory/train/ppo/trainer.py b/src/llamafactory/train/ppo/trainer.py
@@ -13,6 +13,7 @@
 from transformers.utils import SAFE_WEIGHTS_NAME, WEIGHTS_NAME
 from trl import PPOConfig, PPOTrainer
 from trl.core import PPODecorators, logprobs_from_logits
+from trl.models.utils import unwrap_model_for_generation
 
 from ...extras.callbacks import FixValueHeadModelCallback, LogCallback
 from ...extras.logging import get_logger
@@ -322,10 +323,10 @@ def get_inputs(self, batch: Dict[str, torch.Tensor]) -> Tuple[List[torch.Tensor]
             for k, v in batch.items():
                 batch[k] = v[:, start_index:]
 
-        unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
-        generate_output: torch.Tensor = unwrapped_model.generate(
-            generation_config=self.generation_config, logits_processor=get_logits_processor(), **batch
-        )
+        with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model:
+            generate_output: torch.Tensor = unwrapped_model.generate(
+                generation_config=self.generation_config, logits_processor=get_logits_processor(), **batch
+            )
 
         if self.model_args.upcast_layernorm:
             restore_layernorm(self.model, layernorm_params)