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.
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:
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:
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 withMirroredStrategy. - 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.
Related reading
- How does one train multiple models in a single script in TensorFlow when there are GPUs present?
- How does one train multiple models in a single script in TensorFlow when there are GPUs present?
- How does one transfer CUDA constant memory in tensorflow's C API
- How does one use Pytorch cuda with an A100 GPU?
- How does reduce_sum work in tensorflow?
- How does tensorflow batch_matmul work?
- How does one set the pad token correctly not to eos during fine-tuning to avoid model not predicting EOS?
- How does pytorch backprop through argmax?

System Design Fundamentals
Build a strong foundation in designing scalable, reliable distributed systems.
View the courseML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.