Repository: OpenNMT/OpenNMT-tf

Issue:
"Horovod has been shut down" error when training is finished due to early stopping
Hello,

I'm training a Transformer model using OpenNMT-tf 2.29.0 with Horovod on a single node with 4 GPUs, but if the training is finished due to early stopping (in the example below, early stopping kicked in at 600k steps but 1M steps was configured), Horovod complains that it has been shut down and the script exits with a non-zero status.

This is in itself not a huge problem since the model is still saved and exported properly, but this causes troubles when using OpenNMT in Azure ML workspace where the training might be a step followed by others.

I tried to look into it, and the checkpoint averaging triggered here https://github.com/OpenNMT/OpenNMT-tf/blob/v2.29.0/opennmt/runner.py#L304-L307 seems to be done properly, but I don't understand why training steps and gradients computation would still be performed after the checkpoint averaging, but I'm no Horovod expert :)

```
[1,0]<stderr>:2022-10-08 05:23:59.894000: I runner.py:362] Restored checkpoint ./logs/ckpt-600000
[1,0]<stderr>:INFO:tensorflow:Restored checkpoint ./logs/ckpt-600000
[1,0]<stderr>:2022-10-08 05:24:02.028000: I runner.py:365] Averaging 5 checkpoints...
[1,0]<stderr>:INFO:tensorflow:Averaging 5 checkpoints...
[1,0]<stderr>:2022-10-08 05:24:02.028000: I runner.py:365] Reading checkpoint ./logs/ckpt-560000...
[1,0]<stderr>:INFO:tensorflow:Reading checkpoint ./logs/ckpt-560000...
[1,0]<stderr>:2022-10-08 05:24:02.496000: I runner.py:365] Reading checkpoint ./logs/ckpt-570000...
[1,0]<stderr>:INFO:tensorflow:Reading checkpoint ./logs/ckpt-570000...
[1,0]<stderr>:2022-10-08 05:24:02.944000: I runner.py:365] Reading checkpoint ./logs/ckpt-580000...
[1,0]<stderr>:INFO:tensorflow:Reading checkpoint ./logs/ckpt-580000...
[1,0]<stderr>:2022-10-08 05:24:03.396000: I runner.py:365] Reading checkpoint ./logs/ckpt-590000...
[1,0]<stderr>:INFO:tensorflow:Reading checkpoint ./logs/ckpt-590000...
[1,0]<stderr>:2022-10-08 05:24:03.844000: I runner.py:365] Reading checkpoint ./logs/ckpt-600000...
[1,0]<stderr>:INFO:tensorflow:Reading checkpoint ./logs/ckpt-600000...
[1,0]<stderr>:2022-10-08 05:24:06.440000: I runner.py:365] Saved averaged checkpoint to ./logs/avg/ckpt-600000
[1,0]<stderr>:INFO:tensorflow:Saved averaged checkpoint to ./logs/avg/ckpt-600000
[1,1]<stdout>:9913eb9c2e4c4db3bd484228cfe17c1a000000:160:368 [1] NCCL INFO comm 0x7f252b96a910 rank 1 nranks 4 cudaDev 1 busId 70f800000 - Destroy COMPLETE
[1,3]<stdout>:9913eb9c2e4c4db3bd484228cfe17c1a000000:162:366 [3] NCCL INFO comm 0x7f6dcf96a340 rank 3 nranks 4 cudaDev 3 busId a66000000 - Destroy COMPLETE
[1,2]<stdout>:9913eb9c2e4c4db3bd484228cfe17c1a000000:161:362 [2] NCCL INFO comm 0x7f9b5f96a060 rank 2 nranks 4 cudaDev 2 busId 973200000 - Destroy COMPLETE
[1,0]<stdout>:9913eb9c2e4c4db3bd484228cfe17c1a000000:159:355 [0] NCCL INFO comm 0x7f9a7b99eef0 rank 0 nranks 4 cudaDev 0 busId 4f5100000 - Destroy COMPLETE
[1,3]<stderr>:Traceback (most recent call last):
[1,3]<stderr>:  File "/usr/local/bin/onmt-main", line 8, in <module>
[1,3]<stderr>:    sys.exit(main())
[1,3]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/bin/main.py", line 318, in main
[1,3]<stderr>:    runner.train(
[1,3]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/runner.py", line 289, in train
[1,3]<stderr>:    summary = trainer(
[1,3]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 109, in __call__
[1,3]<stderr>:    for i, loss in enumerate(
[1,3]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 260, in _steps
[1,3]<stderr>:    step_fn()
[1,3]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
[1,3]<stderr>:    raise e.with_traceback(filtered_tb) from None
[1,3]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/execute.py", line 54, in quick_execute
[1,3]<stderr>:    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
[1,3]<stderr>:tensorflow.python.framework.errors_impl.UnknownError: Graph execution error:
[1,3]<stderr>:
[1,3]<stderr>:Detected at node 'HorovodAllreduce_ReadVariableOp_22_0' defined at (most recent call last):
[1,3]<stderr>:    File "/usr/local/bin/onmt-main", line 8, in <module>
[1,3]<stderr>:      sys.exit(main())
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/bin/main.py", line 318, in main
[1,3]<stderr>:      runner.train(
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/runner.py", line 289, in train
[1,3]<stderr>:      summary = trainer(
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 109, in __call__
[1,3]<stderr>:      for i, loss in enumerate(
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 260, in _steps
[1,3]<stderr>:      step_fn()
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 239, in _step
[1,3]<stderr>:      return self._step()
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 329, in _step
[1,3]<stderr>:      self._apply_gradients(self._gradient_accumulator.gradients)
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 381, in _apply_gradients
[1,3]<stderr>:      return super()._apply_gradients(map(self._all_reduce_sum, gradients))
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 310, in _apply_gradients
[1,3]<stderr>:      self._optimizer.apply_gradients(
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 388, in _all_reduce_sum
[1,3]<stderr>:      return self._hvd.allreduce(value, op=self._hvd.Sum)
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/horovod/tensorflow/__init__.py", line 125, in allreduce
[1,3]<stderr>:      summed_tensor_compressed = _allreduce(tensor_compressed, op=op,
[1,3]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/horovod/tensorflow/mpi_ops.py", line 127, in _allreduce
[1,3]<stderr>:      return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op,
[1,3]<stderr>:    File "<string>", line 107, in horovod_allreduce
[1,3]<stderr>:Node: 'HorovodAllreduce_ReadVariableOp_22_0'
[1,3]<stderr>:Horovod has been shut down. This was caused by an exception on one of the ranks or an attempt to allreduce, allgather or broadcast a tensor after one of the ranks finished execution. If the shutdown was caused by an exception, you should see the exception in the log before the first shutdown message.
[1,3]<stderr>:	 [[{{node HorovodAllreduce_ReadVariableOp_22_0}}]] [Op:__inference__step_50022]
[1,2]<stderr>:Traceback (most recent call last):
[1,2]<stderr>:  File "/usr/local/bin/onmt-main", line 8, in <module>
[1,2]<stderr>:    sys.exit(main())
[1,2]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/bin/main.py", line 318, in main
[1,2]<stderr>:    runner.train(
[1,2]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/runner.py", line 289, in train
[1,2]<stderr>:    summary = trainer(
[1,2]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 109, in __call__
[1,2]<stderr>:    for i, loss in enumerate(
[1,2]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 260, in _steps
[1,2]<stderr>:    step_fn()
[1,2]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
[1,2]<stderr>:    raise e.with_traceback(filtered_tb) from None
[1,2]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/execute.py", line 54, in quick_execute
[1,2]<stderr>:    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
[1,2]<stderr>:tensorflow.python.framework.errors_impl.UnknownError: Graph execution error:
[1,2]<stderr>:
[1,2]<stderr>:Detected at node 'HorovodAllreduce_ReadVariableOp_165_0' defined at (most recent call last):
[1,2]<stderr>:    File "/usr/local/bin/onmt-main", line 8, in <module>
[1,2]<stderr>:      sys.exit(main())
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/bin/main.py", line 318, in main
[1,2]<stderr>:      runner.train(
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/runner.py", line 289, in train
[1,2]<stderr>:      summary = trainer(
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 109, in __call__
[1,2]<stderr>:      for i, loss in enumerate(
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 260, in _steps
[1,2]<stderr>:      step_fn()
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 239, in _step
[1,2]<stderr>:      return self._step()
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 329, in _step
[1,2]<stderr>:      self._apply_gradients(self._gradient_accumulator.gradients)
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 381, in _apply_gradients
[1,2]<stderr>:      return super()._apply_gradients(map(self._all_reduce_sum, gradients))
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 310, in _apply_gradients
[1,2]<stderr>:      self._optimizer.apply_gradients(
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 388, in _all_reduce_sum
[1,2]<stderr>:      return self._hvd.allreduce(value, op=self._hvd.Sum)
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/horovod/tensorflow/__init__.py", line 125, in allreduce
[1,2]<stderr>:      summed_tensor_compressed = _allreduce(tensor_compressed, op=op,
[1,2]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/horovod/tensorflow/mpi_ops.py", line 127, in _allreduce
[1,2]<stderr>:      return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op,
[1,2]<stderr>:    File "<string>", line 107, in horovod_allreduce
[1,2]<stderr>:Node: 'HorovodAllreduce_ReadVariableOp_165_0'
[1,2]<stderr>:Horovod has been shut down. This was caused by an exception on one of the ranks or an attempt to allreduce, allgather or broadcast a tensor after one of the ranks finished execution. If the shutdown was caused by an exception, you should see the exception in the log before the first shutdown message.
[1,2]<stderr>:	 [[{{node HorovodAllreduce_ReadVariableOp_165_0}}]] [Op:__inference__step_50013]
[1,1]<stderr>:Traceback (most recent call last):
[1,1]<stderr>:  File "/usr/local/bin/onmt-main", line 8, in <module>
[1,1]<stderr>:    sys.exit(main())
[1,1]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/bin/main.py", line 318, in main
[1,1]<stderr>:    runner.train(
[1,1]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/runner.py", line 289, in train
[1,1]<stderr>:    summary = trainer(
[1,1]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 109, in __call__
[1,1]<stderr>:    for i, loss in enumerate(
[1,1]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 260, in _steps
[1,1]<stderr>:    step_fn()
[1,1]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
[1,1]<stderr>:    raise e.with_traceback(filtered_tb) from None
[1,1]<stderr>:  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/execute.py", line 54, in quick_execute
[1,1]<stderr>:    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
[1,1]<stderr>:tensorflow.python.framework.errors_impl.UnknownError: Graph execution error:
[1,1]<stderr>:
[1,1]<stderr>:Detected at node 'HorovodAllreduce_ReadVariableOp_127_0' defined at (most recent call last):
[1,1]<stderr>:    File "/usr/local/bin/onmt-main", line 8, in <module>
[1,1]<stderr>:      sys.exit(main())
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/bin/main.py", line 318, in main
[1,1]<stderr>:      runner.train(
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/runner.py", line 289, in train
[1,1]<stderr>:      summary = trainer(
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 109, in __call__
[1,1]<stderr>:      for i, loss in enumerate(
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 260, in _steps
[1,1]<stderr>:      step_fn()
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 239, in _step
[1,1]<stderr>:      return self._step()
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 329, in _step
[1,1]<stderr>:      self._apply_gradients(self._gradient_accumulator.gradients)
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 381, in _apply_gradients
[1,1]<stderr>:      return super()._apply_gradients(map(self._all_reduce_sum, gradients))
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 310, in _apply_gradients
[1,1]<stderr>:      self._optimizer.apply_gradients(
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/opennmt/training.py", line 388, in _all_reduce_sum
[1,1]<stderr>:      return self._hvd.allreduce(value, op=self._hvd.Sum)
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/horovod/tensorflow/__init__.py", line 125, in allreduce
[1,1]<stderr>:      summed_tensor_compressed = _allreduce(tensor_compressed, op=op,
[1,1]<stderr>:    File "/usr/local/lib/python3.8/dist-packages/horovod/tensorflow/mpi_ops.py", line 127, in _allreduce
[1,1]<stderr>:      return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op,
[1,1]<stderr>:    File "<string>", line 107, in horovod_allreduce
[1,1]<stderr>:Node: 'HorovodAllreduce_ReadVariableOp_127_0'
[1,1]<stderr>:Horovod has been shut down. This was caused by an exception on one of the ranks or an attempt to allreduce, allgather or broadcast a tensor after one of the ranks finished execution. If the shutdown was caused by an exception, you should see the exception in the log before the first shutdown message.
[1,1]<stderr>:	 [[{{node HorovodAllreduce_ReadVariableOp_127_0}}]] [Op:__inference__step_49995]
```


Command used to launch the training:
`horovodrun -np 4 -H localhost:4 onmt-main --model document_context_model.py --config training_config.yml model_config.yml /mnt/azureml/cr/j/84cb509ce83e4ac59d39e94ee62b0ad2/cap/data-capability/wd/INPUT_input__6593fad7/data_conf.yml --data_dir /mnt/azureml/cr/j/84cb509ce83e4ac59d39e94ee62b0ad2/cap/data-capability/wd/INPUT_input__6593fad7 --auto_config --mixed_precision train --with_eval --horovod`


Discussion:
Hi,

The evaluation is only run in the master process, so the other workers are not aware that the early stopping condition is met. This information should somehow be communicated with the other workers (probably with `hvd.broadcast()`).

Do you want to look into that? The related code is located in `opennmt/training.py`. We could extend the method `_evaluate` in `HorovodTrainer` and implement additional logic the send the early stop flag to all workers.
Thanks a lot for the info @guillaumekln , I'll have a look then and will create a PR once it's working :)

Resolving patch:
diff --git a/opennmt/training.py b/opennmt/training.py
--- 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(


Test patch: