Distribute data from tf.data.Dataset to multiple workers e.g. for Horovod
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
Introduction
When training with Horovod, every worker runs the same training code, so the input pipeline must avoid feeding the exact same examples to every worker in the same step. The standard solution is dataset sharding: each worker reads only its own slice of the dataset. In tf.data, that is usually a one-line change with dataset.shard(hvd.size(), hvd.rank()).
Initialize Horovod First
Horovod needs to know how many workers exist and which worker the current process represents.
You usually do this before building the dataset and before configuring GPU visibility.
Shard the Dataset by Rank
The core distribution step is:
That means:
- total number of shards equals total number of workers
- each worker reads only the shard that matches its rank
Without sharding, every worker would often iterate over the full dataset, causing duplicated work and distorted effective batch semantics.
A Minimal Working Example
Here is a simple end-to-end pattern:
Each worker now sees only its own subset of examples.
Shard Early in the Pipeline
A good rule is to shard before expensive per-example preprocessing whenever possible. That way each worker only preprocesses the data it will actually consume.
For example:
If you parse everything first and shard later, every worker may waste CPU doing work for records it will never train on.
File-Level Versus Example-Level Sharding
There are two common strategies:
- shard files across workers
- shard dataset elements across workers
If you have many large input files, sharding the file list first is often efficient:
If you start from one logical dataset object, example-level dataset.shard(...) is often the simplest solution.
Adjust Batch Size Correctly
In distributed training, decide whether your configured batch size is:
- per worker
- global across all workers
Most Horovod examples treat the dataset batch size as per worker. If each worker uses batch size 32 and you have 4 workers, the global batch size is effectively 128.
That affects:
- learning-rate scaling
- steps per epoch
- throughput expectations
Be explicit about which interpretation your training loop uses.
Shuffle Carefully
Shuffling and sharding interact. A common pattern is to shuffle the dataset and then shard it, but depending on the data source and reproducibility requirements, you may instead shard file lists first and then shuffle within each worker’s subset.
The key goal is to avoid every worker processing the exact same element sequence.
Common Pitfalls
The biggest mistake is forgetting to shard the dataset at all, which causes each worker to process the entire dataset redundantly. Another is batching before thinking through whether the size is per worker or global. Developers also place expensive parsing before sharding, wasting preprocessing work on examples that a worker will discard later. Shuffling can be another source of confusion if every worker ends up seeing the same order. The simplest reliable starting point is dataset.shard(hvd.size(), hvd.rank()) early in the input pipeline.
Summary
- Horovod workers should usually read different dataset shards, not the full dataset independently.
- Use
dataset.shard(hvd.size(), hvd.rank())to split input by worker rank. - Initialize Horovod before constructing the distributed input pipeline.
- Shard early to avoid wasted preprocessing.
- Decide clearly whether batch size is per worker or global.
- Verify shuffling and file reading so workers do not see identical data sequences.

