[TF2.4 - HVD Keras Fix] - Horovod worker desynchronization - Potentially deadlock training TF 2.4

[DEKHTIARJonathan]

CC: @romerojosh @tgaddair @nluehr

This PR aims to fix a bug with TF2.4 already reported here: tensorflow/tensorflow#46863

Basically this PR does the following:

• Rollback the allreduce in _aggregate_gradients and make sure to execute the allreduce during gradient computation
• Yes it means that in the eventuality the Grad Step is skipped (by mixed precision) we have allreduced for no reason, but this event is supposed to be very rare. And it's still better than having worker to desynchronize or deadlock.
• self._aggregated_gradients = True has been moved to gradient computation instead of allreduce itself. This is done to work around that gradient tape is used to compute gradient with a CTL

@tgaddair : would be nice to push a new release" to address this bug, that way we can unblock TF2.4 users ;)

Thanks everyone

[tgaddair]

@DEKHTIARJonathan thanks for the PR! Seems this is breaking some tests, can you look into this?

[1]<stdout>:test_tensorflow2_keras.py::Tf2KerasTests::test_gradient_aggregation FAILED
--
  | [1]<stdout>:test_tensorflow2_keras.py::Tf2KerasTests::test_session SKIPPED (Not ...)
  | [0]<stderr>:[2021-02-05 20:33:38. 76607: W /tmp/pip-req-build-a5836dcu/horovod/common/stall_inspector.cc:105] One or more tensors were submitted to be reduced, gathered or broadcasted by subset of ranks and are waiting for remainder of ranks for more than 60 seconds. This may indicate that different ranks are trying to submit different tensors or that only subset of ranks is submitting tensors, which will cause deadlock.
  | [0]<stderr>:Missing ranks:
  | [0]<stderr>:0: [PartitionedCall/DistributedRMSprop_Allreduce/cond/then/_5/DistributedRMSprop_Allreduce/cond/HorovodAllreduce_DistributedRMSprop_Allreduce_UnsortedSegmentSum_0]
  | [0]<stderr>:1: [StatefulPartitionedCall/cond/then/_16/cond/PartitionedCall/DistributedTestingOptimizer_Allreduce/cond/then/_56/DistributedTestingOptimizer_Allreduce/cond/HorovodAllreduce_grads_0]

https://buildkite.com/horovod/horovod/builds/4887#34ba55ce-1e9b-4e4e-a824-efb41fbd7084/216-766

[DEKHTIARJonathan]

@romerojosh @tgaddair I have completely reworked the PR, we found duplicated calls to allreduce() that should address the fews bugs.

In addition, I had to skip with TF 2.4+ the test_gradient_aggregation() from test/parallel/test_tensorflow2_keras which naturally doesn't make any sense anymore since we don't allreduce anymore during _aggregate_gradients, so this failing test was totally expected.

Let's see how the CI comes out

[tgaddair]

Hey @DEKHTIARJonathan, can you explain the skipping of the gradient aggregation test? We still need to support this functionality, even if we don't call _aggregate_gradients explicitly.

[DEKHTIARJonathan]

This PR shifts the gradient allreduce when the gradient is computed (both through the DistributedOptimizer or through the DistributedGradientTape). Which means that this unittest, by design in TF2.4, doesn't call allreduce in _aggregate_gradients().

Since this unittest provide a gradient as constant (without effectively computing it), the gradient will never be allreduced.
The safest bet for this PR (which is to return to a working state) was to skip this unittest that conceptually doesn't make sense in the current direction for TF2.4.

Though if the usecase is still important, I recommend pushing a following PR doing the followings:

• fix test_gradient_aggregation() not to use tf.constant() for the gradient value and actually compute it for real.
• remove the pytest.skip

[tgaddair]

LGTM!

[github-actions]

Unit Test Results

     759 files  +  16       759 suites  +16   4h 51m 28s ⏱️ + 1m 32s
     559 tests ±    0       529 ✔️  -     1       29 💤 ±    0  1 ❌ +1 
15 264 runs  +521  11 614 ✔️ +351  3 649 💤 +169  1 ❌ +1 

For more details on these failures, see this check.

Results for commit 56cc923. ± Comparison against base commit c64b1d6.