Tensorflow model training stalls with MirroredStrategy on multiple GPUs
Hi @JustASquid ,

Please submit a code snippet to replicate the issue.We can't say whether the issue is related to Distribution strategy bug or bug in code it self. As you mentioned Sage maker as GPU I am assuming you are using AWS Sage maker, and also going through the error log it seems the error might be due to NCCL communication problems. Please find the attached sources [link1](https://github.com/NVIDIA/nccl/issues/352) and [link2](https://github.com/NVIDIA/nccl/blob/master/ext-net/README.md) here and let us know if it helps.  This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60363">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60363">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

No

### Source

source

### Tensorflow Version

2.12

### Custom Code

Yes

### OS Platform and Distribution

Ubuntu 20.04

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

11.8

### GPU model and memory

Sagemaker ml.p3.8xlarge - 4x V100

### Current Behaviour?

When using `MirroredStrategy`, initializing training (after calling `fit`) takes an incredibly long time (over an hour), and often it stalls completely. See the attached logs - it stalls after the 4 `Init COMPLETE` messages are logged.

I'm unable to share the specific model and training code as it is proprietary. However, it's based around the Keras model `InceptionResNetV2`.

I'll try to work on putting together a shareable test-case, but in the meantime is there any way I can try and diagnose the issue myself?

### Standalone code to reproduce the issue

```shell
TBD
```


### Relevant log output

