Repository: OpenNMT/OpenNMT-tf

Issue:
Possible error when training on a finite dataset with multiple GPU
When the total number of batches is not a multiple of the number of replicas (finite dataset), the training can stop with an error because some replicas receive an empty batch.

This error can happen on master, or on v2.2.0 when TensorFlow fails to use batch splitting approach to feed the replicas.


Discussion:

Resolving patch:
diff --git a/opennmt/training.py b/opennmt/training.py
--- a/opennmt/training.py
+++ b/opennmt/training.py
@@ -123,15 +123,40 @@ def _forward(next_fn):
       with tf.summary.record_if(should_record_summaries):
         with self._strategy.scope():
           per_replica_source, per_replica_target = next_fn()
-          per_replica_loss, per_replica_words = self._strategy.experimental_run_v2(
-              _accumulate_gradients, args=(per_replica_source, per_replica_target))
 
-          # TODO: these reductions could be delayed until _step is called.
-          loss = self._strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_loss, None)
-          num_words = {
-              k:self._strategy.reduce(tf.distribute.ReduceOp.SUM, v, None)
-              for k, v in six.iteritems(per_replica_words)}
-      return loss, num_words
+          def _run():
+            per_replica_loss, per_replica_words = self._strategy.experimental_run_v2(
+                _accumulate_gradients, args=(per_replica_source, per_replica_target))
+
+            # TODO: these reductions could be delayed until _step is called.
+            loss = self._strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_loss, None)
+            num_words = {
+                k:self._strategy.reduce(tf.distribute.ReduceOp.SUM, v, None)
+                for k, v in six.iteritems(per_replica_words)}
+            return loss, num_words, False
+
+          def _skip():
+            loss = tf.constant(0, dtype=tf.float32)
+            num_words = {}
+            if "length" in per_replica_source:
+              num_words["source"] = tf.constant(0, dtype=tf.int32)
+            if "length" in per_replica_target:
+              num_words["target"] = tf.constant(0, dtype=tf.int32)
+            return loss, num_words, True
+
+          # We verify here that each replica receives a non empty batch. If not,
+          # we skip this iteration. This typically happens at the last iteration
+          # when training on a finite dataset.
+          # TODO: is there a simpler way to handle this case?
+          per_replica_non_empty_batch = self._strategy.experimental_run_v2(
+              lambda tensor: tf.math.count_nonzero(tf.shape(tensor)[0]),
+              args=(tf.nest.flatten(per_replica_source)[0],))
+          non_empty_batch_count = self._strategy.reduce(
+              tf.distribute.ReduceOp.SUM, per_replica_non_empty_batch, None)
+          return tf.cond(
+              tf.math.equal(non_empty_batch_count, self._strategy.num_replicas_in_sync),
+              true_fn=_run,
+              false_fn=_skip)
 
     @tf.function
     def _step():
@@ -147,7 +172,12 @@ def _step():
         self._checkpoint.save(0)
       self._model.visualize(self._checkpoint.model_dir)
 
-      for i, (loss, num_words) in enumerate(_forward()):  # pylint: disable=no-value-for-parameter
+      for i, (loss, num_words, skipped) in enumerate(_forward()):  # pylint: disable=no-value-for-parameter
+        if skipped:
+          # We assume only the last partial batch can possibly be skipped.
+          tf.get_logger().warning("Batch %d is partial, i.e. some training replicas "
+                                  "received an empty batch as input. Skipping.", i + 1)
+          break
         if tf.math.is_nan(loss):
           raise RuntimeError("Model diverged with loss = NaN.")
         if i == 0 or (i + 1) % accum_steps == 0:


Test patch:
diff --git a/opennmt/tests/runner_test.py b/opennmt/tests/runner_test.py
--- a/opennmt/tests/runner_test.py
+++ b/opennmt/tests/runner_test.py
@@ -111,7 +111,8 @@ def testTrainDistribute(self):
         "train": {
             "batch_size": 2,
             "length_bucket_width": None,
-            "max_step": 145002  # Just train for 2 steps.
+            "max_step": 145003,
+            "single_pass": True,  # Test we do not fail when a batch is missing for a replica.
         }
     }
     runner = self._getTransliterationRunner(config)