Fixed issue with horovod and early stopping (#986)

* Fixed issue with horovod and early stopping * applied black Co-authored-by: Arnaud Wéry <[email protected]>
OpenNMT · Dec 5, 2022 · de084e9 · de084e9
1 parent a62b886
commit de084e9
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/opennmt/training.py b/opennmt/training.py
@@ -346,6 +346,15 @@ def is_master(self):
     def num_replicas(self):
         return self._hvd.size()
 
+    def _evaluate(self, evaluator, step, moving_average=None):
+        should_stop = super()._evaluate(evaluator, step, moving_average)
+        # Evaluation is only performed on master, but we want all workers
+        # to be aware of the early stopping decision.
+        should_stop = self._hvd.broadcast_object(
+            should_stop, root_rank=0, name="should_stop"
+        )
+        return should_stop
+
     def _finalize_dataset(self, dataset):
         if callable(dataset):
             dataset = dataset(