Repository: horovod/horovod
Pull request: Support backward_passes_per_step > 1 for TF Keras Graph Mode
## Checklist before submitting

- [x] Did you read the [contributor guide](https://github.com/horovod/horovod/blob/master/CONTRIBUTING.md)?
- [ ] Did you update the docs?
- [x] Did you write any tests to validate this change?  
- [x] Did you update the [CHANGELOG](https://github.com/horovod/horovod/blob/master/CHANGELOG.md), if this change affects users?

## Description
This PR adds support for `backward_passes_per_step > 1` for TF Keras optimizers executing in graph(non-eager) mode. This is one of the features that we have built into [Determined AI's](https://determined.ai/) fork of Horovod that we would like to upstream. After this PR we also plan to add support for `backward_passes_per_step > 1` for eager execution as well as Tensorflow optimizers (LegacyOptimizers).


Problem statement:
Tensorflow Keras support of Gradient Accumulation?
**Is your feature request related to a problem? Please describe.**

No, it's not related to an issue.

**Describe the solution you'd like**

Currently the API supports Gradient Accumulation in Tensorflow and Pytorch module, but not TF.keras.
The proposal is to support GA in tf.keras APIs.

**Describe alternatives you've considered**

Considered to use TF API to do GA, but model was in keras.

**Additional context**

Currently there is no options similar to `backward_passes_per_step` in tf/pytorch.

Hints from issue discussion:

Commit message: Improve TF1 Keras gradient aggregation test.

Signed-off-by: aaron276h <aaron@determined.ai>

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7baebd1b26..6e50679284 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for backward_passes_per_step > 1 for TF Keras graph mode. ([#2346](https://github.com/horovod/horovod/pull/2346))
+
 ### Changed
 
 ### Deprecated
diff --git a/examples/tensorflow/tensorflow_keras_mnist.py b/examples/tensorflow/tensorflow_keras_mnist.py
index fef4cffb23..ebf1d95df8 100644
--- a/examples/tensorflow/tensorflow_keras_mnist.py
+++ b/examples/tensorflow/tensorflow_keras_mnist.py
@@ -68,7 +68,7 @@
 opt = keras.optimizers.Adadelta(1.0 * hvd.size())
 
 # Horovod: add Horovod Distributed Optimizer.
-opt = hvd.DistributedOptimizer(opt)
+opt = hvd.DistributedOptimizer(opt, backward_passes_per_step=1)
 
 model.compile(loss=keras.losses.categorical_crossentropy,
               optimizer=opt,
diff --git a/horovod/_keras/__init__.py b/horovod/_keras/__init__.py
index 3fe1d37f8a..1d669c8f55 100644
--- a/horovod/_keras/__init__.py
+++ b/horovod/_keras/__init__.py
@@ -17,19 +17,23 @@
 
 import horovod.tensorflow as hvd
 import tensorflow as tf
+from horovod.common.gradient_aggregation import LocalGradientAggregationHelper
+from horovod.tensorflow.mpi_ops import rank
 
 
 _PRE_TF_2_4_0 = LooseVersion(tf.__version__) < LooseVersion('2.4.0')
 
 
 def create_distributed_optimizer(keras, optimizer, name, device_dense, device_sparse,
-                                 compression, sparse_as_dense, gradient_predivide_factor):
+                                 compression, sparse_as_dense, gradient_predivide_factor,
+                                 backward_passes_per_step=1, average_aggregated_gradients=True):
     class _DistributedOptimizer(keras.optimizers.Optimizer):
         _HAS_AGGREGATE_GRAD = True
 
         def __init__(self, **kwargs):
             self._name = name or "Distributed%s" % self.__class__.__base__.__name__
             self._aggregated_gradients = False
+
             self._allreduce_grads = hvd._make_allreduce_grads_fn(
                 self._name,
                 device_dense,
@@ -38,6 +42,22 @@ def __init__(self, **kwargs):
                 sparse_as_dense,
                 hvd.Average,
                 gradient_predivide_factor)
+
+            self._agg_helper = None
+            if backward_passes_per_step > 1:
+                if hvd._executing_eagerly():
+                    raise ValueError('backward_passes_per_step > 1 is not supported yet for '
+                                     'eager execution')
+
+                self._agg_helper = LocalGradientAggregationHelper(
+                    backward_passes_per_step=backward_passes_per_step,
+                    allreduce_func=self._allreduce_grads,
+                    sparse_as_dense=sparse_as_dense,
+                    average_aggregated_gradients=average_aggregated_gradients,
+                    rank=rank(),
+                    optimizer_type=LocalGradientAggregationHelper._OPTIMIZER_TYPE_KERAS,
+                )
+
             super(self.__class__, self).__init__(**kwargs)
 
         def get_gradients(self, loss, params):
@@ -63,15 +83,29 @@ def _aggregate_gradients(self, grads_and_vars):
 
         def _allreduce(self, grads):
             self._aggregated_gradients = True
-            return self._allreduce_grads(grads)
+
+            if self._agg_helper:
+                return self._agg_helper.compute_gradients(tuple(grads))
+            else:
+                return self._allreduce_grads(grads)
 
         def apply_gradients(self, *args, **kwargs):
-            results = super(self.__class__, self).apply_gradients(*args, **kwargs)
+            if self._agg_helper:
+                results = self._agg_helper.apply_gradients(
+                    lambda: super(self.__class__, self).apply_gradients(*args, **kwargs),
+                    self,
+                    *args,
+                    **kwargs,
+                )
+            else:
+                results = super(self.__class__, self).apply_gradients(*args, **kwargs)
+
             if not self._aggregated_gradients:
                 raise Exception('`apply_gradients()` was called without a call to '
                                 '`get_gradients()` or `_aggregate_gradients`. If you\'re '
                                 'using TensorFlow 2.0, please specify '
                                 '`experimental_run_tf_function=False` in `compile()`.')
+
             return results
 
     # We dynamically create a new class that inherits from the optimizer that was passed in.
@@ -80,6 +114,7 @@ def apply_gradients(self, *args, **kwargs):
     # model could be easily restored without Horovod.
     cls = type(optimizer.__class__.__name__, (optimizer.__class__,),
                dict(_DistributedOptimizer.__dict__))
+
     return cls.from_config(optimizer.get_config())
 
 
diff --git a/horovod/common/gradient_aggregation.py b/horovod/common/gradient_aggregation.py
new file mode 100644
index 0000000000..06b256a048
--- /dev/null
+++ b/horovod/common/gradient_aggregation.py
@@ -0,0 +1,266 @@
+import tensorflow as tf
+
+
+def apply_op_to_not_none_tensors(tensor_op, tensors, *args):
+    return [
+        tensor_op(
+            tensor,
+            *args
+        ) if tensor is not None else tensor for tensor in tensors]
+
+
+def get_not_none_from_list(tensor_list):
+    return [x for x in tensor_list if x is not None]
+
+
+class LocalGradientAggregationHelper:
+    """
+    LocalGradientAggregationHelper aggregates gradient updates locally,
+    and communicates the updates across machines only once every
+    backward_passes_per_step. Only supports graph mode execution.
+    """
+
+    _OPTIMIZER_TYPE_KERAS = "optimizer_type_keras"
+
+    def __init__(
+            self,
+            backward_passes_per_step,
+            allreduce_func,
+            sparse_as_dense,
+            average_aggregated_gradients,
+            rank,
+            optimizer_type):
+        self._allreduce_grads = allreduce_func
+
+        # backward_passes_per_step controls how often gradient updates are
+        # synchronized.
+        self.backward_passes_per_step = backward_passes_per_step
+        assert self.backward_passes_per_step > 0
+
+        # average_aggregated_gradients controls whether gradient updates that are
+        # aggregated, should be divided by `backward_passes_per_step`.
+        self.average_aggregated_gradients = average_aggregated_gradients
+
+        # This is going to be [N] data structure holding the aggregated gradient updates
+        # N is the number of parameters.
+        self.locally_aggregated_grads = []
+
+        # Used to know when to allreduce and apply gradients. We allreduce when `self.counter`
+        # is equal to `self.backward_passes_per_step`. We apply gradients when `self.counter` is
+        # equal to 0.
+        self.counter = None
+
+        self.sparse_as_dense = sparse_as_dense
+        self.rank = rank
+        self.optimizer_type = optimizer_type
+
+        # Contains the mapping of indexes of grad updates that are not None to their index in
+        # locally_aggregated_grads which only contains not None gradients. When performing
+        # gradient aggregation we have to remove them from the list of grads prior to passing
+        # the list into a tf.cond().
+        self.not_none_indexes = {}
+        self.num_none_grad_updates = 0
+
+    def _init_aggregation_vars(self, grads):
+        """
+        Initializes the counter that is used when to communicate and aggregate gradients
+        and the tensorflow variables that store the locally aggregated gradients.
+        """
+        variable_scope_name = "aggregation_variables_" + str(self.rank)
+        with tf.compat.v1.variable_scope(variable_scope_name, reuse=tf.compat.v1.AUTO_REUSE):
+            self.counter = tf.compat.v1.get_variable(
+                "aggregation_counter", shape=(), dtype=tf.int32,
+                trainable=False, initializer=tf.compat.v1.zeros_initializer(),
+            )
+            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 AssertionError(
+                        "IndexedSlices are not supported when "
+                        "`backward_passes_per_step` > 1 and "
+                        "`sparse_as_dense` is False."
+                    )
+
+                # Handle grads that are None.
+                if grad is None:
+                    self.num_none_grad_updates += 1
+                    continue
+                self.not_none_indexes[idx] = len(self.locally_aggregated_grads)
+
+                # Create shadow variable.
+                grad_aggregation_variable_name = str(idx)
+                grad_aggregation_variable = tf.compat.v1.get_variable(
+                    grad_aggregation_variable_name,
+                    shape=grad.get_shape().as_list(),
+                    trainable=False,
+                    initializer=tf.zeros_initializer(),
+                    dtype=grad.dtype,
+                    collections=[
+                        tf.compat.v1.GraphKeys.LOCAL_VARIABLES,
+                        "aggregating_collection"],
+                )
+                self.locally_aggregated_grads.append(grad_aggregation_variable)
+            assert len(self.locally_aggregated_grads) + \
+                self.num_none_grad_updates == len(grads)
+
+        # We expect to get a `sess` when we need to manually do a `sess.run(...)`
+        # for the variables to be initialized. This is the `tf.keras`
+        # optimizers.
+        if self.optimizer_type == self._OPTIMIZER_TYPE_KERAS:
+            session = tf.compat.v1.keras.backend.get_session(op_input_list=())
+            vars_init_op = tf.compat.v1.variables_initializer(
+                [self.counter, *get_not_none_from_list(self.locally_aggregated_grads)]
+            )
+            session.run(vars_init_op)
+
+    def _clear_grads(self):
+        clear_ops_list = []
+        for idx, grad_aggregator in enumerate(self.locally_aggregated_grads):
+            clear_op = grad_aggregator.assign(grad_aggregator.initial_value)
+            clear_ops_list.append(clear_op)
+        return tf.group(*clear_ops_list)
+
+    def _aggregate_grads(self, grads):
+        aggregation_ops_list = []
+        grads = get_not_none_from_list(grads)
+        assert len(grads) == len(self.locally_aggregated_grads)
+
+        # Apply new gradient updates to the locally copy.
+        for idx, grad in enumerate(grads):
+            if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):
+                grad = tf.convert_to_tensor(grad)
+
+            updated_grad_aggregator = self.locally_aggregated_grads[idx].assign_add(
+                grad)
+            aggregation_ops_list.append(updated_grad_aggregator)
+
+        return aggregation_ops_list
+
+    def _allreduce_grads_helper(self, grads):
+        # Read in latest variables values.
+        aggregated_grads = []
+        aggregation_read_ops_list = []
+        for idx, locally_aggregated_grad in enumerate(
+                self.locally_aggregated_grads):
+            aggregated_grads.append(locally_aggregated_grad.read_value())
+            aggregation_read_ops_list.append(aggregated_grads[idx])
+        aggregation_read_ops = tf.group(*aggregation_read_ops_list)
+
+        with tf.control_dependencies([aggregation_read_ops]):
+            averaged_gradients = self._allreduce_grads(aggregated_grads)
+
+            # Reset counter.
+            with tf.control_dependencies([g.op for g in averaged_gradients if g is not None]):
+                reset_op = self.counter.assign(
+                    tf.constant(0), use_locking=True)
+
+            # Divide by backward_passes_per_step if
+            # average_aggregated_gradients is True.
+            with tf.control_dependencies([reset_op]):
+                gradient_divisor = self.backward_passes_per_step if \
+                    self.average_aggregated_gradients else 1
+
+                averaged_gradients = apply_op_to_not_none_tensors(
+                    tf.divide,
+                    averaged_gradients,
+                    gradient_divisor,
+                )
+                return averaged_gradients
+
+    def compute_gradients(self, grads):
+        """
+        Applies the new gradient updates the locally aggregated gradients, and
+        performs cross-machine communication every backward_passes_per_step
+        times it is called.
+        """
+        self._init_aggregation_vars(grads)
+
+        # Clear the locally aggregated gradients when the counter is at zero.
+        clear_op = tf.cond(
+            pred=tf.equal(self.counter, 0),
+            true_fn=lambda: self._clear_grads(),
+            false_fn=tf.no_op
+        )
+
+        # Add new gradients to the locally aggregated gradients.
+        with tf.control_dependencies([clear_op]):
+            aggregation_ops_list = self._aggregate_grads(grads)
+
+        # Increment the counter once new gradients have been applied.
+        aggregation_ops = tf.group(*aggregation_ops_list)
+        with tf.control_dependencies([aggregation_ops]):
+            update_counter = self.counter.assign_add(tf.constant(1))
+
+        with tf.control_dependencies([update_counter]):
+            grads = get_not_none_from_list(grads)
+            assert len(grads) == len(self.locally_aggregated_grads)
+
+            # 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),
+                lambda: grads,
+            )
+
+            # Handle case where there is only one variable.
+            if not isinstance(allreduced_grads, (list, tuple)):
+                allreduced_grads = (allreduced_grads,)
+            assert len(allreduced_grads) == len(self.locally_aggregated_grads)
+
+            # Insert gradients that are None back in.
+            allreduced_grads = [
+                allreduced_grads[self.not_none_indexes[idx]] if idx in self.not_none_indexes else None
+                for idx in range(len(self.locally_aggregated_grads) + self.num_none_grad_updates)
+            ]
+            assert len(allreduced_grads) == len(
+                self.locally_aggregated_grads) + self.num_none_grad_updates
+
+        # If gradients have not been allreduced this batch, we return the gradients
+        # that were submitted as the updates (the input).
+        return allreduced_grads
+
+    def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):
+        """
+        Apply updates every backward_passes_per_step, which lines up with
+        the batches on which we communicated the locally aggregated gradients.
+        """
+        flattended_args0 = [item for tup in args[0] for item in tup]
+
+        # Since we skip applying updates when the counter is not at zero we
+        # 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()
+            if global_step_counter is None:
+                return tf.no_op()
+            return global_step_counter.assign_add(
+                tf.constant(1, dtype=tf.int64),
+                use_locking=True,
+                read_value=False
+            )
+
+        # Increment global step on iterations where we don't call `apply_gradients()`.
+        cond_increment_global_step_counter = tf.cond(
+            pred=tf.equal(self.counter, 0),
+            true_fn=tf.no_op,
+            false_fn=increment_global_step_counter,
+        )
+        flattended_args0.append(cond_increment_global_step_counter)
+
+        # If optimizer tracks iterations, we increment it on steps where we
+        # are not going to call `apply_gradients()`.
+        def increment_optimizer_iteration():
+            if hasattr(optimizer, "_iterations") and optimizer._iterations is not None:
+                return optimizer._iterations.assign_add(1).op
+            return tf.no_op()
+
+        with tf.control_dependencies([tf.group(*get_not_none_from_list(flattended_args0))]):
+            return tf.cond(
+                pred=tf.equal(self.counter, 0),
+                true_fn=apply_grads_closure,
+                false_fn=increment_optimizer_iteration,
+            )
diff --git a/horovod/tensorflow/__init__.py b/horovod/tensorflow/__init__.py
index c3f711ef31..9d94747e44 100644
--- a/horovod/tensorflow/__init__.py
+++ b/horovod/tensorflow/__init__.py
@@ -20,6 +20,7 @@
 import warnings
 
 from horovod.common.util import check_extension, gpu_available
+from horovod.common.gradient_aggregation import LocalGradientAggregationHelper
 
 check_extension('horovod.tensorflow', 'HOROVOD_WITH_TENSORFLOW', __file__, 'mpi_lib')
 
@@ -434,8 +435,9 @@ def variables(self, *args, **kwargs):
 
 def DistributedOptimizer(optimizer, name=None, use_locking=False, device_dense='',
                          device_sparse='', compression=Compression.none,
-                         sparse_as_dense=False, backward_passes_per_step=1,
-                         op=Average, gradient_predivide_factor=1.0):
+                         backward_passes_per_step=1, sparse_as_dense=False,
+                         op=Average, gradient_predivide_factor=1.0,
+                         average_aggregated_gradients=True):
     """Construct a new DistributedOptimizer, which uses another optimizer
     under the hood for computing single-process gradient values and
     applying gradient updates after the gradient values have been combined
@@ -477,6 +479,10 @@ def DistributedOptimizer(optimizer, name=None, use_locking=False, device_dense='
         before and after the sum. Gradients are scaled by
         1.0 / gradient_predivide_factor before the sum and
         gradient_predivide_factor / size after the sum.
+      average_aggregated_gradients:
+        Whether to average the aggregated gradients that have been accumulated
+        over multiple mini-batches. If true divide gradients by
+        backward_passes_per_step. Only applicable for backward_passes_per_step > 1.
     """
     if gradient_predivide_factor != 1.0:
         if rocm_built():
@@ -490,19 +496,35 @@ def DistributedOptimizer(optimizer, name=None, use_locking=False, device_dense='
                                             device_sparse, compression, backward_passes_per_step)
         else:
             if backward_passes_per_step > 1:
-                raise ValueError('backward_passes_per_step>1 is not supported yet with '
-                                 'op != Adasum')
-            return _DistributedOptimizer(optimizer, name, use_locking, device_dense,
-                                        device_sparse, compression, sparse_as_dense, op,
-                                        gradient_predivide_factor)
+                raise ValueError('backward_passes_per_step > 1 is not supported yet with '
+                                 'op != Adasum for Tensorflow optimizers.')
+            return _DistributedOptimizer(
+                optimizer=optimizer,
+                name=name,
+                use_locking=use_locking,
+                device_dense=device_dense,
+                device_sparse=device_sparse,
+                compression=compression,
+                sparse_as_dense=sparse_as_dense,
+                op=op,
+                gradient_predivide_factor=gradient_predivide_factor,
+            )
     elif isinstance(optimizer, tf.keras.optimizers.Optimizer):
         if op == Adasum:
             raise ValueError('op == Adasum is not supported yet with Keras')
-        if backward_passes_per_step > 1:
-            raise ValueError('backward_passes_per_step > 1 is not supported yet with Keras')
+
         import horovod.tensorflow.keras as hvd_k
-        return hvd_k.DistributedOptimizer(optimizer, name, device_dense, device_sparse,
-                                          compression, sparse_as_dense, gradient_predivide_factor)
+        return hvd_k.DistributedOptimizer(
+            optimizer=optimizer,
+            name=name,
+            device_dense=device_dense,
+            device_sparse=device_sparse,
+            compression=compression,
+            sparse_as_dense=sparse_as_dense,
+            gradient_predivide_factor=gradient_predivide_factor,
+            backward_passes_per_step=backward_passes_per_step,
+            average_aggregated_gradients=average_aggregated_gradients,
+        )
     else:
         raise ValueError('Provided optimizer doesn\'t inherit from either legacy '
                          'TensorFlow or Keras optimizer: %s' % optimizer)
diff --git a/horovod/tensorflow/keras/__init__.py b/horovod/tensorflow/keras/__init__.py
index f58faddcc1..b863722dd4 100644
--- a/horovod/tensorflow/keras/__init__.py
+++ b/horovod/tensorflow/keras/__init__.py
@@ -49,7 +49,9 @@ def DistributedOptimizer(optimizer, name=None,
                          device_dense='', device_sparse='',
                          compression=Compression.none,
                          sparse_as_dense=False,
-                         gradient_predivide_factor=1.0):
+                         gradient_predivide_factor=1.0,
+                         backward_passes_per_step=1,
+                         average_aggregated_gradients=True):
     """
     An optimizer that wraps another keras.optimizers.Optimizer, using an allreduce to
     average gradient values before applying gradients to model weights.
@@ -74,13 +76,30 @@ def DistributedOptimizer(optimizer, name=None,
                                    before and after the sum. Gradients are scaled by
                                    1.0 / gradient_predivide_factor before the sum and
                                    gradient_predivide_factor / size after the sum.
+        backward_passes_per_step: Number of backward passes to perform before calling
+                                  hvd.allreduce. This allows accumulating updates over
+                                  multiple mini-batches before reducing and applying them.
+        average_aggregated_gradients: Whether to average the aggregated gradients that
+                                      have been accumulated over multiple mini-batches.
+                                      If true divides gradient updates by
+                                      backward_passes_per_step
+                                      Only applicable for backward_passes_per_step > 1.
     """
     if gradient_predivide_factor != 1.0 and rocm_built():
             raise ValueError('gradient_predivide_factor not supported yet with ROCm')
 
