Repository: horovod/horovod

Issue:
Failed to use LossScaleOptimizer with DistributedOptimizer.
**Environment:**
Note that I checked this issue on NVIDIA's NGC TF2 container image, `nvcr.io/nvidia/tensorflow:21.05-tf2-py3`. (Release notes: https://docs.nvidia.com/deeplearning/frameworks/tensorflow-release-notes/rel_21-05.html#rel_21-05)

1. Framework: TensorFlow (and its keras API)
2. Framework version: 2.4.0+nv
3. Horovod version: 0.21.3
4. MPI version: 4.1.1rc1
5. CUDA version: V11.3.58
6. NCCL version: 2.9.8
7. Python version: 3.8.5
8. Spark / PySpark version: N/A
9. Ray version: N/A
10. OS and version: Ubuntu 20.04.2 LTS
11. GCC version: 9.3.0
12. CMake version: 3.16.3

**Checklist:**
1. Did you search issues to find if somebody asked this question before?: Y
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?: N/A
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)?: Y
4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)?: Y

**Bug report:**
Please describe erroneous behavior you're observing and steps to reproduce it.

When I tried to use newer mixed precision API, `tf.keras.mixed_precision.LossScaleOptimizer`, with Horovod, the error, `AttributeError: 'LossScaleOptimizer' object has no attribute '_optimizer'`, happened.

Minimal reproduction code is below, and if a line of `LossScaleOptimizer` is commented out, no error happens.

```python
import tensorflow as tf
import horovod.tensorflow.keras as hvd

def main():
    hvd.init()

    gpus = tf.config.experimental.list_physical_devices('GPU')
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)
    if gpus:
        tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')

    policy = tf.keras.mixed_precision.Policy('mixed_float16')
    tf.keras.mixed_precision.set_global_policy(policy)
    opt = tf.keras.optimizers.RMSprop(learning_rate=0.001 * (hvd.size()**0.5))
    opt = tf.keras.mixed_precision.LossScaleOptimizer(opt)
    opt = hvd.DistributedOptimizer(opt)

if __name__ == '__main__':
    main()
```

Full stacktrace message is below.

```
Traceback (most recent call last):
  File "minimal_repro.py", line 21, in <module>
    main()
  File "minimal_repro.py", line 17, in main
    opt = hvd.DistributedOptimizer(opt)
  File "/usr/local/lib/python3.8/dist-packages/horovod/tensorflow/keras/__init__.py", line 97, in DistributedOptimizer
    return _impl.create_distributed_optimizer(
  File "/usr/local/lib/python3.8/dist-packages/horovod/_keras/__init__.py", line 167, in create_distributed_optimizer
    return cls.from_config(config)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py", line 794, in from_config
    return cls(**config)
  File "/usr/local/lib/python3.8/dist-packages/horovod/_keras/__init__.py", line 37, in __init__
    self._name = name or "Distributed%s" % self.__class__.__base__.__name__
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py", line 918, in __setattr__
    if (name != '_optimizer' and name in self._optimizer._hyper
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py", line 888, in __getattribute__
    raise e
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py", line 884, in __getattribute__
    return object.__getattribute__(self, name)
AttributeError: 'LossScaleOptimizer' object has no attribute '_optimizer'
```

As far as I checked, at [`keras/mixed_precision/loss_scale_optimizer.py`#L534](https://github.com/tensorflow/tensorflow/blob/v2.4.0/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py#L534), `_optimizer` is initialized, and Horovod's automatically generated wrapper class, `_DistributedOptimizer`, tries to access `_name` attribute before this initialization ([`horovod/_keras/__init__.py`#L37](https://github.com/horovod/horovod/blob/master/horovod/_keras/__init__.py#L37)). Then, `LossScaleOptimizer` checks if an attribute exists or not when accessing to any attribute, and raise an error for a few attributes like `_optimizer` at ([`keras/mixed_precision/loss_scale_optimizer.py`#L882-L888](https://github.com/tensorflow/tensorflow/blob/v2.4.0/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py#L882-L888))
However, it looks like `_DistributedOptimizer` makes the initialization (= `_optimizer` initialization) delayed to the last of `__init__()` ([`horovod/_keras/__init__.py`#L69](https://github.com/horovod/horovod/blob/master/horovod/_keras/__init__.py#L69)). As a result, this error looks like caused.


Discussion:

Resolving patch:
diff --git a/horovod/_keras/__init__.py b/horovod/_keras/__init__.py
--- a/horovod/_keras/__init__.py
+++ b/horovod/_keras/__init__.py
@@ -34,6 +34,8 @@ class _DistributedOptimizer(keras.optimizers.Optimizer):
         _HAS_AGGREGATE_GRAD = True
 
         def __init__(self, **kwargs):
+            super(self.__class__, self).__init__(**kwargs)
+
             self._name = name or "Distributed%s" % self.__class__.__base__.__name__
             self._aggregated_gradients = False
 
@@ -66,8 +68,6 @@ def __init__(self, **kwargs):
                         optimizer_type=LocalGradientAggregationHelper._OPTIMIZER_TYPE_KERAS,
                     )
 
-            super(self.__class__, self).__init__(**kwargs)
-
         def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
             """
             Compute gradients of all trainable variables.


Test patch: