Repository: tensorflow
Issue: Memory Allocation Issues
Labels: stat:awaiting response,type:bug,stale,TF 2.18

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.18.0

### Custom code

Yes

### OS platform and distribution

Ubuntu 24.04.1 LTS on WSL2

### Mobile device

_No response_

### Python version

3.12.8

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

90300

### GPU model and memory

RTX 4070 12GB

### Current behavior?

I create a virtual gpu with a hard limit of 10GB. I start training the network and it works for bit but then says out of memory and tries to allocate more than the set limit. What I expect to happen is that is stays within the 10GB limit and can train the network successfully.

### Standalone code to reproduce the issue

```shell
import numpy as np
import keras
from keras import layers
import tensorflow as tf
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt

%matplotlib inline

tfds.disable_progress_bar()

gpus = tf.config.list_physical_devices('GPU')

if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=10240)])
    logical_gpus = tf.config.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)

train_ds, validation_ds, test_ds = tfds.load(
    "cats_vs_dogs",
    # Reserve 10% for validation and 10% for test
    split=["train[:40%]", "train[40%:50%]", "train[50%:60%]"],
    as_supervised=True,  # Include labels
)

resize_fn = keras.layers.Resizing(150, 150)
train_ds = train_ds.map(lambda x, y: (resize_fn(x), y))
validation_ds = validation_ds.map(lambda x, y: (resize_fn(x), y))
test_ds = test_ds.map(lambda x, y: (resize_fn(x), y))

augmentation_layers = [
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.1),
]


def data_augmentation(x):
    for layer in augmentation_layers:
        x = layer(x)
    return x


train_ds = train_ds.map(lambda x, y: (data_augmentation(x), y))

from tensorflow import data as tf_data

batch_size = 16

train_ds = train_ds.batch(batch_size).prefetch(tf_data.AUTOTUNE).cache()
validation_ds = validation_ds.batch(batch_size).prefetch(tf_data.AUTOTUNE).cache()
test_ds = test_ds.batch(batch_size).prefetch(tf_data.AUTOTUNE).cache()

base_model = keras.applications.Xception(
    weights="imagenet",  # Load weights pre-trained on ImageNet.
    input_shape=(150, 150, 3),
    include_top=False,
)  # Do not include the ImageNet classifier at the top.

# Freeze the base_model
base_model.trainable = False

# Create new model on top
inputs = keras.Input(shape=(150, 150, 3))

# Pre-trained Xception weights requires that input be scaled
# from (0, 255) to a range of (-1., +1.), the rescaling layer
# outputs: `(inputs * scale) + offset`
scale_layer = keras.layers.Rescaling(scale=1 / 127.5, offset=-1)
x = scale_layer(inputs)

# The base model contains batchnorm layers. We want to keep them in inference mode
# when we unfreeze the base model for fine-tuning, so we make sure that the
# base_model is running in inference mode here.
x = base_model(x, training=False)
x = keras.layers.GlobalAveragePooling2D()(x)
x = keras.layers.Dropout(0.2)(x)  # Regularize with dropout
outputs = keras.layers.Dense(1)(x)
model = keras.Model(inputs, outputs)

model.summary(show_trainable=True)

model.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.BinaryCrossentropy(from_logits=True),
    metrics=[keras.metrics.BinaryAccuracy()],
)

epochs = 2
print("Fitting the top layer of the model")
model.fit(train_ds, epochs=epochs, validation_data=validation_ds)
```


### Relevant log output

```shell
2025-01-09 19:39:57.953074: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] 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:1736469597.967544   14431 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1736469597.971752   14431 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-01-09 19:39:57.986195: 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 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
1 Physical GPUs, 1 Logical GPUs
I0000 00:00:1736469600.052169   14431 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 10240 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 4070, pci bus id: 0000:0a:00.0, compute capability: 8.9
Fitting the top layer of the model
Epoch 1/2
2025-01-09 19:40:04.479339: I tensorflow/core/kernels/data/tf_record_dataset_op.cc:376] The default buffer size is 262144, which is overridden by the user specified `buffer_size` of 8388608
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1736469604.583055   14486 service.cc:148] XLA service 0x7f8df8002230 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1736469604.583109   14486 service.cc:156]   StreamExecutor device (0): NVIDIA GeForce RTX 4070, Compute Capability 8.9
2025-01-09 19:40:04.722034: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.
I0000 00:00:1736469605.234339   14486 cuda_dnn.cc:529] Loaded cuDNN version 90300
  7/582 ━━━━━━━━━━━━━━━━━━━━ 12s 22ms/step - binary_accuracy: 0.5658 - loss: 0.6950 
I0000 00:00:1736469608.599510   14486 device_compiler.h:188] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
243/582 ━━━━━━━━━━━━━━━━━━━━ 16s 48ms/step - binary_accuracy: 0.8715 - loss: 0.2755
2025-01-09 19:40:20.475756: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 1073741824 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469620.475821   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 1073741824
2025-01-09 19:40:20.615636: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 966367744 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469620.615700   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 966367744
2025-01-09 19:40:20.769033: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 869731072 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469620.769095   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 869731072
2025-01-09 19:40:20.906909: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 782758144 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469620.906973   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 782758144
2025-01-09 19:40:21.048863: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 704482304 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469621.048940   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 704482304
2025-01-09 19:40:21.229614: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 634034176 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469621.229682   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 634034176
2025-01-09 19:40:21.371940: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 570630912 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469621.372000   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 570630912
2025-01-09 19:40:21.510751: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 513568000 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469621.510817   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 513568000
2025-01-09 19:40:21.650945: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 462211328 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469621.651034   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 462211328
2025-01-09 19:40:21.814945: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 415990272 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469621.815035   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 415990272
2025-01-09 19:40:21.954790: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 374391296 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469621.954851   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 374391296
2025-01-09 19:40:22.094150: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 336952320 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469622.094219   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 336952320
2025-01-09 19:40:22.267664: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 303257088 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
...
2025-01-09 19:40:23.128022: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 161164032 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469623.128090   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 161164032
2025-01-09 19:40:23.296856: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:1253] failed to alloc 145047808 bytes on host: RESOURCE_EXHAUSTED: : CUDA_ERROR_OUT_OF_MEMORY: out of memory
W0000 00:00:1736469623.296940   14537 device_host_allocator.h:61] could not allocate pinned host memory of size: 145047808
Output is truncated. View as a scrollable element or open in a text editor. Adjust cell output settings...
```