tensorflow: Prefetch to GPU: prefetch_to_device does not do anything (from_generator), tf.data.Dataset API unclear

prefetch_to_device works when the data produced by the input pipeline is allocated in pinned host memory. map().batch() is currently the option to achieve this through tf.data. This is not documented anywhere and is a subtle implementation detail: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc#L517

I plan to investigate whether there is any downside to making all tf.data allocations use pinned host memory and if not, make this available more broadly. Note that this might not address the issue for input pipelines that do not allocate memory through tf.data. For instance, the example in the original description of this issue performs the memory allocation outside of tf.data.

TLDR: prefetch_to_device works as expected, when the data produced by the input pipeline is allocated in pinned host memory.

I have recently investigated this issue.

First, to compare copy_to_device and prefetch_to_device performance, I have modified the original repro as follows:

def main(argv):
  del argv

  @tf.function
  def do_stuff(wmat, tf_var):
    result = tf.matmul(wmat + 1, tf.transpose(wmat))
    error = tf.reduce_mean(tf_var)
    return error, result

  n_batches = 100
  weights = [None] * n_batches
  for i in range(n_batches):
    weights[i] = tf.constant(np.random.rand(2000, 10000), dtype=tf.float32)

  def gen():
    for weight in weights:
      yield weight

  option_names = [
      'prefetch_to_device', 'copy_to_device', 'copy_to_device_and_prefetch'
  ]
  dataset = tf.data.Dataset.from_generator(gen, output_types=(tf.float32))

  if FLAGS.option == 0:
    dataset = dataset.apply(tf.data.experimental.copy_to_device('/gpu:0'))

  if FLAGS.option == 1:
    dataset = dataset.apply(tf.data.experimental.prefetch_to_device('/gpu:0'))

  tf_var = tf.Variable(np.zeros(3))
  adam = tf.keras.optimizers.Adam(1e-4)

  for _, wmat in tqdm(enumerate(dataset)):
    with tf.GradientTape() as tape:
      result = do_stuff(wmat, tf_var)
      grads = tape.gradient(result[0], [tf_var])
      adam.apply_gradients(zip(grads, [tf_var]))

Using this program I have collected the following traces:

copy_to_device

prefetch_to_device

While a single step of the prefetch_to_device version of the program is slightly faster (based on the trace view measurement as opposed to TF profiler overview, which might not be correctly detecting “steps” for this program) – 54.9ms vs 62.6ms – the trace also shows that the data copy seems to be delayed by something.

Next, I investigated whether making sure that the input pipeline produces data that is allocated in pinned host memory, resolves the issue.

Currently, tf.data only allocates pinned host memory for input pipelines that use the map().batch() pattern. Therefore, I needed to modify the program as follows:

def main(argv):
  del argv

  @tf.function
  def do_stuff(wmat, tf_var):
    result = tf.matmul(wmat + 1, tf.transpose(wmat))
    error = tf.reduce_mean(tf_var)
    return error, result

  option_names = [
      'copy_to_device', 'prefetch_to_device', 'copy_to_device_and_prefetch'
  ]

  dataset = tf.data.Dataset.range(1024 * 100)
  dataset = dataset.map(
      lambda _: tf.random.uniform([10000]),
      num_parallel_calls=tf.data.AUTOTUNE)
  dataset = dataset.batch(batch_size=1024)

  if FLAGS.option == 0:
    dataset = dataset.apply(tf.data.experimental.copy_to_device('/gpu:0'))

  if FLAGS.option == 1:
    dataset = dataset.apply(tf.data.experimental.prefetch_to_device('/gpu:0'))

  tf_var = tf.Variable(np.zeros(3))
  adam = tf.keras.optimizers.Adam(1e-4)

  for _, wmat in tqdm(enumerate(dataset)):
    with tf.GradientTape() as tape:
      result = do_stuff(wmat, tf_var)
      grads = tape.gradient(result[0], [tf_var])
      adam.apply_gradients(zip(grads, [tf_var]))

These changes end up resolving the issue:

copy_to_device

prefetch_to_device

@jsimsa

TLDR: prefetch_to_device works as expected, when the data produced by the input pipeline is allocated in pinned host memory.

So the conclusion is: prefetch_to_device actually works, but it needs be to used with map().batch() to get better performance?

Currently, tf.data only allocates pinned host memory for input pipelines that use the map().batch() pattern.

And where can I get such information? From TensorFlow source code?

@aaudiber Thank you for taking the time to review this. It’s true that that was a mistake in this example. I can confirm that option 3 (prefetch_to_device) is as fast as option 2 (copy to device + prefetch).

However, I do still see a noticeable time that the GPU is performing nothing else than MemcpyH2D. I am not sure if I’m reading this correctly, but why would that be the case?

@jsimsa i see that you were able to get the tracer to show prefetching (memcpy overlap with gpu compute)

but did you see an actual improvement in training speed?

ds = tf.data.Dataset.range(2**13)
ds = ds.map(
  lambda _: np.random.randint(0, 10, (4096)).astype(np.float32),
  num_parallel_calls=tf.data.AUTOTUNE)

ds = ds.batch(batch_size=4096).cache().repeat(-1)
pf_ds = ds.prefetch(16)
gpu_pf_ds = ds.apply(tf.data.experimental.prefetch_to_device('/gpu:0', buffer_size=16))


def time_ds(ds, ramp_steps=100, steps=200):
  @tf.function
  def _step(wmat, tf_var):
    with tf.device('/gpu:0'):
      result = tf.matmul(wmat + 1, tf.transpose(wmat))
      error = tf.reduce_mean(tf_var)
      return error, result

  tf_var = tf.Variable(np.zeros(3))
  adam = tf.keras.optimizers.Adam(1e-4)

  assert steps > ramp_steps
  for i, elem in enumerate(ds):
    if i == ramp_steps:
      t = time.time()
    if i == steps:
      print(time.time() - t)
      break
    with tf.GradientTape() as tape:
      result = _step(elem, tf_var)
      grads = tape.gradient(result[0], [tf_var])
      adam.apply_gradients(zip(grads, [tf_var]))

time_ds(ds) time_ds(pf_ds) time_ds(gpu_pf_ds)

all 3 show roughly 8 seconds per 100 steps on colab GPU

However, I do still see a noticeable time that the GPU is performing nothing else than MemcpyH2D. I am not sure if I’m reading this correctly, but why would that be the case?

Same to me. After using dataset = dataset.apply(tf.data.experimental.prefetch_to_device(device='/gpu:0', buffer_size=16)), it is faster than dataset = dataset.prefetch(buffer_size=16).

But there is still noticeable time spent on cudaMemcpyH2D. Please see here for more details.

Maybe prefetch_to_device does not really prefetch?