-    return _impl.create_distributed_optimizer(keras, optimizer, name,
-                                              device_dense, device_sparse, compression,
-                                              sparse_as_dense, gradient_predivide_factor)
+    return _impl.create_distributed_optimizer(
+        keras=keras,
+        optimizer=optimizer,
+        name=name,
+        device_dense=device_dense,
+        device_sparse=device_sparse,
+        compression=compression,
+        sparse_as_dense=sparse_as_dense,
+        gradient_predivide_factor=gradient_predivide_factor,
+        backward_passes_per_step=backward_passes_per_step,
+        average_aggregated_gradients=average_aggregated_gradients
+    )
 
 
 def broadcast_global_variables(root_rank):
diff --git a/test/parallel/test_tensorflow_keras.py b/test/parallel/test_tensorflow_keras.py
index e337bdbcb1..c3a3387863 100644
--- a/test/parallel/test_tensorflow_keras.py
+++ b/test/parallel/test_tensorflow_keras.py
@@ -15,6 +15,7 @@
 
 """Tests for horovod.tensorflow.keras."""
 
+import math
 import numpy as np
 import pytest
 import tensorflow as tf
@@ -23,6 +24,7 @@
 from distutils.version import LooseVersion
 from tensorflow import keras
 from tensorflow.python.keras import backend as K
