We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
hvd.DistributedOptimizer gradient accumulation doesn't clean up infinite gradient correctly
**Environment:**
1. Framework: (TensorFlow, Keras, PyTorch, MXNet) Keras
2. Framework version: 2.4
3. Horovod version: 2.3
4. MPI version: 
5. CUDA version: 
6. NCCL version:
7. Python version:
8. Spark / PySpark version:
9. Ray version:
10. OS and version:
11. GCC version:
12. CMake version:

**Checklist:**
1. Did you search issues to find if somebody asked this question before?
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)?
4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)?

**Bug report:**

We were training in TensorFlow [FP16 mixed precision](https://www.tensorflow.org/guide/mixed_precision) with keras `model.fit()` and with gradient accumulation/aggregation (`backward_pass_per_step` in `hvd.DistributedOptimizer`) and noticed that the [GradientAggregationHelperEager](https://github.com/horovod/horovod/blob/master/horovod/tensorflow/gradient_aggregation_eager.py#L8) doesn't work correctly with FP16 when the loss goes infinite. Details:

It is kind of expected that at the very first 2-15 steps of the training, the gradient out of TF [LossScaleOptimizer](https://github.com/keras-team/keras/blob/v2.8.0/keras/mixed_precision/loss_scale_optimizer.py#L258-L844) is infinite (because the default initial loss scale factor is as large as `2**15`). Dynamic LSO can handle this gracefully, it just skips applying gradient of that step and divides the scale factor by half. However horovod GradientAggregationHelper will anyway add the infinite gradient up locally, and the infinite gradient will never be correctly cleaned up in [this way](https://github.com/horovod/horovod/blob/133ef0725253db83cfb82a4ed4003df76d189829/horovod/tensorflow/gradient_aggregation_eager.py#L119-L123):
```
    def _clear_vars(self):
        self.counter.assign(0)
        for idx in self.locally_aggregated_grads.keys():
            self.locally_aggregated_grads[idx].assign_add(
                -1 * self.locally_aggregated_grads[idx])
```

as the result of adding inf value by its negative val will be NaN. 


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `horovod/tensorflow/gradient_aggregation_eager.py`
Content:
```
1 from distutils.version import LooseVersion
2 
3 import tensorflow as tf
4 
5 _POST_TF_2_4_0 = LooseVersion(tf.__version__) >= LooseVersion('2.4.0')
6 
7 
8 class LocalGradientAggregationHelperEager:
9     """
10    LocalGradientAggregationHelperEager aggregates gradient updates
11    locally, and communicates the updates across machines only once
12    every backward_passes_per_step. Only supports eager execution.
13    """
14 
15     def __init__(
16         self,
17         backward_passes_per_step,
18         allreduce_func,
19         sparse_as_dense,
20         average_aggregated_gradients,
21     ):
22         self.allreduce_grads = allreduce_func
23         self.sparse_as_dense = sparse_as_dense
24 
25         # backward_passes_per_step controls how often gradient updates are
26         # synchronized.
27         self.backward_passes_per_step = backward_passes_per_step
28         if self.backward_passes_per_step <= 0:
29             raise ValueError("backward_passes_per_step must be > 0")
30 
31         # average_aggregated_gradients controls whether gradient updates that are
32         # aggregated, should be divided by `backward_passes_per_step`.
33         self.average_aggregated_gradients = average_aggregated_gradients
34 
35         # This is going to be N data structure holding the aggregated gradient updates
36         # for parameter updates. N is the number of parameters.
37         self.locally_aggregated_grads = {}
38 
39         # Used to know when to allreduce and apply gradients. We allreduce when `self.counter`
40         # is equal to `self.backward_passes_per_step`. We apply gradients when `self.counter`
41         # is equal to 0.
42         self.counter = tf.Variable(initial_value=0)
43 
44     def compute_gradients(self, grads, vars):
45         # On steps where allreduce happens, resulting_grads returns the allreduced
46         # gradients, on other steps it returns the locally aggregated
47         # gradients.
48         resulting_grads = []
49 
50         for idx, grad in enumerate(grads):
51             # Handle IndexedSlices.
52             if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):
53                 grad = tf.convert_to_tensor(grad)
54             elif isinstance(grad, tf.IndexedSlices):
55                 raise ValueError(
56                     "IndexedSlices are not supported when "
57                     "`backward_passes_per_step` > 1 and "
58                     "`sparse_as_dense` is False."
59                 )
60 
61             # Create variables to store to aggregate gradients if they don't
62             # already exist. Skip variables that are None.
63             if idx not in self.locally_aggregated_grads.keys():
64                 if grad is not None:
65                     self.locally_aggregated_grads[idx] = tf.Variable(
66                         initial_value=tf.zeros_like(grad),
67                         trainable=False,
68                         dtype=grad.dtype,
69                     )
70 
71             if grad is None:
72                 resulting_grads.append(None)
73             else:
74                 self.locally_aggregated_grads[idx].assign_add(grad)
75                 resulting_grads.append(
76                     self.locally_aggregated_grads[idx].read_value())
77         assert len(self.locally_aggregated_grads) == len(grads)
78 
79         # Increment counter.
80         self.counter.assign_add(1)
81 
82         def _all_reduce_and_clear_aggregated_variables(aggregated_gradients, vars):
83             # Performs allreduce. If `average_aggregated_gradients` is
84             # set to True divides result by `backward_passes_per_step`.
85             reduced_gradients = self._allreduce_helper(aggregated_gradients, vars)
86             assert len(reduced_gradients) == len(grads)
87 
88             self._clear_vars()
89             return reduced_gradients
90 
91         def _do_nothing(aggregated_gradients):
92             return aggregated_gradients
93 
94         resulting_grads = tf.cond(
95             pred=tf.equal(self.counter, self.backward_passes_per_step),
96             true_fn=lambda: _all_reduce_and_clear_aggregated_variables(resulting_grads, vars),
97             false_fn=lambda: _do_nothing(resulting_grads),
98         )
99 
100         return resulting_grads
101 
102     def _allreduce_helper(self, grads, vars):
103         allreduced_grads = self.allreduce_grads(grads, vars)
104 
105         if not self.average_aggregated_gradients:
106             return allreduced_grads
107 
108         scaled_allreduced_grads = []
109         for grad in allreduced_grads:
110             if grad is None:
111                 scaled_allreduced_grads.append(grad)
112                 continue
113 
114             scaled_allreduced_grads.append(
115                 grad / self.backward_passes_per_step)
116 
117         return scaled_allreduced_grads
118 
119     def _clear_vars(self):
120         self.counter.assign(0)
121         for idx in self.locally_aggregated_grads.keys():
122             self.locally_aggregated_grads[idx].assign_add(
123                 -1 * self.locally_aggregated_grads[idx])
124 
125     def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):
126         def increment_optimizer_iteration():
127             if hasattr(optimizer, "_iterations") and optimizer._iterations is not None:
128                 return optimizer._iterations.assign_add(1).op
129             return tf.no_op()
130 
131         def non_aggregation_step():
132             if _POST_TF_2_4_0:
133                 # In TF 2.4+ `_aggregate_gradients()` is called from inside of `apply_gradients()`.
134                 # We account for this by calling `_aggregate_gradients()` for steps where we do
135                 # not call `apply_gradients()`.
136                 transformed_grads_and_vars = optimizer._transform_unaggregated_gradients(
137                     args[0])
138                 _ = optimizer._aggregate_gradients(transformed_grads_and_vars)
139 
140             return increment_optimizer_iteration()
141 
142         def is_aggregation_step():
143             return tf.equal(self.counter, 0)
144 
145         return tf.cond(
146             pred=is_aggregation_step(),
147             true_fn=apply_grads_closure,
148             false_fn=non_aggregation_step,
149         )
150
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/horovod/tensorflow/gradient_aggregation_eager.py b/horovod/tensorflow/gradient_aggregation_eager.py
--- a/horovod/tensorflow/gradient_aggregation_eager.py
+++ b/horovod/tensorflow/gradient_aggregation_eager.py
@@ -119,8 +119,8 @@
     def _clear_vars(self):
         self.counter.assign(0)
         for idx in self.locally_aggregated_grads.keys():
-            self.locally_aggregated_grads[idx].assign_add(
-                -1 * self.locally_aggregated_grads[idx])
+            self.locally_aggregated_grads[idx].assign(
+                tf.zeros_like(self.locally_aggregated_grads[idx]))
 
     def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):
         def increment_optimizer_iteration():

{"golden_diff": "diff --git a/horovod/tensorflow/gradient_aggregation_eager.py b/horovod/tensorflow/gradient_aggregation_eager.py\n--- a/horovod/tensorflow/gradient_aggregation_eager.py\n+++ b/horovod/tensorflow/gradient_aggregation_eager.py\n@@ -119,8 +119,8 @@\n     def _clear_vars(self):\n         self.counter.assign(0)\n         for idx in self.locally_aggregated_grads.keys():\n-            self.locally_aggregated_grads[idx].assign_add(\n-                -1 * self.locally_aggregated_grads[idx])\n+            self.locally_aggregated_grads[idx].assign(\n+                tf.zeros_like(self.locally_aggregated_grads[idx]))\n \n     def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):\n         def increment_optimizer_iteration():\n", "issue": "hvd.DistributedOptimizer gradient accumulation doesn't clean up infinite gradient correctly\n**Environment:**\r\n1. Framework: (TensorFlow, Keras, PyTorch, MXNet) Keras\r\n2. Framework version: 2.4\r\n3. Horovod version: 2.3\r\n4. MPI version: \r\n5. CUDA version: \r\n6. NCCL version:\r\n7. Python version:\r\n8. Spark / PySpark version:\r\n9. Ray version:\r\n10. OS and version:\r\n11. GCC version:\r\n12. CMake version:\r\n\r\n**Checklist:**\r\n1. Did you search issues to find if somebody asked this question before?\r\n2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?\r\n3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)?\r\n4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)?\r\n\r\n**Bug report:**\r\n\r\nWe were training in TensorFlow [FP16 mixed precision](https://www.tensorflow.org/guide/mixed_precision) with keras `model.fit()` and with gradient accumulation/aggregation (`backward_pass_per_step` in `hvd.DistributedOptimizer`) and noticed that the [GradientAggregationHelperEager](https://github.com/horovod/horovod/blob/master/horovod/tensorflow/gradient_aggregation_eager.py#L8) doesn't work correctly with FP16 when the loss goes infinite. Details:\r\n\r\nIt is kind of expected that at the very first 2-15 steps of the training, the gradient out of TF [LossScaleOptimizer](https://github.com/keras-team/keras/blob/v2.8.0/keras/mixed_precision/loss_scale_optimizer.py#L258-L844) is infinite (because the default initial loss scale factor is as large as `2**15`). Dynamic LSO can handle this gracefully, it just skips applying gradient of that step and divides the scale factor by half. However horovod GradientAggregationHelper will anyway add the infinite gradient up locally, and the infinite gradient will never be correctly cleaned up in [this way](https://github.com/horovod/horovod/blob/133ef0725253db83cfb82a4ed4003df76d189829/horovod/tensorflow/gradient_aggregation_eager.py#L119-L123):\r\n```\r\n    def _clear_vars(self):\r\n        self.counter.assign(0)\r\n        for idx in self.locally_aggregated_grads.keys():\r\n            self.locally_aggregated_grads[idx].assign_add(\r\n                -1 * self.locally_aggregated_grads[idx])\r\n```\r\n\r\nas the result of adding inf value by its negative val will be NaN. \r\n\n", "before_files": [{"content": "from distutils.version import LooseVersion\n\nimport tensorflow as tf\n\n_POST_TF_2_4_0 = LooseVersion(tf.__version__) >= LooseVersion('2.4.0')\n\n\nclass LocalGradientAggregationHelperEager:\n    \"\"\"\n   LocalGradientAggregationHelperEager aggregates gradient updates\n   locally, and communicates the updates across machines only once\n   every backward_passes_per_step. Only supports eager execution.\n   \"\"\"\n\n    def __init__(\n        self,\n        backward_passes_per_step,\n        allreduce_func,\n        sparse_as_dense,\n        average_aggregated_gradients,\n    ):\n        self.allreduce_grads = allreduce_func\n        self.sparse_as_dense = sparse_as_dense\n\n        # backward_passes_per_step controls how often gradient updates are\n        # synchronized.\n        self.backward_passes_per_step = backward_passes_per_step\n        if self.backward_passes_per_step <= 0:\n            raise ValueError(\"backward_passes_per_step must be > 0\")\n\n        # average_aggregated_gradients controls whether gradient updates that are\n        # aggregated, should be divided by `backward_passes_per_step`.\n        self.average_aggregated_gradients = average_aggregated_gradients\n\n        # This is going to be N data structure holding the aggregated gradient updates\n        # for parameter updates. N is the number of parameters.\n        self.locally_aggregated_grads = {}\n\n        # Used to know when to allreduce and apply gradients. We allreduce when `self.counter`\n        # is equal to `self.backward_passes_per_step`. We apply gradients when `self.counter`\n        # is equal to 0.\n        self.counter = tf.Variable(initial_value=0)\n\n    def compute_gradients(self, grads, vars):\n        # On steps where allreduce happens, resulting_grads returns the allreduced\n        # gradients, on other steps it returns the locally aggregated\n        # gradients.\n        resulting_grads = []\n\n        for idx, grad in enumerate(grads):\n            # Handle IndexedSlices.\n            if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):\n                grad = tf.convert_to_tensor(grad)\n            elif isinstance(grad, tf.IndexedSlices):\n                raise ValueError(\n                    \"IndexedSlices are not supported when \"\n                    \"`backward_passes_per_step` > 1 and \"\n                    \"`sparse_as_dense` is False.\"\n                )\n\n            # Create variables to store to aggregate gradients if they don't\n            # already exist. Skip variables that are None.\n            if idx not in self.locally_aggregated_grads.keys():\n                if grad is not None:\n                    self.locally_aggregated_grads[idx] = tf.Variable(\n                        initial_value=tf.zeros_like(grad),\n                        trainable=False,\n                        dtype=grad.dtype,\n                    )\n\n            if grad is None:\n                resulting_grads.append(None)\n            else:\n                self.locally_aggregated_grads[idx].assign_add(grad)\n                resulting_grads.append(\n                    self.locally_aggregated_grads[idx].read_value())\n        assert len(self.locally_aggregated_grads) == len(grads)\n\n        # Increment counter.\n        self.counter.assign_add(1)\n\n        def _all_reduce_and_clear_aggregated_variables(aggregated_gradients, vars):\n            # Performs allreduce. If `average_aggregated_gradients` is\n            # set to True divides result by `backward_passes_per_step`.\n            reduced_gradients = self._allreduce_helper(aggregated_gradients, vars)\n            assert len(reduced_gradients) == len(grads)\n\n            self._clear_vars()\n            return reduced_gradients\n\n        def _do_nothing(aggregated_gradients):\n            return aggregated_gradients\n\n        resulting_grads = tf.cond(\n            pred=tf.equal(self.counter, self.backward_passes_per_step),\n            true_fn=lambda: _all_reduce_and_clear_aggregated_variables(resulting_grads, vars),\n            false_fn=lambda: _do_nothing(resulting_grads),\n        )\n\n        return resulting_grads\n\n    def _allreduce_helper(self, grads, vars):\n        allreduced_grads = self.allreduce_grads(grads, vars)\n\n        if not self.average_aggregated_gradients:\n            return allreduced_grads\n\n        scaled_allreduced_grads = []\n        for grad in allreduced_grads:\n            if grad is None:\n                scaled_allreduced_grads.append(grad)\n                continue\n\n            scaled_allreduced_grads.append(\n                grad / self.backward_passes_per_step)\n\n        return scaled_allreduced_grads\n\n    def _clear_vars(self):\n        self.counter.assign(0)\n        for idx in self.locally_aggregated_grads.keys():\n            self.locally_aggregated_grads[idx].assign_add(\n                -1 * self.locally_aggregated_grads[idx])\n\n    def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):\n        def increment_optimizer_iteration():\n            if hasattr(optimizer, \"_iterations\") and optimizer._iterations is not None:\n                return optimizer._iterations.assign_add(1).op\n            return tf.no_op()\n\n        def non_aggregation_step():\n            if _POST_TF_2_4_0:\n                # In TF 2.4+ `_aggregate_gradients()` is called from inside of `apply_gradients()`.\n                # We account for this by calling `_aggregate_gradients()` for steps where we do\n                # not call `apply_gradients()`.\n                transformed_grads_and_vars = optimizer._transform_unaggregated_gradients(\n                    args[0])\n                _ = optimizer._aggregate_gradients(transformed_grads_and_vars)\n\n            return increment_optimizer_iteration()\n\n        def is_aggregation_step():\n            return tf.equal(self.counter, 0)\n\n        return tf.cond(\n            pred=is_aggregation_step(),\n            true_fn=apply_grads_closure,\n            false_fn=non_aggregation_step,\n        )\n", "path": "horovod/tensorflow/gradient_aggregation_eager.py"}], "after_files": [{"content": "from distutils.version import LooseVersion\n\nimport tensorflow as tf\n\n_POST_TF_2_4_0 = LooseVersion(tf.__version__) >= LooseVersion('2.4.0')\n\n\nclass LocalGradientAggregationHelperEager:\n    \"\"\"\n   LocalGradientAggregationHelperEager aggregates gradient updates\n   locally, and communicates the updates across machines only once\n   every backward_passes_per_step. Only supports eager execution.\n   \"\"\"\n\n    def __init__(\n        self,\n        backward_passes_per_step,\n        allreduce_func,\n        sparse_as_dense,\n        average_aggregated_gradients,\n    ):\n        self.allreduce_grads = allreduce_func\n        self.sparse_as_dense = sparse_as_dense\n\n        # backward_passes_per_step controls how often gradient updates are\n        # synchronized.\n        self.backward_passes_per_step = backward_passes_per_step\n        if self.backward_passes_per_step <= 0:\n            raise ValueError(\"backward_passes_per_step must be > 0\")\n\n        # average_aggregated_gradients controls whether gradient updates that are\n        # aggregated, should be divided by `backward_passes_per_step`.\n        self.average_aggregated_gradients = average_aggregated_gradients\n\n        # This is going to be N data structure holding the aggregated gradient updates\n        # for parameter updates. N is the number of parameters.\n        self.locally_aggregated_grads = {}\n\n        # Used to know when to allreduce and apply gradients. We allreduce when `self.counter`\n        # is equal to `self.backward_passes_per_step`. We apply gradients when `self.counter`\n        # is equal to 0.\n        self.counter = tf.Variable(initial_value=0)\n\n    def compute_gradients(self, grads, vars):\n        # On steps where allreduce happens, resulting_grads returns the allreduced\n        # gradients, on other steps it returns the locally aggregated\n        # gradients.\n        resulting_grads = []\n\n        for idx, grad in enumerate(grads):\n            # Handle IndexedSlices.\n            if self.sparse_as_dense and isinstance(grad, tf.IndexedSlices):\n                grad = tf.convert_to_tensor(grad)\n            elif isinstance(grad, tf.IndexedSlices):\n                raise ValueError(\n                    \"IndexedSlices are not supported when \"\n                    \"`backward_passes_per_step` > 1 and \"\n                    \"`sparse_as_dense` is False.\"\n                )\n\n            # Create variables to store to aggregate gradients if they don't\n            # already exist. Skip variables that are None.\n            if idx not in self.locally_aggregated_grads.keys():\n                if grad is not None:\n                    self.locally_aggregated_grads[idx] = tf.Variable(\n                        initial_value=tf.zeros_like(grad),\n                        trainable=False,\n                        dtype=grad.dtype,\n                    )\n\n            if grad is None:\n                resulting_grads.append(None)\n            else:\n                self.locally_aggregated_grads[idx].assign_add(grad)\n                resulting_grads.append(\n                    self.locally_aggregated_grads[idx].read_value())\n        assert len(self.locally_aggregated_grads) == len(grads)\n\n        # Increment counter.\n        self.counter.assign_add(1)\n\n        def _all_reduce_and_clear_aggregated_variables(aggregated_gradients, vars):\n            # Performs allreduce. If `average_aggregated_gradients` is\n            # set to True divides result by `backward_passes_per_step`.\n            reduced_gradients = self._allreduce_helper(aggregated_gradients, vars)\n            assert len(reduced_gradients) == len(grads)\n\n            self._clear_vars()\n            return reduced_gradients\n\n        def _do_nothing(aggregated_gradients):\n            return aggregated_gradients\n\n        resulting_grads = tf.cond(\n            pred=tf.equal(self.counter, self.backward_passes_per_step),\n            true_fn=lambda: _all_reduce_and_clear_aggregated_variables(resulting_grads, vars),\n            false_fn=lambda: _do_nothing(resulting_grads),\n        )\n\n        return resulting_grads\n\n    def _allreduce_helper(self, grads, vars):\n        allreduced_grads = self.allreduce_grads(grads, vars)\n\n        if not self.average_aggregated_gradients:\n            return allreduced_grads\n\n        scaled_allreduced_grads = []\n        for grad in allreduced_grads:\n            if grad is None:\n                scaled_allreduced_grads.append(grad)\n                continue\n\n            scaled_allreduced_grads.append(\n                grad / self.backward_passes_per_step)\n\n        return scaled_allreduced_grads\n\n    def _clear_vars(self):\n        self.counter.assign(0)\n        for idx in self.locally_aggregated_grads.keys():\n            self.locally_aggregated_grads[idx].assign(\n                tf.zeros_like(self.locally_aggregated_grads[idx]))\n\n    def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):\n        def increment_optimizer_iteration():\n            if hasattr(optimizer, \"_iterations\") and optimizer._iterations is not None:\n                return optimizer._iterations.assign_add(1).op\n            return tf.no_op()\n\n        def non_aggregation_step():\n            if _POST_TF_2_4_0:\n                # In TF 2.4+ `_aggregate_gradients()` is called from inside of `apply_gradients()`.\n                # We account for this by calling `_aggregate_gradients()` for steps where we do\n                # not call `apply_gradients()`.\n                transformed_grads_and_vars = optimizer._transform_unaggregated_gradients(\n                    args[0])\n                _ = optimizer._aggregate_gradients(transformed_grads_and_vars)\n\n            return increment_optimizer_iteration()\n\n        def is_aggregation_step():\n            return tf.equal(self.counter, 0)\n\n        return tf.cond(\n            pred=is_aggregation_step(),\n            true_fn=apply_grads_closure,\n            false_fn=non_aggregation_step,\n        )\n", "path": "horovod/tensorflow/gradient_aggregation_eager.py"}]}