Repository: tensorflow
Issue: Distributed training using multiple GPUs hangs when enabling eager execution
Labels: stat:awaiting tensorflower,type:bug,comp:gpu,TF 2.11

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

Bug

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

Yes

### Source

binary

### Tensorflow Version

tf 2.11.0

### Custom Code

Yes

### OS Platform and Distribution

Linux Ubuntu 20.04.5 LTS

### Mobile device

_No response_

### Python version

3.8.10 (docker tensorflow:devel-gpu)

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

11.7

### GPU model and memory

_No response_

### Current Behaviour?

```shell
After adding the following code to enable eager execution:
`
tf.config.run_functions_eagerly(True)
tf.data.experimental.enable_debug_mode()
`

Distributed training using MirroredStrategy with more than 2 GPUs hangs forever.

Using 1 GPU does not trigger the bug. Using MirroredStrategy with multiple logical CPUs also does not trigger the bug.

Whether adding the second line (`tf.data.experimental.enable_debug_mode()`) or not does not matter.

The model we used is a simple Sequential model trained on CIFAR-10 dataset.
The model structure is given as follows:
`
Input - GlobalMaxPool2D - BatchNormalization- Dense
`

Current behavior:
Using tf 2.11.0, the program hangs at the line of code:
`loss = model.fit(train_input, train_label)`
Also, the program cannot be stopped gracefully by keyboard interruption.

Using tf nightly, the program reports an error (see below).

Expected behavior:
The program should print the values of loss as an eager tensor.
```


### Standalone code to reproduce the issue

```shell
The code is also available at https://colab.research.google.com/drive/1BiAxZEO9IGak_4XOeZcyXRhJEByZ-X6B?usp=sharing

import keras
import tensorflow as tf

tf.config.run_functions_eagerly(True)
tf.data.experimental.enable_debug_mode()

if __name__ == "__main__":
    # Using at least 2 GPU triggers the bug
    strategy = tf.distribute.MirroredStrategy(["/GPU:0", "/GPU:1"])

    # Training data, batch_size=240
    train_input = tf.random.uniform(shape=(240, 32, 32, 3))
    train_label = tf.random.uniform(shape=(240, 10), minval=0, maxval=2, dtype=tf.int32)

    with strategy.scope():
        model = keras.Sequential([
            keras.layers.Input(shape=(32, 32, 3)),
            keras.layers.GlobalMaxPool2D(),
            keras.layers.BatchNormalization(axis=-1),
            keras.layers.Dense(10, ), ])
        optimizer = keras.optimizers.Adam(learning_rate=0.1)
        loss = tf.keras.losses.MeanSquaredError(
            reduction=tf.keras.losses.Reduction.NONE)
        model.compile(optimizer=optimizer, loss=loss)

    # One step training
    loss = model.fit(train_input, train_label)
    print(loss)
```


### Relevant log output

```shell
-------------- Output using tf 2.11.0 -----------------------------
tf-docker /mnt/src/reproduce > python ./reproduce_eager.py 
2023-04-05 03:23:45.884503: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-04-05 03:23:48.001421: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-04-05 03:23:53.074754: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 9636 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:3d:00.0, compute capability: 7.5
2023-04-05 03:23:53.076188: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 9636 MB memory:  -> device: 1, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:3f:00.0, compute capability: 7.5
2023-04-05 03:23:53.077442: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 9636 MB memory:  -> device: 2, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:40:00.0, compute capability: 7.5
2023-04-05 03:23:53.078666: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 9636 MB memory:  -> device: 3, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:41:00.0, compute capability: 7.5
2023-04-05 03:23:53.079887: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:4 with 9636 MB memory:  -> device: 4, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:61:00.0, compute capability: 7.5
2023-04-05 03:23:53.081109: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:5 with 9636 MB memory:  -> device: 5, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:62:00.0, compute capability: 7.5
2023-04-05 03:23:53.082252: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:6 with 9636 MB memory:  -> device: 6, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:63:00.0, compute capability: 7.5
2023-04-05 03:23:53.083395: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:7 with 9636 MB memory:  -> device: 7, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:64:00.0, compute capability: 7.5
WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
(program hangs here)
--------------------------------------------------------------------

-------------- Output using tf nightly -----------------------------
tf-docker /mnt/src/reproduce > python ./reproduce_eager.py 
2023-04-05 03:29:24.607254: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-04-05 03:29:25.559092: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2023-04-05 03:29:32.898888: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 9598 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:3d:00.0, compute capability: 7.5
2023-04-05 03:29:32.900404: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 9598 MB memory:  -> device: 1, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:3f:00.0, compute capability: 7.5
2023-04-05 03:29:32.901638: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 9598 MB memory:  -> device: 2, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:40:00.0, compute capability: 7.5
2023-04-05 03:29:32.902870: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 9598 MB memory:  -> device: 3, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:41:00.0, compute capability: 7.5
2023-04-05 03:29:32.904072: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:4 with 9598 MB memory:  -> device: 4, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:61:00.0, compute capability: 7.5
2023-04-05 03:29:32.905328: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:5 with 9598 MB memory:  -> device: 5, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:62:00.0, compute capability: 7.5
2023-04-05 03:29:32.906501: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:6 with 9598 MB memory:  -> device: 6, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:63:00.0, compute capability: 7.5
2023-04-05 03:29:32.907653: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1635] Created device /job:localhost/replica:0/task:0/device:GPU:7 with 9598 MB memory:  -> device: 7, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:64:00.0, compute capability: 7.5
2023-04-05 03:29:33.578635: I tensorflow/core/common_runtime/executor.cc:1210] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_9' with dtype int32 and shape [240,10]
         [[{{node Placeholder/_9}}]]
2023-04-05 03:29:33.578971: I tensorflow/core/common_runtime/executor.cc:1210] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'Placeholder/_9' with dtype int32 and shape [240,10]
         [[{{node Placeholder/_9}}]]
WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance.
2023-04-05 03:29:35.301668: E tensorflow/core/common_runtime/base_collective_executor.cc:249] BaseCollectiveExecutor::StartAbort UNKNOWN: Error invoking NCCL: unhandled cuda error
Traceback (most recent call last):
  File "./reproduce_eager.py", line 27, in <module>
    loss = model.fit(train_input, train_label)
  File "/usr/local/lib/python3.8/dist-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/lib/python3.8/multiprocessing/pool.py", line 364, in map
    return self._map_async(func, iterable, mapstar, chunksize).get()
  File "/usr/lib/python3.8/multiprocessing/pool.py", line 771, in get
    raise self._value
  File "/usr/lib/python3.8/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
  File "/usr/lib/python3.8/multiprocessing/pool.py", line 48, in mapstar
    return list(map(*args))
tensorflow.python.framework.errors_impl.UnknownError: {{function_node __wrapped__CollectiveReduceV2_Nordering_token_1_device_/job:localhost/replica:0/task:0/device:GPU:1}} Collective ops is aborted by: Error invoking NCCL: unhandled cuda error
The error could be from a previous operation. Restart your program to reset. [Op:CollectiveReduceV2] name:
--------------------------------------------------------------------
```
</details>