+from tensorflow.python.keras.optimizer_v2 import optimizer_v2
 
 import horovod.tensorflow.keras as hvd
 
@@ -40,16 +42,17 @@ def __init__(self, *args, **kwargs):
         warnings.simplefilter('module')
         hvd.init()
 
-        self.config = tf.ConfigProto()
+        self.config = tf.compat.v1.ConfigProto()
         self.config.gpu_options.allow_growth = True
         self.config.gpu_options.visible_device_list = str(hvd.local_rank())
 
-    def test_train_model(self):
+    def train_model(self, backward_passes_per_step):
         with self.test_session(config=self.config) as sess:
             K.set_session(sess)
 
             opt = keras.optimizers.RMSprop(lr=0.0001)
-            opt = hvd.DistributedOptimizer(opt)
+            opt = hvd.DistributedOptimizer(
+                opt, backward_passes_per_step=backward_passes_per_step)
 
             model = keras.models.Sequential()
             model.add(keras.layers.Dense(2, input_shape=(3,)))
@@ -77,6 +80,12 @@ def generator():
                                 workers=4,
                                 initial_epoch=1)
 
+    def test_train_model(self):
+        self.train_model(backward_passes_per_step=1)
+
+    def test_train_model_with_gradient_aggregation(self):
+        self.train_model(backward_passes_per_step=2)
+
     def test_sparse_as_dense(self):
         with self.test_session(config=self.config) as sess:
             K.set_session(sess)
