add synchronization after loss

Luca-Calabria · Jan 16, 2025 · 1ac26cf · 1ac26cf
1 parent f48dda8
commit 1ac26cf
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py
@@ -1020,6 +1020,8 @@ def hpu_deepspeed_checkpointing(function, *checkpoint_args, use_reentrant: Optio
                 # TODO: keep syncs for fast DDP?
                 with self.accelerator.accumulate(model):
                     tr_loss_step = self.training_step(model, inputs)
+                    import habana_frameworks.torch.hpu as hthpu
+                    hthpu.synchronize()
 
                 is_last_step_and_steps_less_than_grad_acc = (
                     steps_in_epoch <= args.gradient_accumulation_steps and (step + 1) == steps_in_epoch