Keras.fit stuck/error in TensorFlow 2.13/2.14 (TPU is fine, inference on GPU is fine, 2.11 GPU is fine)
Hi @edwardyehuang ,

Could you please confirm whether keras imported as stand alone package or from tf.keras? Also please submit some reproducible code snippet along with environment details.

Thanks! > Hi @edwardyehuang ,
> 
> Could you please confirm whether keras imported as stand alone package or from tf.keras? Also please submit some reproducible code snippet along with environment details.
> 
> Thanks!

1) tf.keras
2) It is hard to produce a code snippet because I am unsure what can cause this issue. I will continue to test more environments and see if I can figure out the root cause or limit it to a narrow range (after that, a code snippet is possible). I will update my progress on this issue. @SuryanarayanaY  Any place to download earlier versions of the tf-nightly wheel?  The [pip repo](https://pypi.org/project/tf-nightly/) only contains the post 2.14 package. Currently, I do not have time to build from sources. It looks like the commit that caused this issue is between 2.12 and 2.13rc0 @SuryanarayanaY 
Update:
After I use
 ```
tf.distribute.experimental.MultiWorkerMirroredStrategy(
    communication= tf.distribute.experimental.CollectiveCommunication.RING
)
```
to replace 
```
tf.distribute.MirroredStrategy
```
 2.13 and 2.14 training on GPU is no longer stuck.

It seems some commits caused issues in NCCL (maybe also XLA). Related issues : https://github.com/tensorflow/tensorflow/issues/41539 Pull requests may cause this issue: #60001  #59424   @SuryanarayanaY  > It looks like the commit that caused this issue is between 2.12 and 2.13rc0

Same issue here. I have 2 GPU. It works on 2.12 but fails on 2.13/2.14 I found two related settings.
 1. use strategy:   

 tf.distribute.experimental.MultiWorkerMirroredStrategy(
   communication= tf.distribute.experimental.CollectiveCommunication.RING
) 

2. add jit_compile=False to tf.keras.optimizers.Adam or use tf.keras.optimizers.legacy.Adam




Ugly solutions. But it works.
 > I found two related settings.
> 
> 1. use strategy:
> 
> tf.distribute.experimental.MultiWorkerMirroredStrategy( communication= tf.distribute.experimental.CollectiveCommunication.RING )
> 
> 2. add jit_compile=False to tf.keras.optimizers.Adam or use tf.keras.optimizers.legacy.Adam
> 
> Ugly solutions. But it works.

Seems like #59424 cause this issue @cheshire  @reedwm  Is it possible to revert 59424 in the next release of TensorFlow (e.g. 2.15 RC-1) ? @edwardyehuang do you have a way for us to reproduce this issue? Have you verified that #59424 is causing the issue? We cannot address this without some way to reproduce it. CC @jurahul @reedwm  See issue https://github.com/tensorflow/tensorflow/issues/61314 I have the same issue; I use tf 2.14 with Ubuntu 22.04 and Python 3.9 as well, but am using a GeForce RTX 3090, running a conda environment, installed tf 2.14 with pip and am using jupyter notebooks in VSCode. Sometimes the issue occurs when I train a model, but other times the issue occurs when I initialize a model - setting 


```python
tf.distribute.experimental.MultiWorkerMirroredStrategy(
   communication= tf.distribute.experimental.CollectiveCommunication.RING
)
``` 

does not work in my case (though I believe it removes the "Cannot start spawn child process" error); adding "jit_compile=false" doesn't seem to work well either.

I do have some steps to reproduce (assuming the command nvidia-smi "shows" a gpu); i.e. I have a notebook with code that one can run given the setup and the error should again occur, given one downloads the kaggle datasets that the notebook "trains" a model on.

1. conda deactivate (remove base environment)
2. conda create --name tfgittesting python=3.9
3. pip install --upgrade pip
4. pip install tensorflow[and-cuda]
5. pip install ipython ipykernel
6. pip install matplotlib
7. pip install pandas
8. Download the [Fashion MNIST dataset.](https://www.kaggle.com/datasets/zalando-research/fashionmnist?select=fashion-mnist_train.csv)
9. Download [this txt file](https://github.com/tensorflow/tensorflow/files/13325968/fashionmnisttest.txt) and replace txt with ipynb to make it a notebook.
10. Replace the words "Fashion MNIST test dataset" and "Fashion MNIST train dataset" with the paths to the test and train datasets, respectively.
11. Run the notebook; add the "mirrored strategy code" and the "jit_compile=false code" at some locations.
 *Could someone else test the code to see if the error can be reproduced? @reedwm @SuryanarayanaY ^ @reedwm  Is there any update on this issue? The issue is causing TensorFlow XLA and multi-gpu to become unusable for versions >= 2.13. This could be caused by switch to PJRT interface from TF, which has broken a number of multi-GPU workflows internally as well. 

@swachhandl @kenfranko for pointers. It is unlikely that this is due to the switch to PJRT - it hasn't been rolled out for GPU on OSS yet.  I just tested the conda-forge-packed TensorFlow GPU 2.14 (https://github.com/conda-forge/tensorflow-feedstock , note that the 2.15 is not working), and it is working; it may be easier for you to locate the error in the pip package and a docker image. The issue was resolved by upgrading the NVIDIA and CUDA drivers to the latest version.
The reason that the conda version can work is because it is compiled with a lower CUDA version.

However, it also means that before TensorFlow 2.16, the driver version check had BUG, but it is fine for now as TensorFlow 2.16 forces the NVIDIA driver >= 545 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/62234">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/62234">No</a>

### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

tf 2.14

### Custom code

Yes

### OS platform and distribution

Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.9

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

Titan RTX * 2 or 4090 * 4

### Current behavior?

Below are the situations in different environments:
  1. TPU + TensorFlow 2.10: OK
  2. TPU + TensorFlow 2.11: OK
  3. TPU + TensorFlow 2.14: OK
  4. GPU (Titan RTX * 2) + TensorFlow 2.10 (conda): OK
  5. GPU (Titan RTX * 2) + TensorFlow 2.11 (conda): OK
  6. GPU (Titan RTX * 2) + TensorFlow 2.11 (docker): OK
  7. GPU (Titan RTX * 2) + TensorFlow 2.12 (docker): OK
  8. GPU (Titan RTX * 2) + TensorFlow 2.13 (conda): OK
  9. GPU (Titan RTX * 2) + TensorFlow 2.13 (docker): Prediction OK. keras.fit stuck at `Loaded cuDNN version 8600`
10. GPU (Titan RTX * 2) + TensorFlow 2.14 (pip): Prediction OK. keras.fit stuck at `Start cannot spawn child process: No such file or directory` after `Loaded cuDNN version 8600`
11. GPU (Titan RTX * 2) + TensorFlow 2.14 (docker): Prediction OK. keras.fit stuck at `Loaded cuDNN version 8600`
12. GPU (4090 * 4) + TensorFlow 2.11 (conda): OK
13. GPU (4090 * 4) + TensorFlow 2.14 (docker): Prediction OK. keras.fit stuck at `Loaded cuDNN version 8600`
14. GPU (4090 * 4) + TensorFlow 2.14 (pip): Prediction OK. keras.fit stuck at `Start cannot spawn child process: No such file or directory` after `Loaded cuDNN version 8600`

I am currently testing more environments and see if I can narrow down the search space of problematic TensorFlow commits.


### Standalone code to reproduce the issue

```shell
See behavior
```


### Relevant log output

_No response_