```shell
INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK
Your GPUs will likely run quickly with dtype policy mixed_float16 as they all have compute capability of at least 7.0
INFO:tensorflow:Mixed precision compatibility check (mixed_float16): OK
Your GPUs will likely run quickly with dtype policy mixed_float16 as they all have compute capability of at least 7.0
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).

...

INFO:tensorflow:batch_all_reduce: 499 all-reduces with algorithm = nccl, num_packs = 1
INFO:tensorflow:batch_all_reduce: 499 all-reduces with algorithm = nccl, num_packs = 1
INFO:tensorflow:batch_all_reduce: 499 all-reduces with algorithm = nccl, num_packs = 1
INFO:tensorflow:batch_all_reduce: 499 all-reduces with algorithm = nccl, num_packs = 1
4e409fc75fa0:30:152 [3] NCCL INFO Bootstrap : Using eth0:172.17.0.2<0>
4e409fc75fa0:30:152 [3] NCCL INFO NET/Plugin: Failed to find ncclCollNetPlugin_v6 symbol.
4e409fc75fa0:30:152 [3] NCCL INFO NET/Plugin: Failed to find ncclCollNetPlugin symbol (v4 or v5).
4e409fc75fa0:30:152 [0] NCCL INFO cudaDriverVersion 11080
NCCL version 2.16.5+cuda11.8
4e409fc75fa0:30:536 [0] NCCL INFO NET/OFI Using aws-ofi-nccl 1.5.0aws
4e409fc75fa0:30:536 [0] NCCL INFO NET/OFI Setting FI_EFA_FORK_SAFE environment variable to 1

4e409fc75fa0:30:536 [0] nccl_net_ofi_init:1444 NCCL WARN NET/OFI Only EFA provider is supported

4e409fc75fa0:30:536 [0] nccl_net_ofi_init:1483 NCCL WARN NET/OFI aws-ofi-nccl initialization failed
4e409fc75fa0:30:536 [0] NCCL INFO NET/IB : No device found.
4e409fc75fa0:30:536 [0] NCCL INFO NET/Socket : Using [0]eth0:172.17.0.2<0>
4e409fc75fa0:30:536 [0] NCCL INFO Using network Socket
4e409fc75fa0:30:539 [3] NCCL INFO Using network Socket
4e409fc75fa0:30:537 [1] NCCL INFO Using network Socket
4e409fc75fa0:30:538 [2] NCCL INFO Using network Socket
4e409fc75fa0:30:538 [2] NCCL INFO Trees [0] 1/-1/-1->2->3 [1] 3/-1/-1->2->1 [2] 1/-1/-1->2->3 [3] 3/-1/-1->2->1 [4] 1/-1/-1->2->3 [5] 3/-1/-1->2->1 [6] 1/-1/-1->2->3 [7] 3/-1/-1->2->1
4e409fc75fa0:30:538 [2] NCCL INFO P2P Chunksize set to 524288
4e409fc75fa0:30:537 [1] NCCL INFO Trees [0] -1/-1/-1->1->2 [1] 2/-1/-1->1->-1 [2] -1/-1/-1->1->2 [3] 2/-1/-1->1->-1 [4] -1/-1/-1->1->2 [5] 2/-1/-1->1->-1 [6] -1/-1/-1->1->2 [7] 2/-1/-1->1->-1
4e409fc75fa0:30:537 [1] NCCL INFO P2P Chunksize set to 524288
4e409fc75fa0:30:536 [0] NCCL INFO Channel 00/08 :    0   1   2   3
4e409fc75fa0:30:536 [0] NCCL INFO Channel 01/08 :    0   3   2   1
4e409fc75fa0:30:536 [0] NCCL INFO Channel 02/08 :    0   3   1   2
4e409fc75fa0:30:536 [0] NCCL INFO Channel 03/08 :    0   2   1   3
4e409fc75fa0:30:536 [0] NCCL INFO Channel 04/08 :    0   1   2   3
4e409fc75fa0:30:536 [0] NCCL INFO Channel 05/08 :    0   3   2   1
4e409fc75fa0:30:536 [0] NCCL INFO Channel 06/08 :    0   3   1   2
4e409fc75fa0:30:539 [3] NCCL INFO Trees [0] 2/-1/-1->3->0 [1] 0/-1/-1->3->2 [2] 2/-1/-1->3->0 [3] 0/-1/-1->3->2 [4] 2/-1/-1->3->0 [5] 0/-1/-1->3->2 [6] 2/-1/-1->3->0 [7] 0/-1/-1->3->2
4e409fc75fa0:30:539 [3] NCCL INFO P2P Chunksize set to 524288
4e409fc75fa0:30:536 [0] NCCL INFO Channel 07/08 :    0   2   1   3
4e409fc75fa0:30:536 [0] NCCL INFO Trees [0] 3/-1/-1->0->-1 [1] -1/-1/-1->0->3 [2] 3/-1/-1->0->-1 [3] -1/-1/-1->0->3 [4] 3/-1/-1->0->-1 [5] -1/-1/-1->0->3 [6] 3/-1/-1->0->-1 [7] -1/-1/-1->0->3
4e409fc75fa0:30:536 [0] NCCL INFO P2P Chunksize set to 524288
4e409fc75fa0:30:539 [3] NCCL INFO Channel 00/0 : 3[1e0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 00/0 : 1[1c0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 03/0 : 3[1e0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 02/0 : 1[1c0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 00/0 : 2[1d0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 04/0 : 3[1e0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 04/0 : 1[1c0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 00/0 : 0[1b0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 04/0 : 2[1d0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 07/0 : 3[1e0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 06/0 : 1[1c0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 04/0 : 0[1b0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 02/0 : 2[1d0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 02/0 : 3[1e0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 03/0 : 1[1c0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 03/0 : 0[1b0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 06/0 : 2[1d0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 06/0 : 3[1e0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 07/0 : 1[1c0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 07/0 : 0[1b0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 01/0 : 2[1d0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 01/0 : 0[1b0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 03/0 : 2[1d0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 02/0 : 0[1b0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 01/0 : 3[1e0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 01/0 : 1[1c0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 05/0 : 2[1d0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 05/0 : 0[1b0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 05/0 : 3[1e0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 05/0 : 1[1c0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 07/0 : 2[1d0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 06/0 : 0[1b0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Connected all rings
4e409fc75fa0:30:537 [1] NCCL INFO Connected all rings
4e409fc75fa0:30:537 [1] NCCL INFO Channel 01/0 : 1[1c0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Connected all rings
4e409fc75fa0:30:536 [0] NCCL INFO Connected all rings
4e409fc75fa0:30:537 [1] NCCL INFO Channel 03/0 : 1[1c0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 05/0 : 1[1c0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Channel 07/0 : 1[1c0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 01/0 : 2[1d0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 02/0 : 2[1d0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 03/0 : 2[1d0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 01/0 : 3[1e0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 05/0 : 2[1d0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 02/0 : 3[1e0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 06/0 : 2[1d0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 05/0 : 3[1e0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 07/0 : 2[1d0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 06/0 : 3[1e0] -> 0[1b0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 00/0 : 0[1b0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 03/0 : 0[1b0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 04/0 : 0[1b0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:536 [0] NCCL INFO Channel 07/0 : 0[1b0] -> 3[1e0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 00/0 : 3[1e0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 02/0 : 3[1e0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 00/0 : 2[1d0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 03/0 : 3[1e0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 02/0 : 2[1d0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 04/0 : 3[1e0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 04/0 : 2[1d0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 06/0 : 3[1e0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:538 [2] NCCL INFO Channel 06/0 : 2[1d0] -> 1[1c0] via P2P/direct pointer
4e409fc75fa0:30:539 [3] NCCL INFO Channel 07/0 : 3[1e0] -> 2[1d0] via P2P/direct pointer
4e409fc75fa0:30:537 [1] NCCL INFO Connected all trees
4e409fc75fa0:30:537 [1] NCCL INFO threadThresholds 8/8/64 | 32/8/64 | 512 | 512
4e409fc75fa0:30:537 [1] NCCL INFO 8 coll channels, 8 p2p channels, 2 p2p channels per peer
4e409fc75fa0:30:536 [0] NCCL INFO Connected all trees
4e409fc75fa0:30:538 [2] NCCL INFO Connected all trees
4e409fc75fa0:30:538 [2] NCCL INFO threadThresholds 8/8/64 | 32/8/64 | 512 | 512
4e409fc75fa0:30:538 [2] NCCL INFO 8 coll channels, 8 p2p channels, 2 p2p channels per peer
4e409fc75fa0:30:536 [0] NCCL INFO threadThresholds 8/8/64 | 32/8/64 | 512 | 512
4e409fc75fa0:30:536 [0] NCCL INFO 8 coll channels, 8 p2p channels, 2 p2p channels per peer
4e409fc75fa0:30:539 [3] NCCL INFO Connected all trees
4e409fc75fa0:30:539 [3] NCCL INFO threadThresholds 8/8/64 | 32/8/64 | 512 | 512
4e409fc75fa0:30:539 [3] NCCL INFO 8 coll channels, 8 p2p channels, 2 p2p channels per peer
4e409fc75fa0:30:538 [2] NCCL INFO comm 0x7f7cd93df490 rank 2 nranks 4 cudaDev 2 busId 1d0 commId 0xb631c8061e2f5303 - Init COMPLETE
4e409fc75fa0:30:536 [0] NCCL INFO comm 0x7f7cf2520cf0 rank 0 nranks 4 cudaDev 0 busId 1b0 commId 0xb631c8061e2f5303 - Init COMPLETE
4e409fc75fa0:30:539 [3] NCCL INFO comm 0x7f7cd93e1f20 rank 3 nranks 4 cudaDev 3 busId 1e0 commId 0xb631c8061e2f5303 - Init COMPLETE
4e409fc75fa0:30:537 [1] NCCL INFO comm 0x7f7ce1c85c50 rank 1 nranks 4 cudaDev 1 busId 1c0 commId 0xb631c8061e2f5303 - Init COMPLETE
```
</details>