Fix lazy_mode assignment (huggingface#1558)

Co-authored-by: Yaser Afshar <yaser.afshar@intel.com>
HabanaAI · Jan 20, 2025 · db68958 · db68958
1 parent 25306d7
commit db68958
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py
@@ -986,7 +986,13 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio
                             inputs["flash_attention_causal_mask"] = True
                 if self.model.config is not None:
                     if self.model.config.model_type in ["llama", "qwen2", "mistral", "starcoder2"]:
-                        inputs["lazy_mode"] = args.use_lazy_mode
+                        if _is_peft_model(model):
+                            forward_method = getattr(model.get_base_model(), "forward")
+                        else:
+                            forward_method = getattr(model, "forward")
+                        signature = inspect.signature(forward_method)
+                        if "lazy_mode" in signature.parameters:
+                            inputs["lazy_mode"] = args.use_lazy_mode
                 # TODO: keep syncs for fast DDP?
                 with self.accelerator.accumulate(model):
                     tr_loss_step = self.training_step(model, inputs)