Repository: horovod/horovod

Issue:
error when set backward_passes_per_step > 1 in tensorflow keras
**Environment:**
1. Framework: (TensorFlow, Keras, PyTorch, MXNet) 
2. Framework version: (1.15, 2.3.1, 1.6.0, 1.6.0)
3. Horovod version: 0.21.3
4. MPI version:
5. CUDA version:10.0
6. NCCL version:
7. Python version:3.7
8. Spark / PySpark version:
9. Ray version:
10. OS and version:
11. GCC version:
12. CMake version:

**Bug report:**
when set backward_passes_per_step > 1 and sparse_as_dense True , execute mpirun -np 2 --allow-run-as-root python test.py   it output error
here is code 
 ```python
from __future__ import print_function
import numpy as np
from tqdm import tqdm
from bert4keras.backend import keras, K
from bert4keras.layers import Loss
from bert4keras.models import build_transformer_model
from bert4keras.tokenizers import Tokenizer, load_vocab
from bert4keras.optimizers import Adam
from bert4keras.snippets import sequence_padding, open
from bert4keras.snippets import DataGenerator, AutoRegressiveDecoder
from keras.models import Model
from rouge import Rouge  # pip install rouge
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
import pdb
import os
import tensorflow as tf
import keras
import horovod.tensorflow.keras as hvd

hvd.init()

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = str(hvd.local_rank())
K.set_session(tf.Session(config=config))

opt = Adam(1e-5 * hvd.size())
opt = hvd.DistributedOptimizer(opt, backward_passes_per_step=2, sparse_as_dense=True)
```


output error log :

```bash
More specifically: Substructure "type=IndexedSlices str=IndexedSlices(indices=Tensor("training/Adam/gradients/Embedding-Segment/embedding_lookup_grad/Reshape_1:0", shape=(?,), dtype=int32), values=Tensor("training/Adam/gradients/Embedding-Segment/embedding_lookup_grad/Reshape:0", shape=(?, 768), dtype=float32), dense_shape=Tensor("training/Adam/gradients/Embedding-Segment/embedding_lookup_grad/Const:0", shape=(2,), dtype=int32))" is a sequence, while substructure "type=Tensor str=Tensor("training/Adam/cond_1/truediv_1:0", shape=(2, 768), dtype=float32)" is not
Entire first structure:
[., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., .]
Entire second structure:
[., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., ., .]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "test.py", line 223, in <module>
    verbose=1 if hvd.rank() == 0 else 0
  File "/usr/local/lib/python3.7/dist-packages/keras/legacy/interfaces.py", line 91, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1732, in fit_generator
    initial_epoch=initial_epoch)
  File "/usr/local/lib/python3.7/dist-packages/keras/engine/training_generator.py", line 42, in fit_generator
    model._make_train_function()
  File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 316, in _make_train_function
    loss=self.total_loss)
  File "/usr/local/lib/python3.7/dist-packages/keras/legacy/interfaces.py", line 91, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/keras/optimizers.py", line 504, in get_updates
    grads = self.get_gradients(loss, params)
  File "/usr/local/lib/python3.7/dist-packages/horovod/_keras/__init__.py", line 106, in get_gradients
    return self._allreduce(gradients, params)
  File "/usr/local/lib/python3.7/dist-packages/horovod/_keras/__init__.py", line 121, in _allreduce
    return self._agg_helper.compute_gradients(tuple(grads), tuple(vars))
  File "/usr/local/lib/python3.7/dist-packages/horovod/tensorflow/gradient_aggregation.py", line 208, in compute_gradients
    lambda: grads,
  File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/ops/control_flow_ops.py", line 1260, in cond
    "Incompatible return values of true_fn and false_fn: {}".format(e))
ValueError: Incompatible return values of true_fn and false_fn: The two structures don't have the same nested structure.
```





Discussion:
@rb-determined-ai could you take a look?
Gladly.  I'll try to dig into it some today or tomorrow but it might take me until next week... busy week this week.
Forgive me for the delay.  I am now up-to-speed with running horovod locally and would be happy to help out with fixing this.

I noticed that the example provided does not seem to include enough code to reproduce the error... looks from the stacktrace that it happens in model.fit().

It would help me out if you could share enough of an example that I could reproduce it.

But, you are also free to just go edit your `/usr/local/lib/python3.7/dist-packages/horovod/tensorflow/gradient_aggregation.py` file.  Based on reading through the feature I think the solution might be something like this:

```diff
            # Allreduce locally aggregated gradients when the counter is equivalent to
            # `backward_passes_per_step`. This the condition is true, it also resets
            # the counter back to 0.
            allreduced_grads = tf.cond(
                tf.equal(self.counter, self.backward_passes_per_step),
                lambda: self._allreduce_grads_helper(grads, vars),
-               lambda: grads,
+               lambda: [tf.convert_to_tensor(g) if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices) else g for g in grads],
            )
```

