Repository: tensorflow
Issue: TensorFlow keeps creating threads when multi-GPU training （thread leak）
Labels: stat:awaiting tensorflower,type:bug,comp:dist-strat,TF 2.11

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

2.11.0

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 20.04

### Mobile device

_No response_

### Python version

3.9.13

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

11.4

### GPU model and memory

Nvidia A800 

### Current behavior?

I was using this machine to train a GPT-2 example from (the data I used can also be found in this link, also here https://github.com/chinese-poetry/chinese-poetry.git) https://keras.io/examples/generative/gpt2_text_generation_with_kerasnlp/

Before we start, I would give a baseline amount of this machine's threads :
![10](https://github.com/user-attachments/assets/53e16381-e50f-4054-8c2a-b790fe2e077b)

When starting with the multi-GPU training of tensorflow by calling the tf.distribute.MirroredStrategy, the training process worked fine as usual. 

But with the time went by, I found the amount of threads increased with the training process going, here is the evidence of thread increasing when processing 3354th batch (I used cat /proc/"this programs' pid"/status to check the number of threads):
![3](https://github.com/user-attachments/assets/56bbfaf0-f59c-48f1-ae19-a98d563ad000)
![tf](https://github.com/user-attachments/assets/c546ea4b-855a-44e5-84c1-95d6d3f3aba4)
![2](https://github.com/user-attachments/assets/0d07b24d-e93e-451e-9e1e-51a26f7db60d)

Then, the evidence of thread increasing when processing 3791st batch (the amount of threads reached 22178):
![6](https://github.com/user-attachments/assets/508fed5c-8aec-4deb-8e68-74dbf6aa5613)
![4](https://github.com/user-attachments/assets/77507790-70fd-4abf-b72c-e9d831565879)
![5](https://github.com/user-attachments/assets/d8945e53-23d9-4fb7-b50d-37962b93a692)

When calculating  the 5054th batch, the training program got an error, and I captured a count of threads before the error (achieved around 31120 threads):
![8](https://github.com/user-attachments/assets/15b433cc-eefb-451f-a31c-2a286e17afec)
![8](https://github.com/user-attachments/assets/813835d2-3834-400a-b8fe-ec27f15d89ad)

I checked an similar issue in https://github.com/tensorflow/tensorflow/issues/62466, but I cannot find a solution, moreover, I have run other examples like diffusion model using this machine and the same tensorflow env with multi-GPU training, which worked fine and no any problems. So, could you please give me a help for this problem, very appreciated.





### Standalone code to reproduce the issue

```shell
Here is my code

import os

os.environ["KERAS_BACKEND"] = "tensorflow"  # or "tensorflow" or "torch"

import keras_nlp
import keras
import tensorflow as tf
import time

keras.mixed_precision.set_global_policy("mixed_float16")

import os
import json
import datetime

train_ds = (
    tf.data.Dataset.from_tensor_slices(paragraphs)
    .batch(36)
    .cache()
    .prefetch(tf.data.AUTOTUNE)
)

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1", "GPU:2", "GPU:3", "GPU:4", "GPU:5"])
print("Number of devices: {}".format(strategy.num_replicas_in_sync))

preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
    "gpt2_base_en",
    sequence_length=128,
)

# Open a strategy scope.
with strategy.scope():
# To speed up training and generation, we use preprocessor of length 128
# instead of full length 1024.
    gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
        "gpt2_base_en", preprocessor=preprocessor
    )
    num_epochs = 5

    checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
        filepath=checkpoint_path,
        save_weights_only=True,
        monitor="accuracy",
        # monitor="i_loss",
        mode="min",
        save_best_only=True,
        save_freq="epoch"
    )
    learning_rate = 5e-4
    
    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    gpt2_lm.compile(
        optimizer=keras.optimizers.Adam(learning_rate),
        loss=loss,
        weighted_metrics=["accuracy"],
    )

    gpt2_lm.fit(train_ds, epochs=num_epochs, callbacks=[
        checkpoint_callback,
        tensorboard_callback,
    ],
                )
```


### Relevant log output

```shell
2024-09-21 00:24:25.840848: W tensorflow/compiler/tf2xla/kernels/random_ops.cc:57] Warning: Using tf.random.uniform with XLA compilation will ignore seeds; consider using tf.random.stateless_uniform instead if reproducible behavior is desired. gpt2_causal_lm/gpt2_backbone/embeddings_dropout/dropout/random_uniform/RandomUniform
2024-09-21 00:24:25.846135: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
2024-09-21 00:24:26.559075: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/assert_equal_1/Assert/Assert
2024-09-21 00:24:26.572007: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/assert_equal_1/Assert/Assert
2024-09-21 00:24:26.573183: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/assert_equal_1/Assert/Assert
2024-09-21 00:24:26.581128: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/assert_equal_1/Assert/Assert
2024-09-21 00:24:26.581197: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/assert_equal_1/Assert/Assert
2024-09-21 00:24:26.588190: W tensorflow/compiler/tf2xla/kernels/assert_op.cc:38] Ignoring Assert operator sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/assert_equal_1/Assert/Assert
12024-09-21 00:25:01.986918: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:325] ptxas warning : Registers are spilled to local memory in function 'fusion_1267'
ptxas warning : Registers are spilled to local memory in function 'fusion_1225'
ptxas warning : Registers are spilled to local memory in function 'fusion_1111'
ptxas warning : Registers are spilled to local memory in function 'fusion_1220'
ptxas warning : Registers are spilled to local memory in function 'fusion_1124'
ptxas warning : Registers are spilled to local memory in function 'fusion_1175'
ptxas warning : Registers are spilled to local memory in function 'fusion_1128'
ptxas warning : Registers are spilled to local memory in function 'fusion_1006'
ptxas warning : Registers are spilled to local memory in function 'fusion_1015'

2024-09-21 00:25:02.215951: I tensorflow/compiler/jit/xla_compilation_cache.cc:477] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
2024-09-21 00:25:02.420703: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:325] ptxas warning : Registers are spilled to local memory in function 'fusion_1267'
ptxas warning : Registers are spilled to local memory in function 'fusion_1225'
ptxas warning : Registers are spilled to local memory in function 'fusion_1111'
ptxas warning : Registers are spilled to local memory in function 'fusion_1220'
ptxas warning : Registers are spilled to local memory in function 'fusion_1124'
ptxas warning : Registers are spilled to local memory in function 'fusion_1175'
ptxas warning : Registers are spilled to local memory in function 'fusion_1128'
ptxas warning : Registers are spilled to local memory in function 'fusion_1006'
ptxas warning : Registers are spilled to local memory in function 'fusion_1015'

2024-09-21 00:25:02.542130: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:325] ptxas warning : Registers are spilled to local memory in function 'fusion_1267'
ptxas warning : Registers are spilled to local memory in function 'fusion_1225'
ptxas warning : Registers are spilled to local memory in function 'fusion_1111'
ptxas warning : Registers are spilled to local memory in function 'fusion_1220'
ptxas warning : Registers are spilled to local memory in function 'fusion_1124'
ptxas warning : Registers are spilled to local memory in function 'fusion_1175'
ptxas warning : Registers are spilled to local memory in function 'fusion_1128'
ptxas warning : Registers are spilled to local memory in function 'fusion_1006'
ptxas warning : Registers are spilled to local memory in function 'fusion_1015'

2024-09-21 00:25:03.595774: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:325] ptxas warning : Registers are spilled to local memory in function 'fusion_1267'
ptxas warning : Registers are spilled to local memory in function 'fusion_1225'
ptxas warning : Registers are spilled to local memory in function 'fusion_1111'
ptxas warning : Registers are spilled to local memory in function 'fusion_1220'
ptxas warning : Registers are spilled to local memory in function 'fusion_1124'
ptxas warning : Registers are spilled to local memory in function 'fusion_1175'
ptxas warning : Registers are spilled to local memory in function 'fusion_1128'
ptxas warning : Registers are spilled to local memory in function 'fusion_1006'
ptxas warning : Registers are spilled to local memory in function 'fusion_1015'

12024-09-21 00:25:04.071526: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:325] ptxas warning : Registers are spilled to local memory in function 'fusion_1267'
ptxas warning : Registers are spilled to local memory in function 'fusion_1225'
ptxas warning : Registers are spilled to local memory in function 'fusion_1111'
ptxas warning : Registers are spilled to local memory in function 'fusion_1220'
ptxas warning : Registers are spilled to local memory in function 'fusion_1124'
ptxas warning : Registers are spilled to local memory in function 'fusion_1175'
ptxas warning : Registers are spilled to local memory in function 'fusion_1128'
ptxas warning : Registers are spilled to local memory in function 'fusion_1006'
ptxas warning : Registers are spilled to local memory in function 'fusion_1015'

2024-09-21 00:25:04.130504: I tensorflow/compiler/xla/stream_executor/gpu/asm_compiler.cc:325] ptxas warning : Registers are spilled to local memory in function 'fusion_1267'
ptxas warning : Registers are spilled to local memory in function 'fusion_1225'
ptxas warning : Registers are spilled to local memory in function 'fusion_1111'
ptxas warning : Registers are spilled to local memory in function 'fusion_1220'
ptxas warning : Registers are spilled to local memory in function 'fusion_1124'
ptxas warning : Registers are spilled to local memory in function 'fusion_1175'
ptxas warning : Registers are spilled to local memory in function 'fusion_1128'
ptxas warning : Registers are spilled to local memory in function 'fusion_1006'
ptxas warning : Registers are spilled to local memory in function 'fusion_1015'

2024-09-21 00:25:05.333297: I tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:630] TensorFloat-32 will be used for the matrix multiplication. This will only be logged once.
5054/8663 [================>.............] - ETA: 9:14 - loss: 11.8715 - accuracy: 2.2702terminate called after throwing an instance of 'std::system_error'
terminate called recursively
  what():  Resource temporarily unavailable
Aborted (core dumped)
```