Multi-GPU training with gradient propagation in FP16 with XLA fails.
@cyanic-selkie The error in the outcome would be because of the following reasons;
a. Incorrect loss scaling factor leading to overflows or underflow.
b. Incompatible data types between layers or operations.
c. Casting issues between FP16 and FP32 during gradient propagation.
Please check if XLA not enabled or not being used for all relevant operations?
Thank you!
 Hi, @sushreebarsa. How do I check that? @cyanic-selkie To check if XLA is enabled or not please use , `tf.config.list_logical_devices()'
 If XLA is disabled, you won't see any XLA devices listed in the output.
```
import tensorflow as tf

print(tf.config.list_logical_devices())
```
@sachinprasadhs Please have a look at this issue.
Thank you! Sorry for the delayed response, here is the result @sachinprasadhs:

```
[LogicalDevice(name='/job:worker/replica:0/task:0/device:CPU:0', device_type='CPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:GPU:0', device_type='GPU'), LogicalDevice(name='/job:worker/replica:0/task:1/device:CPU:0', device_type='CPU'), LogicalDevice(name='/job:worker/replica:0/task:1/device:GPU:0', device_type='GPU')]
``` Hi. I have the same problem when using Mirrored strategy for ditributed training and jit compiling some functions called inside the training loop step. Any news regarding this @cyanic-selkie @sushreebarsa ? @andremfreitas Unfortunately, no. In fact, all of my opened issues ultimately went unanswered. It's no secret that TF is internally deprecated, and I would suggest switching to PyTorch or JAX/Flax, perhaps using Keras 3 as a stepping stone if legacy code is a big issue for you.
### Issue type

Bug

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

Yes

### Source

binary

### TensorFlow version

tf 2.14

### Custom code

Yes

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

11.8/8.7

### GPU model and memory

A100 40GB (20GB MiG)

### Current behavior?

Following advice given [here](https://www.youtube.com/watch?v=6ovfZW8pepo), I implemented gradient propagation in FP16 and concurrent backpropagation with gradient propagation. I am using mixed precision with XLA enabled.

The custom training step works fine when XLA is not enabled. However, when I enable it, I get the error that you can see in the log.

This is run on 2 GPUs using MultiWorkerMirroredStrategy.

### Standalone code to reproduce the issue

```python
import json
import os

import tensorflow as tf


def get_replica_hostnames():
    ...


def get_replica_id():
    ...


def set_multiworker_env_config():
    hostnames = get_replica_hostnames()
    replica_index = get_replica_id()

    os.environ["TF_CONFIG"] = json.dumps(
        {
            "cluster": {
                "worker": hostnames,
            },
            "task": {"type": "worker", "index": replica_index},
        }
    )


class Model(tf.keras.models.Model):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)

        self._embedder = tf.keras.Sequential(
            [
                tf.keras.layers.Conv2D(
                    filters=8,
                    kernel_size=3,
                    padding="same",
                    activation=tf.keras.activations.relu,
                    use_bias=False,
                ),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.Conv2D(
                    filters=8,
                    kernel_size=3,
                    padding="same",
                    activation=tf.keras.activations.relu,
                    use_bias=False,
                ),
                tf.keras.layers.BatchNormalization(),
                tf.keras.layers.MaxPool2D(),
                tf.keras.layers.GlobalAveragePooling2D(),
            ]
        )

        self._classifier = tf.keras.layers.Dense(550)

    def call(self, x: tf.Tensor) -> tf.Tensor:
        x = self._embedder(x)
        x = self._classifier(x)
        x = tf.keras.layers.Activation("linear", dtype="float32")(x)
        return x

    def train_step(self, data):
        x, y, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data)

        # Run forward pass.
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compute_loss(x, y, y_pred, sample_weight)
            loss = self.optimizer.get_scaled_loss(loss)

        self._validate_target_and_loss(y, loss)

        # Run backward pass.
        grads = tape.gradient(loss, self.trainable_variables)
        grads = [tf.cast(grad, "float16") for grad in grads]
        all_reduced_grads = tf.distribute.get_replica_context().all_reduce(
            tf.distribute.ReduceOp.SUM,
            grads,
            tf.distribute.experimental.CommunicationOptions(
                bytes_per_pack=50 * 1024 * 1024,
            ),
        )
        all_reduced_grads = [tf.cast(grad, "float32") for grad in all_reduced_grads]
        all_reduced_grads = self.optimizer.get_unscaled_gradients(all_reduced_grads)

        self.optimizer.apply_gradients(
            zip(all_reduced_grads, self.trainable_variables),
            skip_gradients_aggregations=True,
        )

        return self.compute_metrics(x, y, y_pred, sample_weight)


def create_dummy_dataset(batch_size: int) -> tf.data.Dataset:
    X = np.random.rand(batch_size, 384, 640, 1)
    y = np.random.randint(550, size=batch_size)
    return tf.data.Dataset.from_tensor_slices((X, y)).batch(batch_size, True).repeat()


def train():
    tf.keras.mixed_precision.set_global_policy("mixed_float16")

    set_multiworker_env_config()
    strategy = tf.distribute.MultiWorkerMirroredStrategy()
    num_replicas = strategy.num_replicas_in_sync

    batch_size = 384 * num_replicas
    dataset = create_dummy_dataset(batch_size)
    dataset = strategy.experimental_distribute_dataset(dataset)

    with strategy.scope():
        model = Model()
        model.compile(
            loss=tf.keras.losses.SparseCategoricalCrossentropy(
                from_logits=True,
            ),
            optimizer=tf.keras.optimizers.Adam(
                learning_rate=1e-3,
                weight_decay=1e-5,
            ),
            metrics=[
                "accuracy",
            ],
            jit_compile=True,
        )

    model.fit(
        dataset,
        epochs=10,
        steps_per_epoch=100,
    )


if __name__ == "__main__":
    train()
```


### Relevant log output

```shell
2024-01-15 14:06:48.704206: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-01-15 14:06:48.704294: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-01-15 14:06:48.704350: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-01-15 14:06:48.712047: 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 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-01-15 14:06:49.711421: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2024-01-15 14:06:51.459926: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 18184 MB memory:  -> device: 0, name: NVIDIA A100-SXM4-40GB MIG 3g.20gb, pci bus id: 0000:01:00.0, compute capability: 8.0
2024-01-15 14:06:51.469378: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1886] Created device /job:worker/replica:0/task:0/device:GPU:0 with 18184 MB memory:  -> device: 0, name: NVIDIA A100-SXM4-40GB MIG 3g.20gb, pci bus id: 0000:01:00.0, compute capability: 8.0
2024-01-15 14:06:51.490217: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:457] Started server with target: grpc://gen-svc-82ee2ab7-aae8-44b2-a5ab-1f9fd0401075:80
2024-01-15 14:06:51.495668: I tensorflow/tsl/distributed_runtime/coordination/coordination_service.cc:551] /job:worker/replica:0/task:0 has connected to coordination service. Incarnation: 3072789976930509670
2024-01-15 14:06:51.495948: I tensorflow/tsl/distributed_runtime/coordination/coordination_service_agent.cc:299] Coordination agent has successfully connected.
2024-01-15 14:06:52.263966: I tensorflow/tsl/distributed_runtime/coordination/coordination_service.cc:551] /job:worker/replica:0/task:1 has connected to coordination service. Incarnation: 14088291762770279744
2024-01-15 14:07:16.278224: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 1509949440 exceeds 10% of free system memory.
2024-01-15 14:07:17.178980: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 1509949440 exceeds 10% of free system memory.
2024-01-15 14:07:17.905980: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 1509949440 exceeds 10% of free system memory.
2024-01-15 14:07:22.444540: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 1509949440 exceeds 10% of free system memory.
Epoch 1/10
2024-01-15 14:07:24.533902: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 1509949440 exceeds 10% of free system memory.
2024-01-15 14:07:27.793986: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7f080000c190 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2024-01-15 14:07:27.794120: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): NVIDIA A100-SXM4-40GB MIG 3g.20gb, Compute Capability 8.0
2024-01-15 14:07:28.936911: 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-01-15 14:07:30.163125: I tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:442] Loaded cuDNN version 8906
2024-01-15 14:08:24.307063: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:328] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_232', 12 bytes spill stores, 12 bytes spill loads

2024-01-15 14:08:24.726556: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:328] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_232', 12 bytes spill stores, 12 bytes spill loads

2024-01-15 14:08:24.789581: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:328] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_232', 8 bytes spill stores, 8 bytes spill loads

2024-01-15 14:08:26.317751: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:328] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_392', 92 bytes spill stores, 92 bytes spill loads

2024-01-15 14:08:26.687122: W tensorflow/core/framework/op_kernel.cc:1839] OP_REQUIRES failed at xla_ops.cc:624 : UNKNOWN: <unknown>:0: error: loc("all-reduce-start.1"): 'lmhlo_gpu.all_reduce_start' op requires the same element type for all operands
<unknown>:0: note: loc("all-reduce-start.1"): see current operation: 
%225 = "lmhlo_gpu.all_reduce_start"(%222, %8, %224, %18, %220, %222, %8, %224, %18, %220) ({
^bb0(%arg58: tensor<f32>, %arg59: tensor<f32>):
  %251 = "mhlo.add"(%arg58, %arg59) : (tensor<f32>, tensor<f32>) -> tensor<f32>
  "mhlo.return"(%251) : (tensor<f32>) -> ()
}) {channel_id = #mhlo.channel_handle<handle = 4294967300, type = 0>, constrain_layout = false, is_sync = true, replica_groups = dense<> : tensor<0x0xi64>, use_global_device_ids = false} : (memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<5630xf16>, memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<5630xf16>) -> !mhlo.token

2024-01-15 14:08:26.687331: I tensorflow/core/framework/local_rendezvous.cc:421] Local rendezvous recv item cancelled. Key hash: 10204104846648882243
2024-01-15 14:08:26.687358: I tensorflow/core/framework/local_rendezvous.cc:421] Local rendezvous recv item cancelled. Key hash: 5394287870414942276
2024-01-15 14:08:26.687399: I tensorflow/core/framework/local_rendezvous.cc:421] Local rendezvous recv item cancelled. Key hash: 8969947325169437454
Traceback (most recent call last):
  File "/kirax_source/train.py", line 133, in <module>
    train()
  File "/kirax_source/train.py", line 125, in train
    model.fit(
  File "/usr/local/lib/python3.11/dist-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.11/dist-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
tensorflow.python.framework.errors_impl.UnknownError: Graph execution error:

Detected at node StatefulPartitionedCall defined at (most recent call last):
  File "/usr/lib/python3.11/threading.py", line 1002, in _bootstrap

  File "/usr/lib/python3.11/threading.py", line 1045, in _bootstrap_inner

Detected at node StatefulPartitionedCall defined at (most recent call last):
  File "/usr/lib/python3.11/threading.py", line 1002, in _bootstrap

  File "/usr/lib/python3.11/threading.py", line 1045, in _bootstrap_inner

2 root error(s) found.
  (0) UNKNOWN:  <unknown>:0: error: loc("all-reduce-start.1"): 'lmhlo_gpu.all_reduce_start' op requires the same element type for all operands
<unknown>:0: note: loc("all-reduce-start.1"): see current operation: 
%225 = "lmhlo_gpu.all_reduce_start"(%222, %8, %224, %18, %220, %222, %8, %224, %18, %220) ({
^bb0(%arg58: tensor<f32>, %arg59: tensor<f32>):
  %251 = "mhlo.add"(%arg58, %arg59) : (tensor<f32>, tensor<f32>) -> tensor<f32>
  "mhlo.return"(%251) : (tensor<f32>) -> ()
}) {channel_id = #mhlo.channel_handle<handle = 4294967300, type = 0>, constrain_layout = false, is_sync = true, replica_groups = dense<> : tensor<0x0xi64>, use_global_device_ids = false} : (memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<5630xf16>, memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<5630xf16>) -> !mhlo.token

         [[{{node StatefulPartitionedCall}}]]
         [[Reshape_3/_24]]
  (1) UNKNOWN:  <unknown>:0: error: loc("all-reduce-start.1"): 'lmhlo_gpu.all_reduce_start' op requires the same element type for all operands
<unknown>:0: note: loc("all-reduce-start.1"): see current operation: 
%225 = "lmhlo_gpu.all_reduce_start"(%222, %8, %224, %18, %220, %222, %8, %224, %18, %220) ({
^bb0(%arg58: tensor<f32>, %arg59: tensor<f32>):
  %251 = "mhlo.add"(%arg58, %arg59) : (tensor<f32>, tensor<f32>) -> tensor<f32>
  "mhlo.return"(%251) : (tensor<f32>) -> ()
}) {channel_id = #mhlo.channel_handle<handle = 4294967300, type = 0>, constrain_layout = false, is_sync = true, replica_groups = dense<> : tensor<0x0xi64>, use_global_device_ids = false} : (memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<5630xf16>, memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<1xf32>, memref<5630xf16>) -> !mhlo.token

         [[{{node StatefulPartitionedCall}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_2528]
```
