fix errors in trainer save (#1213)

Signed-off-by: Dillon Laird <dillonalaird@gmail.com>
intel · Jan 31, 2024 · ff501d0 · ff501d0
1 parent 9bc38ae
commit ff501d0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/...l_extension_for_transformers/neural_chat/examples/finetuning/multi_modal/llava_trainer.py b/...l_extension_for_transformers/neural_chat/examples/finetuning/multi_modal/llava_trainer.py
@@ -428,10 +428,10 @@ def _save_checkpoint(self, model, trial, metrics=None):
                     self.model.config.save_pretrained(output_dir)
                     torch.save(weight_to_save, os.path.join(output_dir, f'mm_projector.bin'))
             else:
-                super(LLaVATrainer, self)._save_checkpoint(model, trial, metrics)
+                super(GaudiLLaVATrainer, self)._save_checkpoint(model, trial, metrics)
 
         def _save(self, output_dir: Optional[str] = None, state_dict=None):
             if getattr(self.args, 'tune_mm_mlp_adapter', False):
                 pass
             else:
-                super(LLaVATrainer, self)._save(output_dir, state_dict)
+                super(GaudiLLaVATrainer, self)._save(output_dir, state_dict)
diff --git a/intel_extension_for_transformers/neural_chat/examples/finetuning/multi_modal/llava_utils.py b/intel_extension_for_transformers/neural_chat/examples/finetuning/multi_modal/llava_utils.py
@@ -149,7 +149,11 @@ def safe_save_model_for_hf_trainer(trainer: transformers.Trainer,
         return
 
     if trainer.deepspeed:
-        torch.cuda.synchronize()
+        if is_hpu_available:
+            import habana_frameworks.torch as ht
+            ht.hpu.synchronize()
+        else:
+            torch.cuda.synchronize()
         trainer.save_model(output_dir)
         return