tf.distribute.MirroredStrategy giving NaN in loss and accuracy with 4090 Multi-GPU (x2-4)
@maziarzamani,
I tried to execute the mentioned code on both single-gpu & multi-gpu and observed that the Loss and accuracy is not Nan values. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/9fcda1c4363966e1151afc003220194d/untitled1725.ipynb) and screen shot of the multi-gpu for the reference.


![62915_GitHub](https://github.com/tensorflow/tensorflow/assets/81610181/c928fc1a-3d5e-4ed0-8821-1b233dd1bade)


Also please have a look at this similar issue for the reference.
https://github.com/tensorflow/tensorflow/issues/41657
https://github.com/tensorflow/tensorflow/issues/36224
Thank you! 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/62915">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/62915">No</a>

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

v2.15.0-2-g0b15fdfcb3f 2.15.0

### Custom code

Yes

### OS platform and distribution

Ubuntu 22.04.3 LTS

### Mobile device

_No response_

### Python version

Python 3.10.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

12.3

### GPU model and memory

4090 / 24GB

### Current behavior?

Hardware configuration:
- Intel W3475
- 4* NVIDIA RTX 4090
- 128GB DDR5 ECC RAM
- Asus W790E-SAGE SE

When running Multi-GPU training with MirroredStrategy I get NaN in loss after a short while of running the first epoch.

Running a Single-GPU works fine.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import keras


def get_compiled_model():
    # Make a simple 2-layer densely-connected neural network.
    inputs = keras.Input(shape=(784,))
    x = keras.layers.Dense(256, activation="relu")(inputs)
    x = keras.layers.Dense(256, activation="relu")(x)
    outputs = keras.layers.Dense(10)(x)
    model = keras.Model(inputs, outputs)
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )
    return model


def get_dataset():
    batch_size = 32
    num_val_samples = 10000

    # Return the MNIST dataset in the form of a `tf.data.Dataset`.
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

    # Preprocess the data (these are Numpy arrays)
    x_train = x_train.reshape(-1, 784).astype("float32") / 255
    x_test = x_test.reshape(-1, 784).astype("float32") / 255
    y_train = y_train.astype("float32")
    y_test = y_test.astype("float32")

    # Reserve num_val_samples samples for validation
    x_val = x_train[-num_val_samples:]
    y_val = y_train[-num_val_samples:]
    x_train = x_train[:-num_val_samples]
    y_train = y_train[:-num_val_samples]
    return (
        tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size),
        tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(batch_size),
        tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size),
    )


# Create a MirroredStrategy.
strategy = tf.distribute.MirroredStrategy()
print("Number of devices: {}".format(strategy.num_replicas_in_sync))

# Open a strategy scope.
with strategy.scope():
    # Everything that creates variables should be under the strategy scope.
    # In general this is only model construction & `compile()`.
    model = get_compiled_model()

# Train the model on all available devices.
train_dataset, val_dataset, test_dataset = get_dataset()
model.fit(train_dataset, epochs=2, validation_data=val_dataset)

# Test the model on all available devices.
model.evaluate(test_dataset)
```


### Relevant log output

```shell
python3 test.py 
2024-02-07 12:18:53.214106: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-02-07 12:18:53.239167: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-02-07 12:18:53.239189: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-02-07 12:18:53.239879: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-02-07 12:18:53.243863: 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 AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-02-07 12:18:53.654794: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2024-02-07 12:18:54.616202: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1929] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 22284 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 4090, pci bus id: 0000:16:00.0, compute capability: 8.9
2024-02-07 12:18:54.616782: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1929] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 22287 MB memory:  -> device: 1, name: NVIDIA GeForce RTX 4090, pci bus id: 0000:34:00.0, compute capability: 8.9
2024-02-07 12:18:54.617432: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1929] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 22287 MB memory:  -> device: 2, name: NVIDIA GeForce RTX 4090, pci bus id: 0000:52:00.0, compute capability: 8.9
2024-02-07 12:18:54.617945: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1929] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 22287 MB memory:  -> device: 3, name: NVIDIA GeForce RTX 4090, pci bus id: 0000:70:00.0, compute capability: 8.9
Number of devices: 4
Epoch 1/2
2024-02-07 12:18:57.622003: I external/local_tsl/tsl/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory
2024-02-07 12:18:58.617348: I external/local_xla/xla/service/service.cc:168] XLA service 0x7fdad6459df0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2024-02-07 12:18:58.617381: I external/local_xla/xla/service/service.cc:176]   StreamExecutor device (0): NVIDIA GeForce RTX 4090, Compute Capability 8.9
2024-02-07 12:18:58.617384: I external/local_xla/xla/service/service.cc:176]   StreamExecutor device (1): NVIDIA GeForce RTX 4090, Compute Capability 8.9
2024-02-07 12:18:58.617387: I external/local_xla/xla/service/service.cc:176]   StreamExecutor device (2): NVIDIA GeForce RTX 4090, Compute Capability 8.9
2024-02-07 12:18:58.617390: I external/local_xla/xla/service/service.cc:176]   StreamExecutor device (3): NVIDIA GeForce RTX 4090, Compute Capability 8.9
2024-02-07 12:18:58.622639: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2024-02-07 12:18:58.650080: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:454] Loaded cuDNN version 8904
2024-02-07 12:18:58.652158: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:454] Loaded cuDNN version 8904
2024-02-07 12:18:58.653847: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:454] Loaded cuDNN version 8904
2024-02-07 12:18:58.655977: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:454] Loaded cuDNN version 8904
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1707308338.717054   15378 device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
1563/1563 [==============================] - 10s 5ms/step - loss: nan - sparse_categorical_accuracy: 0.2065 - val_loss: nan - val_sparse_categorical_accuracy: 2.6972
Epoch 2/2
1563/1563 [==============================] - 7s 4ms/step - loss: nan - sparse_categorical_accuracy: nan - val_loss: nan - val_sparse_categorical_accuracy: nan
313/313 [==============================] - 1s 2ms/step - loss: nan - sparse_categorical_accuracy: nan
```
```
