tf.data.experimental.prefetch_to_device has no effect inside tf.distribute.Strategy.distribute_datasets_from_function.

Issue type

Performance

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

tf 2.19.0

Custom code

No

OS platform and distribution

RHEL 9.4

Python version

3.11

CUDA/cuDNN version

12.5

Current behavior?

MemcpyH2D does not overlap with model computation when using tf.data.experimental.prefetch_to_device inside tf.distribute.MirroredStrategy.distribute_datasets_from_function. I would expect this operations to overlap.

Standalone code to reproduce the issue

import tensorflow as tf

class Model(tf.keras.Model):
    def call(self, x):
        y = x / 1000
        for i in range(3):
            y = tf.matmul(y, x / 1000)
        return tf.reduce_sum(y, axis=[1, 2])

def get_dataset(ictx):
    ds = tf.data.Dataset.range(1, 1001, output_type=tf.float32)
    ds = ds.map(lambda i: (tf.ones((1024 * 5, 1024 * 5)) / i, 0.0))
    ds = ds.batch(8)
    ds = ds.apply(tf.data.experimental.prefetch_to_device('gpu'))
    return ds

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    ds = strategy.distribute_datasets_from_function(get_dataset)
    model = Model()
    model.compile(loss='mse')
    model.fit(
        ds,
        epochs=1,
        steps_per_epoch=30,
        callbacks=tf.keras.callbacks.TensorBoard(profile_batch=(15, 25)))

gist