Repository: horovod/horovod

Issue:
MetricAverageCallback() broken with TensorFlow 2.3
**Environment:**
1. Framework: TensorFlow
2. Framework version: 2.3.0
3. Horovod version: 0.20.3
4. MPI version: 4.0.3
5. ~~CUDA version:~~
6. ~~NCCL version:~~
7. Python version: 3.6.8
8. ~~Spark / PySpark version:~~
9. OS and version: Red Hat Enterprise Linux release 8.3 (Ootpa)
10. GCC version: 9.2.1 20191120
11. CMake version: 3.11.4

**Bug report:**

`horovod.tensorflow.keras.callbacks.MetricAverageCallback()` is broken with TensorFlow 2.3.0. The logs displayed to stdout during the training are different for every worker, and if every worker logs to TensorBoard, the values also differ.

Minimal example script:
```python
# train.py
import numpy as np
import tensorflow as tf
import horovod.tensorflow.keras as hvd
hvd.init()
# Create toy dataset.
n = 1024
X = np.random.randn(n, 128)
Y = np.random.randn(n, 32)
# Distributed optimizer.
optimizer = tf.keras.optimizers.Adam()
optimizer = hvd.DistributedOptimizer(optimizer)
# Build model.
input_layer = tf.keras.layers.Input(shape=(128,))
next_layer = tf.keras.layers.Dense(128, activation='relu')(input_layer)
next_layer = tf.keras.layers.Dense(64, activation='relu')(next_layer)
next_layer = tf.keras.layers.Dense(64, activation='relu')(next_layer)
output_layer = tf.keras.layers.Dense(32)(next_layer)
model = tf.keras.models.Model(inputs=input_layer, outputs=output_layer)
model.compile(optimizer=optimizer, loss="mse", experimental_run_tf_function=False)
# Training.
train_history = model.fit(
    X, Y,
    batch_size=256,
    epochs=10,
    callbacks=[hvd.callbacks.BroadcastGlobalVariablesCallback(0),
               hvd.callbacks.MetricAverageCallback()],
    verbose=2)
print(train_history.history)
```

Training with two workers:
```
$ horovodrun -np 2 python3 train.py 
[1,1]<stdout>:Epoch 1/10
[1,0]<stdout>:Epoch 1/10
[1,1]<stdout>:4/4 - 0s - loss: 1.1202
[1,0]<stdout>:4/4 - 0s - loss: 1.1049
[1,1]<stdout>:Epoch 2/10
[1,0]<stdout>:Epoch 2/10
[1,1]<stdout>:4/4 - 0s - loss: 1.0446
[1,0]<stdout>:4/4 - 0s - loss: 1.0419
[1,1]<stdout>:Epoch 3/10
[1,0]<stdout>:Epoch 3/10
[1,1]<stdout>:4/4 - 0s - loss: 1.0196
[1,0]<stdout>:4/4 - 0s - loss: 1.0168
[1,0]<stdout>:Epoch 4/10
[1,1]<stdout>:Epoch 4/10
[1,1]<stdout>:4/4 - 0s - loss: 1.0082
[1,0]<stdout>:4/4 - 0s - loss: 1.0057
[1,1]<stdout>:Epoch 5/10
[1,0]<stdout>:Epoch 5/10
[1,1]<stdout>:4/4 - 0s - loss: 1.0016
[1,0]<stdout>:4/4 - 0s - loss: 0.9990
[1,0]<stdout>:Epoch 6/10
[1,1]<stdout>:Epoch 6/10
[1,0]<stdout>:4/4 - 0s - loss: 0.9946
[1,1]<stdout>:4/4 - 0s - loss: 0.9963
[1,1]<stdout>:Epoch 7/10
[1,0]<stdout>:Epoch 7/10
[1,1]<stdout>:4/4 - 0s - loss: 0.9922
[1,0]<stdout>:4/4 - 0s - loss: 0.9904
[1,1]<stdout>:Epoch 8/10
[1,0]<stdout>:Epoch 8/10
[1,1]<stdout>:4/4 - 0s - loss: 0.9882
[1,0]<stdout>:4/4 - 0s - loss: 0.9866
[1,1]<stdout>:Epoch 9/10
[1,0]<stdout>:Epoch 9/10
[1,1]<stdout>:4/4 - 0s - loss: 0.9845
[1,0]<stdout>:4/4 - 0s - loss: 0.9828
[1,1]<stdout>:Epoch 10/10
[1,0]<stdout>:Epoch 10/10
[1,1]<stdout>:4/4 - 0s - loss: 0.9806
[1,0]<stdout>:4/4 - 0s - loss: 0.9788
[1,1]<stdout>:{'loss': [1.1125493, 1.0432737, 1.0181885, 1.0069536, 1.0002735, 0.99541533, 0.9913291, 0.9873961, 0.98364025, 0.97974634]}
[1,0]<stdout>:{'loss': [1.1125493, 1.0432737, 1.0181885, 1.0069536, 1.0002735, 0.99541533, 0.9913291, 0.9873961, 0.98364025, 0.97974634]}
```
The losses reported by each worker differ, however the losses stored in `model.history` are averaged.
I think this is related to https://github.com/tensorflow/tensorflow/issues/41851, where it is said:

> Some callbacks (e.g. ProgbarLogger, ModelCheckpoint, ...) have the flag self._supports_tf_logs = True. If other callbacks (especially custom Callback) don't have this property, then those callbacks do not have acces to the same logs.

Adding `self._supports_tf_logs = True` in the `__init__` method of `MetricAverageCallbackImpl` in `horovod/_keras/callbacks.py` fixes the issue. However, I am not sure about the consequences in other use cases, and I assume that other callbacks may be broken, too.



Discussion:
Hey @loic-ehrhardt, thanks for reporting this.  I looked into the issue, and tried applying the `self._supports_tf_logs` solution you described.  However, this had the strange effect of fixing the logs at the cost of breaking the history.  Do you observe similar behavior?

Here's my simple test script to prove out this behavior:

```
def test():
        opt = tf.keras.optimizers.Adam(0.01)
        opt = hvd.DistributedOptimizer(opt)

        def test_metric(y_true, y_pred):
            return hvd.rank()

        model = keras.models.Sequential()
        model.add(keras.layers.Dense(2, input_shape=(3,)))
        model.compile(loss=keras.losses.mean_squared_error,
                      optimizer=opt,
                      metrics=[test_metric],
                      experimental_run_tf_function=False)

        x = np.random.random((1, 3))
        y = np.random.random((1, 3, 2))

        callbacks = [
            hvd.callbacks.BroadcastGlobalVariablesCallback(0),
            hvd.callbacks.MetricAverageCallback(),
        ]

        train_history = model.fit(
            x,
            y,
            steps_per_epoch=10,
            callbacks=callbacks,
            epochs=1
        )

        expected = sum(range(hvd.size())) / hvd.size()
        results = train_history.history.get('test_metric')

        assert results is not None, 'no results for test_metric'
        assert len(results) == 1
        assert results[0] == expected
```
Hi @tgaddair, thanks for looking into this. Oh no! This is indeed the case. I think this is because the two `logs` dictionaries (the ones that the callback sees if `self._supports_tf_logs` is `True` and the other ones that the callback sees if `self._supports_tf_logs` is `False`) are somewhat independent. I tried to add two `hvd.callbacks.MetricAverageCallback()` objects to the `callbacks`, one with and one without `self._supports_tf_logs` and this seems to work. It is very ugly though...
```python
# Training.
metric_average_callback_1 = hvd.callbacks.MetricAverageCallback()
metric_average_callback_1._supports_tf_logs = False
metric_average_callback_2 = hvd.callbacks.MetricAverageCallback()
metric_average_callback_2._supports_tf_logs = True
train_history = model.fit(
    X, Y,
    batch_size=256,
    epochs=10,
    callbacks=[hvd.callbacks.BroadcastGlobalVariablesCallback(0),
               metric_average_callback_1,
               metric_average_callback_2],
    verbose=2)
```
Is this working for you?
Thanks @loic-ehrhardt, it seems this works, but I agree it is not a good long term solution, particularly because this TensorFlow behavior could change at any time.  I think the right course of action here is to file this as an issue with TensorFlow and see if they can find a solution on their end.  Does that sound reasonable to you?
This sounds good, thanks again. Is there by chance a Horovod/TensorFlow compatibility table or is it mentioned somewhere the tested or advised TensorFlow version one should use with a given version of Horovod?

What about also logging a warning when a MetricAverageCallback object is instantiated with TF 2.3 that this may be broken?
Hey @loic-ehrhardt, we officially support all version of TF >= 1.15.  I agree that adding a warning about this is a good idea.  Are you interested in adding a PR for that?  Otherwise, we can of course add it to the backlog.
Sure, I'd be interested in creating a PR for this. I will do that ASAP.
PR: https://github.com/horovod/horovod/pull/2549
Bests, Loïc
Thanks!  I'll take a look.

Resolving patch:
diff --git a/horovod/_keras/callbacks.py b/horovod/_keras/callbacks.py
--- a/horovod/_keras/callbacks.py
+++ b/horovod/_keras/callbacks.py
@@ -54,7 +54,7 @@ def __init__(self, backend, device='', *args):
         self.allreduce_ops = {}
         self.device = device
 
-        if LooseVersion(tf.__version__) >= LooseVersion("2.3"):
+        if LooseVersion("2.3") <= LooseVersion(tf.__version__) < LooseVersion("2.5"):
             warnings.warn(
                 "Some callbacks may not have access to the averaged metrics, "
                 "see https://github.com/horovod/horovod/issues/2440")


Test patch: