Repository: tensorflow
Issue: Multi-GPU training with gradient propagation in FP16 with XLA fails.
Labels: stat:awaiting tensorflower,type:bug,comp:dist-strat,TF2.14

### 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]
```