Repository: horovod/horovod

Issue:
Bug in global_step update when local gradient aggregation is used
**Environment:**
1. Framework: TensorFlow
2. Framework version: 1.15.3
3. Horovod version: 0.21.3
4. MPI version: 3.1.6
5. CUDA version: 11.0.194
6. NCCL version: 2.7.6
7. Python version: 3.6.9
8. Spark / PySpark version:
9. OS and version: 
10. GCC version:
11. CMake version:

**Bug report:**
There seems to be a bug in how the `global_step` variable is updated when using local gradient aggregation with `backward_passes_per_step>1`. 

Here is a simple example to reproduce the issue:
```
import tensorflow as tf
import horovod.tensorflow as hvd

tf.disable_eager_execution()
LR=0.01
# data/label is always just 1
data = tf.constant([[1]], dtype=tf.float32)
# one-variable linear model, no bias (y = w*x)
#
# gradient update is hand-calculatable if needed for testing:
#
# model:            yhat = w*x
# mse loss:         L = (ytrue - w*x)**2
# gradient:         dL/dw = -2*x*(ytrue - `w*x)`
# let ytrue = x:    dL/dw = -2*x*x*(1-w)
def predict(x):
    dense = tf.layers.Dense(1, use_bias=False, kernel_initializer=tf.zeros_initializer())
    return dense(x)
pred = predict(data)
loss = tf.losses.mean_squared_error(data, pred)
step = tf.contrib.framework.get_or_create_global_step()
optimizer = tf.train.GradientDescentOptimizer(LR)
hvd.init()
optimizer = hvd.DistributedOptimizer(optimizer, backward_passes_per_step=2, average_aggregated_gradients=True)

with tf.control_dependencies([tf.print("step:", step, "loss:", loss)]):
    train_op = optimizer.minimize(loss, global_step=step)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())
    # training loop
    for i in range(10):
        sess.run([train_op])
```

Which gives this output:
```
step: 0 loss: 1
step: 0 loss: 1
step: 0 loss: 1
step: 2 loss: 0.960400045
step: 2 loss: 0.960400045
step: 4 loss: 0.922368109
step: 4 loss: 0.922368109
step: 6 loss: 0.885842443
step: 6 loss: 0.885842443
step: 8 loss: 0.850763
step: 8 loss: 0.850763
```





Discussion:
Hi @mnabian , thanks for reporting.
I don't think this is a real bug, it is a really interesting thing related to Tensorflow's internal graph scheduling though.

I've injected the below diff in hovorod (master), and ran your test again. It did show that `global_step` is updated correctly:
- Diff:
```
chongxiaoc@chongxiaoc-C02DL7WPMD6R horovod % git diff .
diff --git a/horovod/tensorflow/gradient_aggregation.py b/horovod/tensorflow/gradient_aggregation.py
index c6857d5..1a2f2be 100644
--- a/horovod/tensorflow/gradient_aggregation.py
+++ b/horovod/tensorflow/gradient_aggregation.py
@@ -261,8 +261,9 @@ class LocalGradientAggregationHelper:
             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,
-            )
+            with tf.control_dependencies([tf.print("in hvd apply_grads counter:", self.counter, "global step:", tf.train.get_global_step())]):
+                return tf.cond(
+                    pred=tf.equal(self.counter, 0),
+                    true_fn=apply_grads_closure,
+                    false_fn=increment_optimizer_iteration,
+                )
```

- Test output:
```

step: 0 loss: 1
step: 0 loss: 1
in hvd apply_grads counter: 1 global step: 0
step: 0 loss: 1
in hvd apply_grads counter: 0 global step: 1
step: 2 loss: 0.960400045
in hvd apply_grads counter: 1 global step: 2
step: 2 loss: 0.960400045
in hvd apply_grads counter: 0 global step: 3
step: 4 loss: 0.922368109
in hvd apply_grads counter: 1 global step: 4
step: 4 loss: 0.922368109
in hvd apply_grads counter: 0 global step: 5
step: 6 loss: 0.885842443
in hvd apply_grads counter: 1 global step: 6
step: 6 loss: 0.885842443
in hvd apply_grads counter: 0 global step: 7
step: 8 loss: 0.850763
in hvd apply_grads counter: 1 global step: 8
step: 8 loss: 0.850763
in hvd apply_grads counter: 0 global step: 9
```