Resolving patch:
diff --git a/horovod/tensorflow/__init__.py b/horovod/tensorflow/__init__.py
--- a/horovod/tensorflow/__init__.py
+++ b/horovod/tensorflow/__init__.py
@@ -475,17 +475,19 @@ def compute_gradients(self, *args, **kwargs):
                 avg_grads = self._allreduce_grads(grads, vars)
             return list(zip(avg_grads, vars))
 
-        def apply_gradients(self, *args, **kwargs):
+        def apply_gradients(self, grads_and_vars, global_step=None, name=None):
             """Calls this same method on the underlying optimizer."""
             if self._agg_helper:
                 return self._agg_helper.apply_gradients(
-                    lambda: self._optimizer.apply_gradients(*args, **kwargs),
+                    lambda: self._optimizer.apply_gradients(
+                        grads_and_vars, global_step=global_step, name=name),
                     self._optimizer,
-                    *args,
-                    **kwargs,
+                    grads_and_vars,
+                    global_step=global_step,
+                    name=name,
                 )
-
-            return self._optimizer.apply_gradients(*args, **kwargs)
+            return self._optimizer.apply_gradients(
+                grads_and_vars, global_step=global_step, name=name)
 
         def get_slot(self, *args, **kwargs):
             """Calls this same method on the underlying optimizer."""
