Repository: horovod/horovod

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. 



Discussion:

Resolving patch:
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 _allreduce_helper(self, grads, vars):
     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():


Test patch:
diff --git a/test/parallel/test_tensorflow2_keras.py b/test/parallel/test_tensorflow2_keras.py
--- a/test/parallel/test_tensorflow2_keras.py
+++ b/test/parallel/test_tensorflow2_keras.py
@@ -182,6 +182,33 @@ def test_sparse_as_dense_with_grad_aggregation(self):
         aggregation_counter = opt._agg_helper.counter.numpy()
         assert aggregation_counter == training_steps % backward_passes_per_step
 
+    def test_grad_aggregation_with_inf_grad(self):
+        backward_passes_per_step = 2
+        step_count = tf.Variable(0, trainable=False, dtype=tf.int32)
+        opt = tf.optimizers.SGD()
+        opt = hvd.DistributedOptimizer(
+            opt, 
+            backward_passes_per_step=backward_passes_per_step,
+            sparse_as_dense=True
+        )
+        x = tf.Variable(0.)
+        var = [x]
+
+        def loss():
+            step_count.assign_add(1)
+            return tf.cond(
+                pred=tf.greater(step_count, 1),
+                true_fn=lambda: x,
+                false_fn=lambda: x * float('inf')
+            )
+        for _ in range(2 * backward_passes_per_step):
+            # in the first aggregation cycle the gradient is infinite,
+            # and it should be cleaned up to zero after apply_gradients
+            # and doesn't affect the 2nd aggregation cycle
+            grads_and_vars = opt._compute_gradients(loss=loss, var_list=var)
+            opt.apply_gradients(grads_and_vars)
+        assert tf.math.is_finite(grads_and_vars[0][0])
+
     def test_from_config(self):
         opt = keras.optimizers.Adam()
         hopt = hvd.DistributedOptimizer(opt)