Both `hvd.optimzier`'s `counter` and TF's `global step` are updated correctly by following `backward_passes_per_step`in horovod internally.

But why `tf.print` in the test has different value of `global_step`?
My two cents (I'm still learning TF), might be wrong:
- In horovod, TF optimizer's `apply_grads` is only connected in graph when `self.counter == 0` (resetting self.count to 0 if self.counter == backward_passes_per_step) as below:
https://github.com/horovod/horovod/blob/e9dba2bc06153217b3dd066eb5063d1ec15e0c8f/horovod/tensorflow/gradient_aggregation.py#L265

Horovod doesn't call TF optimizer's `apply_grads()` if `backward_passes_per_step` is not met, so next `forward run` will still use old weights without applying gradients. 

- For `backward_passes_per_step` not met cases, the graph execution of `forward run` and `previous` hvd optimizer `minimize` doesn't have sequential order in graph, so `tf.print` in test doesn't have  `global_step` value up to date.

@aaron276h is the author of `backward_passes_per_step` in TF, he might give better explanation.

@tgaddair Please feel free to correct me if I'm wrong.

Thanks.  

Hi, I helped @mnabian develop the reproducing example code.

I spent some time looking into the default `Optimizer.minimize` method.  There, we have a dependency graph (`-->` = "depends on) like:

```
# when there is no global_step
optimizer.apply_gradients() = gradient update               

# when there is a global_step  
optimizer.apply_gradients() = assign_add --> gradient update  
```

You can see that in this section of code [here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/optimizer.py#L625):

```python
      if global_step is None:
        apply_updates = self._finish(update_ops, name)
      else:
        with ops.control_dependencies([self._finish(update_ops, "update")]):
          with ops.colocate_with(global_step):
            if isinstance(
                global_step, resource_variable_ops.BaseResourceVariable):
              # TODO(apassos): the implicit read in assign_add is slow; consider
              # making it less so.
              apply_updates = resource_variable_ops.assign_add_variable_op(
                  global_step.handle,
                  ops.convert_to_tensor(1, dtype=global_step.dtype),
                  name=name)
            else:
              apply_updates = state_ops.assign_add(global_step, 1, name=name)
      ...
      return apply_updates
```

But it seems that what we have with this implementation is a reversal of the dependency relationship between the `train_op` returned by `minimize()` and the `assign_add`:

```
# when counter == 0, obviously behaves correctly
optimizer.apply_gradients() = wrapped_optimizer.apply_gradients() 

# when counter != 0, reversed dependency; it's valid for assign_add to run way too early
optimizer.apply_gradients() =  increment_optimizer_iteration --> assign_add  
```

In short, normally the `train_op = optimizer.apply_gradients()` *is* the `assign_add` and depends on the the gradient update happening first, but with this implementation, the `train_op` is the gradient update and the `assign_add` is declared to happen first.

I don't understand why the `assign_add` is not honoring the declared dependency on `tf.print()` in the example as written, but there's a lot of things about tensorflow I don't understand.

But this looks to me like a bug where `assign_add` is not happening at a well-defined time.

I think that what's going on in the provided counter-example is that for the `counter!=0` cases, the step is being incremented too early, while in the `counter==0` cases the step is incremented at the right time; after the user's print and before the injected print.  Or in other words, the injected print runs always after the assign_add, but the user's print sometimes runs before and sometimes after.

Would it be wrong to return the `cond_increment_global_step_counter` with an explicit depenency on the train op?

```diff
     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(
+            train_op = tf.cond(
                 pred=tf.equal(self.counter, 0),
                 true_fn=apply_grads_closure,
                 false_fn=increment_optimizer_iteration,
             )

+        # 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
+            )

+        with tf.control_dependencies([train_op]):
+            # Increment global step on iterations where we don't call `apply_gradients()`.
+            return = tf.cond(
+                pred=tf.equal(self.counter, 0),
+                true_fn=tf.no_op,
+                false_fn=increment_global_step_counter,
+            )
```

I confess I've had trouble building horovod from source to actually test this solution or make a PR...
@rb-determined-ai 
Thanks for your contribution, your explanation makes sense to me. Applying the above diff works on my local test:
```
step: 0 loss: 1
step: 0 loss: 1
after train step: 1 loss: 1
step: 1 loss: 1
after train step: 2 loss: 1
step: 2 loss: 0.960400045
after train step: 3 loss: 0.960400045
step: 3 loss: 0.960400045
after train step: 4 loss: 0.960400045
step: 4 loss: 0.922368109
after train step: 5 loss: 0.922368109
step: 5 loss: 0.922368109
after train step: 6 loss: 0.922368109
step: 6 loss: 0.885842443
after train step: 7 loss: 0.885842443
step: 7 loss: 0.885842443
after train step: 8 loss: 0.885842443
step: 8 loss: 0.850763
after train step: 9 loss: 0.850763
step: 9 loss: 0.850763
after train step: 10 loss: 0.850763
```

Please create PR to upstream above. Thanks for contributing to horovod.

Resolving patch:
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
@@ -232,6 +232,20 @@ def apply_gradients(self, apply_grads_closure, optimizer, *args, **kwargs):
         """
         flattended_args0 = [item for tup in args[0] for item in tup]
 
+        # 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))]):
+            train_op = tf.cond(
+                pred=tf.equal(self.counter, 0),
+                true_fn=apply_grads_closure,
+                false_fn=increment_optimizer_iteration,
+            )
+
         # 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).
@@ -245,24 +259,10 @@ def increment_global_step_counter():
                 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))]):
+        with tf.control_dependencies([train_op]):
+            # Increment global step on iterations where we don't call `apply_gradients()`.
             return tf.cond(
                 pred=tf.equal(self.counter, 0),
-                true_fn=apply_grads_closure,
-                false_fn=increment_optimizer_iteration,
+                true_fn=tf.no_op,
+                false_fn=increment_global_step_counter,
             )


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
@@ -51,7 +51,7 @@
     # tests from running in the graph mode.
     tf.enable_eager_execution(config=config)
 
-ccl_supported_types = set([tf.uint8, tf.int8, tf.uint16, tf.int16, 
+ccl_supported_types = set([tf.uint8, tf.int8, tf.uint16, tf.int16,
                            tf.int32, tf.int64, tf.float32, tf.float64])
 
 _IS_TF2 = LooseVersion(tf.__version__) >= LooseVersion('2.0.0')
@@ -2689,7 +2689,7 @@ def test_elastic_state(self):
                 self.assertAllClose(w1, w2)
             assert state.batch == 21
             assert state.epoch == 11
-            
+
     def test_horovod_join_allreduce(self):
         """Test that the hvd.join with allreduce works on GPUs."""
         # Only do this test if there are GPUs available.
@@ -2719,13 +2719,13 @@ def test_horovod_join_allreduce(self):
                         [17] * dim, -100, 100, dtype=dtype)
                 if local_rank == first_join_rank:
                     self.evaluate(hvd.join())
-                else:		
+                else:
                     summed = hvd.allreduce(tensor, average=False)
                     multiplied = tensor * (size-1)
                     max_difference = tf.reduce_max(tf.abs(summed - multiplied))
 
                     if size <= 3 or dtype in [tf.int32, tf.int64]:
-                        threshold = 0 
+                        threshold = 0
                     elif size < 10:
                         threshold = 1e-4
                     elif size < 15:
@@ -2857,7 +2857,17 @@ def _apply_dense(self, grad, var):
             )
 
             grads_and_vars = opt.compute_gradients()
-            update_op = opt.apply_gradients(grads_and_vars)
+
+            # 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
+                global_step_before_update = global_step + 0
+            with tf.compat.v1.control_dependencies([global_step_before_update]):
+                update_op = opt.apply_gradients(grads_and_vars, global_step=global_step)
+
             sess.run(tf.compat.v1.global_variables_initializer())
             sess.run(tf.compat.v1.local_variables_initializer())
 
@@ -2877,11 +2887,15 @@ def compute_expected_value(batch_id):
                 aggregations_completed = math.floor((batch_id + 1) / backward_passes_per_step)
                 return aggregations_completed * sum_per_aggregation
 
+            steps_seen = []
             for idx in range(10):
-                _ = sess.run(update_op)
+                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)
 
+            assert steps_seen == list(range(10))
+
 
 from tensorflow.python.framework.test_util import run_all_in_graph_and_eager_modes
 run_all_in_graph_and_eager_modes(TensorFlowTests)