diff --git a/horovod/tensorflow/gradient_aggregation.py b/horovod/tensorflow/gradient_aggregation.py
--- a/horovod/tensorflow/gradient_aggregation.py
+++ b/horovod/tensorflow/gradient_aggregation.py
@@ -63,6 +63,20 @@ def __init__(
         self.not_none_indexes = {}
         self.num_none_grad_updates = 0
 
+    def _maybe_convert_grad(self, grad):
+        # Handle IndexedSlices.
+        if isinstance(grad, tf.IndexedSlices):
+            if self.sparse_as_dense:
+                return tf.convert_to_tensor(grad)
+            else:
+                raise ValueError(
+                    "IndexedSlices are not supported when "
+                    "`backward_passes_per_step` > 1 and "
+                    "`sparse_as_dense` is False."
+                )
+
+        return grad
+
     def _init_aggregation_vars(self, grads):
         """
         Initializes the counter that is used when to communicate and aggregate gradients
@@ -76,15 +90,7 @@ def _init_aggregation_vars(self, grads):
                 collections=[tf.compat.v1.GraphKeys.LOCAL_VARIABLES],
             )
             for idx, grad in enumerate(grads):
-                # Handle IndexedSlices.
-                if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):
-                    grad = tf.convert_to_tensor(grad)
-                elif isinstance(grad, tf.IndexedSlices):
-                    raise ValueError(
-                        "IndexedSlices are not supported when "
-                        "`backward_passes_per_step` > 1 and "
-                        "`sparse_as_dense` is False."
-                    )
+                grad = self._maybe_convert_grad(grad)
 
                 # Handle grads that are None.
                 if grad is None:
@@ -140,7 +146,7 @@ def _aggregate_grads(self, grads):
 
         return aggregation_ops_list
 
-    def _allreduce_grads_helper(self, grads, vars):
+    def _allreduce_grads_helper(self, vars):
         # Read in latest variables values.
         aggregated_grads = []
         aggregation_read_ops_list = []
@@ -204,8 +210,8 @@ def compute_gradients(self, grads, vars):
             # the counter back to 0.
             allreduced_grads = tf.cond(
                 tf.equal(self.counter, self.backward_passes_per_step),
-                lambda: self._allreduce_grads_helper(grads, vars),
-                lambda: grads,
+                lambda: self._allreduce_grads_helper(vars),
+                lambda: [self._maybe_convert_grad(g) for g in grads]
             )
 
             # Handle case where there is only one variable.
@@ -250,7 +256,7 @@ def increment_optimizer_iteration():
         # still want to increment the global step if it is being tracked
         # (e.g., Tensorflow Estimators).
         def increment_global_step_counter():
-            global_step_counter = tf.compat.v1.train.get_global_step()
+            global_step_counter = kwargs.get('global_step')
             if global_step_counter is None:
                 return tf.no_op()
             return global_step_counter.assign_add(


Test patch:
diff --git a/test/parallel/test_tensorflow.py b/test/parallel/test_tensorflow.py
--- a/test/parallel/test_tensorflow.py
+++ b/test/parallel/test_tensorflow.py
@@ -2828,6 +2828,23 @@ def test_horovod_syncbn_cpu(self):
                 self.assertAllClose(self.evaluate(sync_bn.moving_mean), self.evaluate(bn.moving_mean))
                 self.assertAllClose(self.evaluate(sync_bn.moving_variance), self.evaluate(bn.moving_variance))
 
+    def _grad_agg_compute_expected_value(self, backward_passes_per_step, batch_id):
+        sum_per_aggregation = 0.0
+        for _ in range(backward_passes_per_step):
+            grads_for_batch = 0.0
+            for rank in range(hvd.size()):
+                grads_for_batch += rank
+
+            # Apply `average_aggregated_gradients`.
+            grads_for_batch /= float(backward_passes_per_step)
+
+            # Averages across workers.
+            sum_per_aggregation += grads_for_batch / float(hvd.size())
+
+        aggregations_completed = math.floor((batch_id + 1) / backward_passes_per_step)
+        return aggregations_completed * sum_per_aggregation
+
+
     def test_local_gradient_aggregation(self):
         """Test that local gradient aggregation works as expected."""
 
@@ -2858,9 +2875,14 @@ def _apply_dense(self, grad, var):
 
             grads_and_vars = opt.compute_gradients()
 
+            # Use custom variable for global_step instead of default one from
+            # `tf.compat.v1.train.get_or_create_global_step` to ensure
+            # the correct counter is being incremented.
+            global_step = tf.compat.v1.get_variable(
+                'test_global_step', initializer=tf.constant(0, dtype=tf.int64),
+                trainable=False)
             # Check the global_step before the update_op to ensure that the global step is
             # well-ordered in the graph, occuring within apply_gradients().
-            global_step = tf.compat.v1.train.get_or_create_global_step()
             with tf.compat.v1.control_dependencies(g for g, _ in grads_and_vars):
                 # Use +0 instead of tf.identity() since tf.identity() has some weird semantics
                 # with control_dependencies: github.com/tensorflow/tensorflow/issues/4663
@@ -2871,31 +2893,59 @@ def _apply_dense(self, grad, var):
             sess.run(tf.compat.v1.global_variables_initializer())
             sess.run(tf.compat.v1.local_variables_initializer())
 
-            def compute_expected_value(batch_id):
-                sum_per_aggregation = 0.0
-                for _ in range(backward_passes_per_step):
-                    grads_for_batch = 0.0
-                    for rank in range(hvd.size()):
-                        grads_for_batch += rank
-
-                    # Apply `average_aggregated_gradients`.
-                    grads_for_batch /= float(backward_passes_per_step)
-
-                    # Averages across workers.
-                    sum_per_aggregation += grads_for_batch / float(hvd.size())
-
-                aggregations_completed = math.floor((batch_id + 1) / backward_passes_per_step)
-                return aggregations_completed * sum_per_aggregation
-
             steps_seen = []
             for idx in range(10):
                 step, _ = sess.run((global_step_before_update, update_op))
                 computed_value = sess.run(opt._optimizer.variable.read_value())[0]
                 steps_seen.append(step)
-                assert computed_value == compute_expected_value(idx)
+                self.assertEquals(computed_value, self._grad_agg_compute_expected_value(
+                    backward_passes_per_step, idx))
 
             assert steps_seen == list(range(10))
 
+    def test_local_gradient_aggregation_sparse(self):
+        """Test the the gradient aggregation works for sparse gradients."""
+
+        if _executing_eagerly() or not hasattr(tf, 'ConfigProto'):
+            pytest.skip("Gradient aggregation for Legacy Optimizer only support graph mode.")
+
+        hvd.init()
+        with self.test_session(config=config) as sess:
+            class SparseGradientOptimizer(tf.compat.v1.train.Optimizer):
+                def __init__(self, *args, **kwargs):
+                    super().__init__(*args, **kwargs)
+                    self.variable = tf.Variable([0.0])
+                    self.grad = tf.IndexedSlices(
+                        tf.constant([float(hvd.rank())]), tf.constant([0]), tf.constant([1])
+                    )
+
+                def compute_gradients(self, *args, **kwargs):
+                    return [(self.grad, self.variable)]
+
+                def _apply_dense(self, grad, var):
+                    return var.assign_add(grad)
+
+            backward_passes_per_step = 4
+            opt = SparseGradientOptimizer(name="sparse_test", use_locking=False)
+            opt = hvd.DistributedOptimizer(
+                optimizer=opt,
+                backward_passes_per_step=backward_passes_per_step,
+                average_aggregated_gradients=True,
+                sparse_as_dense=True,
+            )
+
+            grads_and_vars = opt.compute_gradients()
+            update_op = opt.apply_gradients(grads_and_vars)
+
+            sess.run(tf.compat.v1.global_variables_initializer())
+            sess.run(tf.compat.v1.local_variables_initializer())
+
+            for idx in range(10):
+                sess.run(update_op)
+                computed_value = sess.run(opt._optimizer.variable.read_value())[0]
+                self.assertEquals(computed_value, self._grad_agg_compute_expected_value(
+                    backward_passes_per_step, idx))
+
 
 from tensorflow.python.framework.test_util import run_all_in_graph_and_eager_modes
 run_all_in_graph_and_eager_modes(TensorFlowTests)