TensorFlow
Dataset API
Multi-GPU
Data Pipeline
Machine Learning

How does one move data to multiple GPU towers using Tensorflow's Dataset API

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

With multi-GPU training, you usually do not manually copy individual examples to each GPU tower yourself. In modern TensorFlow, the normal pattern is to build one tf.data.Dataset, batch it with the global batch size, and let a distribution strategy such as MirroredStrategy shard each batch across the GPUs.

Think in terms of global batch and per-replica batch

If you have 2 GPUs and want each GPU to process 32 examples at a time, the dataset should typically produce a global batch size of 64. TensorFlow then splits that batch into per-replica chunks under the distribution strategy.

A simple input pipeline looks like this:

python
1import tensorflow as tf
2
3GLOBAL_BATCH_SIZE = 64
4
5dataset = tf.data.Dataset.from_tensor_slices((
6    tf.random.normal([1024, 32]),
7    tf.random.uniform([1024], maxval=10, dtype=tf.int32)
8))
9dataset = dataset.shuffle(1024).batch(GLOBAL_BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

The dataset itself stays conceptually device-agnostic. You do not annotate one item for GPU 0 and another for GPU 1. That splitting is handled later.

Let MirroredStrategy distribute the dataset

In TensorFlow 2, the standard approach is:

python
1strategy = tf.distribute.MirroredStrategy()
2
3with strategy.scope():
4    model = tf.keras.Sequential([
5        tf.keras.layers.Dense(64, activation="relu"),
6        tf.keras.layers.Dense(10)
7    ])
8    model.compile(
9        optimizer="adam",
10        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
11        metrics=["accuracy"]
12    )
13
14dist_dataset = strategy.experimental_distribute_dataset(dataset)
15model.fit(dist_dataset, epochs=2)

experimental_distribute_dataset prepares the dataset for distributed consumption. Each replica gets its slice of each global batch. This is the modern replacement for a lot of older manual "tower" plumbing.

In older tower-style code, the same principle still applied

The phrase "GPU towers" usually comes from TensorFlow 1 multi-GPU code, where people manually created one model tower per GPU. Even there, the right mental model was still "produce batches centrally, then split them per tower," not "have every tower run its own unrelated dataset reader."

If you manually move tensors to each GPU, you are often rebuilding logic that the distribution system already knows how to do more efficiently and with less bug risk.

Place expensive preprocessing in the dataset pipeline

What you should focus on is making the dataset fast enough to keep all GPUs busy:

  • parallel file reads
  • vectorized parsing
  • batching before heavy work when possible
  • 'prefetch(tf.data.AUTOTUNE)'

If GPUs are underutilized, the problem is usually the input pipeline or batch sizing, not that you forgot to issue manual device copies to each tower.

Sharding data across workers is different from sharding across GPUs

Inside one machine, MirroredStrategy handles replica splitting. Across multiple workers, strategies such as MultiWorkerMirroredStrategy add worker-level sharding concerns too.

That distinction matters because:

  • per-GPU splitting happens inside the worker
  • worker sharding determines which worker sees which dataset partition

It is easy to mix those two levels together when reading older distributed-training examples.

Use the right batching rule

A frequent source of confusion is batch size. If you want each of 4 GPUs to process 16 examples, the dataset batch size should usually be 64, not 16. The smaller number is the per-replica batch. The larger number is the global batch that TensorFlow distributes.

Getting that wrong can make the model look like it is running on multiple GPUs while actually underfeeding them.

Common Pitfalls

  • Trying to manually place each dataset element on a specific GPU tower.
  • Batching by per-replica size instead of global batch size.
  • Forgetting to prefetch, which can leave GPUs idle.
  • Confusing multi-GPU splitting on one machine with multi-worker dataset sharding.
  • Reusing old TensorFlow 1 tower examples without translating them into TensorFlow 2 distribution concepts.

Summary

  • In modern TensorFlow, you usually do not manually move data to each GPU tower.
  • Build one tf.data.Dataset, batch it with the global batch size, and distribute it with MirroredStrategy.
  • TensorFlow splits each batch into per-replica chunks for the GPUs.
  • Input-pipeline efficiency matters more than manual device placement in most cases.
  • Keep global batch size and per-replica batch size clearly separated.

Course
Beginner
27 lessons
10 hours
System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the course
ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design