Repository: tensorflow
Issue: Tutorial "Multi-worker training with Keras" fails to complete
Labels: type:bug,TF 2.18

### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

v1.12.1-120353-gc5bd67bc56f 2.19.0-dev20250107

### Custom code

No

### OS platform and distribution

Debian 6.1.123-1 (2025-01-02) x86_64 GNU/Linux

### Mobile device

_No response_

### Python version

Python 3.12.8

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Following the tutorial everything goes well until you start the second worker. Then the below failure occures.

2025-01-20 07:19:35.283801: I tensorflow/core/util/port.cc:153] 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`.
2025-01-20 07:19:35.290192: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1737379175.297785    4595 cuda_dnn.cc:8501] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1737379175.300054    4595 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-01-20 07:19:35.307721: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2025-01-20 07:19:36.510476: E external/local_xla/xla/stream_executor/cuda/cuda_platform.cc:51] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
2025-01-20 07:19:36.510494: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:167] env: CUDA_VISIBLE_DEVICES="-1"
2025-01-20 07:19:36.510499: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:170] CUDA_VISIBLE_DEVICES is set to -1 - this hides all GPUs from CUDA
2025-01-20 07:19:36.510501: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:178] verbose logging is disabled. Rerun with verbose logging (usually --v=1 or --vmodule=cuda_diagnostics=1) to get more diagnostic output from this module
2025-01-20 07:19:36.510505: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:183] retrieving CUDA diagnostic information for host: michael
2025-01-20 07:19:36.510507: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:190] hostname: michael
2025-01-20 07:19:36.510562: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:197] libcuda reported version is: 565.77.0
2025-01-20 07:19:36.510572: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:201] kernel reported version is: 565.77.0
2025-01-20 07:19:36.510574: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:291] kernel version seems to match DSO: 565.77.0
2025-01-20 07:19:36.519175: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:637] Initializing CoordinationService
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1737379176.519611    4595 grpc_server_lib.cc:465] Started server with target: grpc://localhost:12345
2025-01-20 07:19:36.524874: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:378] /job:worker/replica:0/task:0 has connected to coordination service. Incarnation: 4677280066871850635
2025-01-20 07:19:36.524894: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:816] Waiting for 1/2 tasks to connect.
2025-01-20 07:19:36.524898: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:819] Example stragglers:
/job:worker/replica:0/task:1
I0000 00:00:1737379176.525022    4595 coordination_service_agent.cc:369] Coordination agent has successfully connected.
2025-01-20 07:22:27.996664: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:378] /job:worker/replica:0/task:1 has connected to coordination service. Incarnation: 13530699364709055870
2025-01-20 07:22:27.996686: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:816] Waiting for 0/2 tasks to connect.
/home/chad/anaconda3/lib/python3.12/site-packages/keras/src/layers/core/input_layer.py:27: UserWarning: Argument `input_shape` is deprecated. Use `shape` instead.
  warnings.warn(
2025-01-20 07:22:28.461733: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
Traceback (most recent call last):
  File "/home/chad/Documents/McCueFiles/NeuralNetworks/TensorFlowProject/TensorFlowDocExample/main.py", line 21, in <module>
    multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70)
  File "/home/chad/anaconda3/lib/python3.12/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/chad/anaconda3/lib/python3.12/site-packages/tensorflow/python/framework/constant_op.py", line 108, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: Attempt to convert a value (PerReplica:{
  0: <tf.Tensor: shape=(64, 28, 28), dtype=float32, numpy=
array([[[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       ...,

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]]], dtype=float32)>
}) with an unsupported type (<class 'tensorflow.python.distribute.values.PerReplica'>) to a Tensor.

### Standalone code to reproduce the issue

```shell
python main.py &> job_1.log
```

### Relevant log output

```shell
2025-01-20 07:19:35.283801: I tensorflow/core/util/port.cc:153] 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`.
2025-01-20 07:19:35.290192: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1737379175.297785    4595 cuda_dnn.cc:8501] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1737379175.300054    4595 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-01-20 07:19:35.307721: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2025-01-20 07:19:36.510476: E external/local_xla/xla/stream_executor/cuda/cuda_platform.cc:51] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
2025-01-20 07:19:36.510494: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:167] env: CUDA_VISIBLE_DEVICES="-1"
2025-01-20 07:19:36.510499: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:170] CUDA_VISIBLE_DEVICES is set to -1 - this hides all GPUs from CUDA
2025-01-20 07:19:36.510501: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:178] verbose logging is disabled. Rerun with verbose logging (usually --v=1 or --vmodule=cuda_diagnostics=1) to get more diagnostic output from this module
2025-01-20 07:19:36.510505: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:183] retrieving CUDA diagnostic information for host: michael
2025-01-20 07:19:36.510507: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:190] hostname: michael
2025-01-20 07:19:36.510562: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:197] libcuda reported version is: 565.77.0
2025-01-20 07:19:36.510572: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:201] kernel reported version is: 565.77.0
2025-01-20 07:19:36.510574: I external/local_xla/xla/stream_executor/cuda/cuda_diagnostics.cc:291] kernel version seems to match DSO: 565.77.0
2025-01-20 07:19:36.519175: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:637] Initializing CoordinationService
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1737379176.519611    4595 grpc_server_lib.cc:465] Started server with target: grpc://localhost:12345
2025-01-20 07:19:36.524874: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:378] /job:worker/replica:0/task:0 has connected to coordination service. Incarnation: 4677280066871850635
2025-01-20 07:19:36.524894: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:816] Waiting for 1/2 tasks to connect.
2025-01-20 07:19:36.524898: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:819] Example stragglers:
/job:worker/replica:0/task:1
I0000 00:00:1737379176.525022    4595 coordination_service_agent.cc:369] Coordination agent has successfully connected.
2025-01-20 07:22:27.996664: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:378] /job:worker/replica:0/task:1 has connected to coordination service. Incarnation: 13530699364709055870
2025-01-20 07:22:27.996686: I external/local_xla/xla/tsl/distributed_runtime/coordination/coordination_service.cc:816] Waiting for 0/2 tasks to connect.
/home/chad/anaconda3/lib/python3.12/site-packages/keras/src/layers/core/input_layer.py:27: UserWarning: Argument `input_shape` is deprecated. Use `shape` instead.
  warnings.warn(
2025-01-20 07:22:28.461733: W tensorflow/core/framework/dataset.cc:993] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
Traceback (most recent call last):
  File "/home/chad/Documents/McCueFiles/NeuralNetworks/TensorFlowProject/TensorFlowDocExample/main.py", line 21, in <module>
    multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70)
  File "/home/chad/anaconda3/lib/python3.12/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/chad/anaconda3/lib/python3.12/site-packages/tensorflow/python/framework/constant_op.py", line 108, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: Attempt to convert a value (PerReplica:{
  0: <tf.Tensor: shape=(64, 28, 28), dtype=float32, numpy=
array([[[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       ...,

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]]], dtype=float32)>
}) with an unsupported type (<class 'tensorflow.python.distribute.values.PerReplica'>) to a Tensor.
```