TensorFlow
tf.data.Dataset
data parallelism
distributed computing
Horovod

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.

python
1import horovod.tensorflow.keras as hvd
2
3hvd.init()
4print("world size:", hvd.size())
5print("rank:", hvd.rank())
6print("local rank:", hvd.local_rank())

You usually do this before building the dataset and before configuring GPU visibility.

Shard the Dataset by Rank

The core distribution step is:

python
dataset = dataset.shard(num_shards=hvd.size(), index=hvd.rank())

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:

python
1import tensorflow as tf
2import horovod.tensorflow.keras as hvd
3
4hvd.init()
5
6features = tf.random.normal((1000, 10))
7labels = tf.random.uniform((1000,), maxval=2, dtype=tf.int32)
8
9dataset = tf.data.Dataset.from_tensor_slices((features, labels))
10dataset = dataset.shuffle(1000)
11dataset = dataset.shard(hvd.size(), hvd.rank())
12dataset = dataset.batch(32)
13dataset = dataset.prefetch(tf.data.AUTOTUNE)
14
15for x_batch, y_batch in dataset.take(1):
16    print(x_batch.shape, y_batch.shape)

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:

python
1dataset = tf.data.TFRecordDataset(files)
2dataset = dataset.shard(hvd.size(), hvd.rank())
3dataset = dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
4dataset = dataset.batch(64)

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:

python
1files = tf.data.Dataset.list_files("/data/train-*.tfrecord", shuffle=True)
2files = files.shard(hvd.size(), hvd.rank())
3dataset = files.interleave(
4    tf.data.TFRecordDataset,
5    cycle_length=tf.data.AUTOTUNE,
6    num_parallel_calls=tf.data.AUTOTUNE
7)

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.

Course illustration
Course illustration

All Rights Reserved.