@@ -342,3 +351,61 @@ def test_elastic_state(self):
                 self.assertAllClose(w1, w2)
             assert state.batch == 21
             assert state.epoch == 11
+
+    def test_gradient_aggregation(self):
+        with self.test_session(config=self.config) as sess:
+
+            class TestingOptimizer(optimizer_v2.OptimizerV2):
+                """
+                Custom optimizer we use for testing gradient aggregation.
+                """
+                def get_config(self):
+                    config = super(TestingOptimizer, self).get_config()
+                    return config
+
+                def _create_slots(self, var_list):
+                    pass
+
+                def _resource_apply_dense(self, grad, var, apply_state=None):
+                    return var.assign_add(grad)
+
+            K.set_session(sess)
+            session = tf.compat.v1.keras.backend.get_session(op_input_list=())
+
+            backward_passes_per_step = 4
+            hvd_optimizer = hvd.DistributedOptimizer(
+                optimizer=TestingOptimizer("test"),
+                backward_passes_per_step=backward_passes_per_step,
+            )
+            iterations = hvd_optimizer.iterations
+            session.run(iterations.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
+
+            grads = [tf.constant([float(hvd.rank())])]
+            variables = [tf.Variable([0.0])]
+            session.run(variables[0].initializer)
+
+            allreduce_op = hvd_optimizer._allreduce(grads)
+            grads_and_vars = [(allreduce_op[0], variables[0])]
+            apply_grads_op = hvd_optimizer.apply_gradients(grads_and_vars)
+
+            for idx in range(10):
+                _ = session.run(apply_grads_op)
+
+                assert idx + 1 == session.run(hvd_optimizer.iterations)
+                assert session.run(variables[0].read_value()) == compute_expected_value(idx)


Code review comments:
diff_hunk: @@ -434,8 +435,9 @@ def variables(self, *args, **kwargs):
 
 def DistributedOptimizer(optimizer, name=None, use_locking=False, device_dense='',
                          device_sparse='', compression=Compression.none,
-                         sparse_as_dense=False, backward_passes_per_step=1,
-                         op=Average, gradient_predivide_factor=1.0):
+                         backward_passes_per_step=1, sparse_as_dense=False,
path: horovod/tensorflow/__init__.py
text: @user1:
For backwards compatibility (in case there's someone out there not using keyword args), we should keep `backward_passes_per_step` coming just after `sparse_as_dense`.

@author:
Thanks for catching this, didn't mean to change the order.

diff_hunk: @@ -0,0 +1,266 @@
+import tensorflow as tf
+
+
+def apply_op_to_not_none_tensors(tensor_op, tensors, *args):
+    return [
+        tensor_op(
+            tensor,
+            *args
+        ) if tensor is not None else tensor for tensor in tensors]
+
+
+def get_not_none_from_list(tensor_list):
+    return [x for x in tensor_list if x is not None]
+
+
+class LocalGradientAggregationHelper:
+    """
+    LocalGradientAggregationHelper aggregates gradient updates locally,
+    and communicates the updates across machines only once every
+    backward_passes_per_step. Only supports graph mode execution.
+    """
+
+    _OPTIMIZER_TYPE_KERAS = "optimizer_type_keras"
+
+    def __init__(
+            self,
+            backward_passes_per_step,
+            allreduce_func,
+            sparse_as_dense,
+            average_aggregated_gradients,
+            rank,
+            optimizer_type):
+        self._allreduce_grads = allreduce_func
+
+        # backward_passes_per_step controls how often gradient updates are
+        # synchronized.
+        self.backward_passes_per_step = backward_passes_per_step
+        assert self.backward_passes_per_step > 0
+
+        # average_aggregated_gradients controls whether gradient updates that are
+        # aggregated, should be divided by `backward_passes_per_step`.
+        self.average_aggregated_gradients = average_aggregated_gradients
+
+        # This is going to be [N] data structure holding the aggregated gradient updates
+        # N is the number of parameters.
+        self.locally_aggregated_grads = []
+
+        # Used to know when to allreduce and apply gradients. We allreduce when `self.counter`
+        # is equal to `self.backward_passes_per_step`. We apply gradients when `self.counter` is
+        # equal to 0.
+        self.counter = None
+
+        self.sparse_as_dense = sparse_as_dense
+        self.rank = rank
+        self.optimizer_type = optimizer_type
+
+        # Contains the mapping of indexes of grad updates that are not None to their index in
+        # locally_aggregated_grads which only contains not None gradients. When performing
+        # gradient aggregation we have to remove them from the list of grads prior to passing
+        # the list into a tf.cond().
+        self.not_none_indexes = {}
+        self.num_none_grad_updates = 0
+
+    def _init_aggregation_vars(self, grads):
+        """
+        Initializes the counter that is used when to communicate and aggregate gradients
+        and the tensorflow variables that store the locally aggregated gradients.
+        """
+        variable_scope_name = "aggregation_variables_" + str(self.rank)
+        with tf.compat.v1.variable_scope(variable_scope_name, reuse=tf.compat.v1.AUTO_REUSE):
+            self.counter = tf.compat.v1.get_variable(
+                "aggregation_counter", shape=(), dtype=tf.int32,
+                trainable=False, initializer=tf.compat.v1.zeros_initializer(),
+            )
+            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 AssertionError(
+                        "IndexedSlices are not supported when "
+                        "`backward_passes_per_step` > 1 and "
+                        "`sparse_as_dense` is False."
+                    )
+
+                # Handle grads that are None.
+                if grad is None:
+                    self.num_none_grad_updates += 1
+                    continue
+                self.not_none_indexes[idx] = len(self.locally_aggregated_grads)
+
+                # Create shadow variable.
+                grad_aggregation_variable_name = str(idx)
+                grad_aggregation_variable = tf.compat.v1.get_variable(
+                    grad_aggregation_variable_name,
+                    shape=grad.get_shape().as_list(),
+                    trainable=False,
+                    initializer=tf.zeros_initializer(),
+                    dtype=grad.dtype,
+                    collections=[
+                        tf.compat.v1.GraphKeys.LOCAL_VARIABLES,
+                        "aggregating_collection"],
+                )
+                self.locally_aggregated_grads.append(grad_aggregation_variable)
+            assert len(self.locally_aggregated_grads) + \
+                self.num_none_grad_updates == len(grads)
+
+        # We expect to get a `sess` when we need to manually do a `sess.run(...)`
+        # for the variables to be initialized. This is the `tf.keras`
+        # optimizers.
+        if self.optimizer_type == self._OPTIMIZER_TYPE_KERAS:
+            session = tf.compat.v1.keras.backend.get_session(op_input_list=())
+            vars_init_op = tf.compat.v1.variables_initializer(
+                [self.counter, *get_not_none_from_list(self.locally_aggregated_grads)]
+            )
+            session.run(vars_init_op)
+
+    def _clear_grads(self):
+        clear_ops_list = []
+        for idx, grad_aggregator in enumerate(self.locally_aggregated_grads):
+            clear_op = grad_aggregator.assign(grad_aggregator.initial_value)
+            clear_ops_list.append(clear_op)
+        return tf.group(*clear_ops_list)
+
+    def _aggregate_grads(self, grads):
+        aggregation_ops_list = []
+        grads = get_not_none_from_list(grads)
+        assert len(grads) == len(self.locally_aggregated_grads)
+
+        # Apply new gradient updates to the locally copy.
path: horovod/common/gradient_aggregation.py
text: @user1:
Typo: `locally` -> `local`

@author:
Thanks for catching this.

diff_hunk: @@ -0,0 +1,266 @@
+import tensorflow as tf
+
+
+def apply_op_to_not_none_tensors(tensor_op, tensors, *args):
+    return [
+        tensor_op(
+            tensor,
+            *args
+        ) if tensor is not None else tensor for tensor in tensors]
+
+
+def get_not_none_from_list(tensor_list):
+    return [x for x in tensor_list if x is not None]
+
+
+class LocalGradientAggregationHelper:
+    """
+    LocalGradientAggregationHelper aggregates gradient updates locally,
+    and communicates the updates across machines only once every
+    backward_passes_per_step. Only supports graph mode execution.
+    """
+
+    _OPTIMIZER_TYPE_KERAS = "optimizer_type_keras"
+
+    def __init__(
+            self,
+            backward_passes_per_step,
+            allreduce_func,
+            sparse_as_dense,
+            average_aggregated_gradients,
+            rank,
+            optimizer_type):
+        self._allreduce_grads = allreduce_func
+
+        # backward_passes_per_step controls how often gradient updates are
+        # synchronized.
+        self.backward_passes_per_step = backward_passes_per_step
+        assert self.backward_passes_per_step > 0
+
+        # average_aggregated_gradients controls whether gradient updates that are
+        # aggregated, should be divided by `backward_passes_per_step`.
+        self.average_aggregated_gradients = average_aggregated_gradients
+
+        # This is going to be [N] data structure holding the aggregated gradient updates
+        # N is the number of parameters.
+        self.locally_aggregated_grads = []
+
+        # Used to know when to allreduce and apply gradients. We allreduce when `self.counter`
+        # is equal to `self.backward_passes_per_step`. We apply gradients when `self.counter` is
+        # equal to 0.
+        self.counter = None
+
+        self.sparse_as_dense = sparse_as_dense
+        self.rank = rank
+        self.optimizer_type = optimizer_type
+
+        # Contains the mapping of indexes of grad updates that are not None to their index in
+        # locally_aggregated_grads which only contains not None gradients. When performing
+        # gradient aggregation we have to remove them from the list of grads prior to passing
+        # the list into a tf.cond().
+        self.not_none_indexes = {}
+        self.num_none_grad_updates = 0
+
+    def _init_aggregation_vars(self, grads):
+        """
+        Initializes the counter that is used when to communicate and aggregate gradients
+        and the tensorflow variables that store the locally aggregated gradients.
+        """
+        variable_scope_name = "aggregation_variables_" + str(self.rank)
+        with tf.compat.v1.variable_scope(variable_scope_name, reuse=tf.compat.v1.AUTO_REUSE):
+            self.counter = tf.compat.v1.get_variable(
+                "aggregation_counter", shape=(), dtype=tf.int32,
+                trainable=False, initializer=tf.compat.v1.zeros_initializer(),
+            )
+            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 AssertionError(
path: horovod/common/gradient_aggregation.py
text: @user1:
Let's change this to `ValueError`.

@author:
Done

diff_hunk: @@ -0,0 +1,266 @@
+import tensorflow as tf
+
+
+def apply_op_to_not_none_tensors(tensor_op, tensors, *args):
+    return [
+        tensor_op(
+            tensor,
+            *args
+        ) if tensor is not None else tensor for tensor in tensors]
+
+
+def get_not_none_from_list(tensor_list):
+    return [x for x in tensor_list if x is not None]
+
+
+class LocalGradientAggregationHelper:
+    """
+    LocalGradientAggregationHelper aggregates gradient updates locally,
+    and communicates the updates across machines only once every
+    backward_passes_per_step. Only supports graph mode execution.
+    """
+
+    _OPTIMIZER_TYPE_KERAS = "optimizer_type_keras"
+
+    def __init__(
+            self,
+            backward_passes_per_step,
+            allreduce_func,
+            sparse_as_dense,
+            average_aggregated_gradients,
+            rank,
+            optimizer_type):
+        self._allreduce_grads = allreduce_func
+
+        # backward_passes_per_step controls how often gradient updates are
+        # synchronized.
+        self.backward_passes_per_step = backward_passes_per_step
+        assert self.backward_passes_per_step > 0
path: horovod/common/gradient_aggregation.py
text: @user1:
Instead of an assert, it would be safer here to raise a `ValueError` so the user can handle it.

@author:
Done.