Distributed training using multiple GPUs hangs when enabling eager execution
@Liqi1003 

The issue you're experiencing is likely due to a combination of factors: eager execution, MirroredStrategy, and multiple GPUs. As the warning message indicates, using MirroredStrategy eagerly has significant overhead and is not currently recommended.

To resolve this issue, you can try the following:

Wrap the code inside a tf.function. This will help reduce the overhead of running MirroredStrategy eagerly, as suggested by the warning message. 

``

@tf.function
def train_step(inputs, labels):
    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)

    return model.fit(inputs, labels)

loss = train_step(train_input, train_label)

``
Disable eager execution by removing the following lines:
tf.config.run_functions_eagerly(True)
tf.data.experimental.enable_debug_mode()

Disabling eager execution may help resolve the issue as well, but it is essential to understand that your code will not run eagerly, and you may need to adjust your code accordingly if it relies on eager execution.

Please try these solutions and see if they resolve the issue. If the problem persists, further investigation may be required to determine the root cause Thanks @Adesoji1 for your pointers. 
@Liqi1003 , Usually when you run the multi GPU distributed training in the eager mode, for each GPU it creates a new session and these sessions will not be synchronized since these are running in the eager mode and executes immediately. 
Even though `tf.distribute.Srategy` works well with both `eager` and `tf.function`, it works best with `tf.function` and `eager` mode is recommended for debug mode.   Thanks for the responses.

@Adesoji1 
1. Wrapping the function inside `@tf.function`, I see the same behavior in both TF2.11 and TF nightly.
2. Turning off eager execution does solve the problem, but I wanted to use eager execution to help me debug the program, which is not working now.

@sachinprasadhs, I understand that `tf.distribute.Strategy` works best with `tf.function` rather than `eager` mode. My previous code was using function mode. However, I'm trying to use eager execution to help me debug. I think it is officially supported, but currently it simply hangs/crashes and I cannot use it.
 @Liqi1003 glad to hear that Turning off eager execution does solve the problem. @sachinprasadhs Is this issue still under investigation or do I need to provide further information? Currently I am still unable to use eager mode with distributed training in the nightly version.
<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>