Tensorflow
Keras
Auto-sharding
Dataset
Machine Learning

Tensorflow - Keras Consider either turning off auto-sharding or switching the auto_shard_policy to DATA to shard this dataset

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

This TensorFlow warning usually appears during distributed training when TensorFlow tries to shard the input dataset automatically and cannot do it the way it wants. In practice, the dataset pipeline does not support the current auto-sharding strategy, often because the pipeline is not file-based in a way TensorFlow can split per worker. The fix is usually to change the auto-shard policy to DATA or turn it off deliberately.

What Auto-Sharding Is Doing

In multi-worker or distributed training, TensorFlow tries to prevent every worker from reading the exact same dataset. Auto-sharding splits the dataset so different workers process different data.

Common policies include:

  • 'AUTO'
  • 'FILE'
  • 'DATA'
  • 'OFF'

The warning means the default behavior could not shard your current dataset cleanly using the preferred strategy.

Why the Warning Happens

A common trigger is a dataset pipeline that does not originate from a clean file-based source that TensorFlow can partition worker by worker.

For example, pipelines built from:

  • 'from_tensor_slices'
  • Python generators
  • custom mapped datasets
  • already batched or transformed sources

can make FILE-style sharding impossible or ineffective.

TensorFlow then suggests either:

  • disable auto-sharding
  • or use DATA sharding instead

Switch to DATA Sharding

If your dataset can be split at the element level, DATA is often the right fix.

python
1import tensorflow as tf
2
3options = tf.data.Options()
4options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
5
6dataset = dataset.with_options(options)

This tells TensorFlow to shard by dataset elements rather than by input files.

For many in-memory or non-file pipelines, this is the most practical choice.

Turn Auto-Sharding Off When Appropriate

If every worker intentionally needs the full dataset, or if sharding is genuinely inappropriate for the workflow, you can disable it.

python
1import tensorflow as tf
2
3options = tf.data.Options()
4options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
5
6dataset = dataset.with_options(options)

Do this only when you understand the consequence: workers may read duplicate data, which can waste work or distort training depending on the strategy.

Example With Keras fit

python
1import tensorflow as tf
2
3x = tf.random.uniform((1000, 10))
4y = tf.random.uniform((1000,), maxval=2, dtype=tf.int32)
5
6dataset = tf.data.Dataset.from_tensor_slices((x, y)).batch(32)
7
8options = tf.data.Options()
9options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
10dataset = dataset.with_options(options)
11
12model = tf.keras.Sequential([
13    tf.keras.layers.Dense(16, activation='relu'),
14    tf.keras.layers.Dense(1, activation='sigmoid')
15])
16
17model.compile(optimizer='adam', loss='binary_crossentropy')
18model.fit(dataset, epochs=1)

This kind of pipeline often benefits from DATA sharding because there are no source files to shard per worker.

Choose the Policy Based on the Dataset Source

A useful rule is:

  • use FILE when workers can safely split input files
  • use DATA when the dataset is better split by elements
  • use OFF only when duplication is acceptable or sharding is genuinely wrong

The best policy depends on how the data enters the pipeline, not just on Keras itself.

Do Not Ignore the Distributed Context

If you are not actually training in a distributed multi-worker setting, the warning may appear because of the runtime environment or strategy configuration. In that case, the first debugging question is whether auto-sharding matters at all in the current run.

If you are using a distribution strategy, then the sharding choice becomes part of correctness and efficiency, not just noise suppression.

Common Pitfalls

A common mistake is turning sharding off just to silence the warning without understanding that all workers may now see the same data.

Another issue is assuming FILE sharding should always work. It only makes sense when the dataset really comes from a shardable file source.

Developers also sometimes apply the option to the wrong dataset object. The configured options must be attached to the dataset actually passed into training.

Finally, do not treat this as a Keras-only issue. The underlying behavior comes from tf.data and distributed input handling.

Summary

  • The warning means TensorFlow could not auto-shard the dataset the way it expected.
  • 'DATA sharding is often the right fix for non-file or in-memory dataset pipelines.'
  • 'OFF disables sharding entirely and should be used only deliberately.'
  • Pick the policy based on how the dataset is built, not by guesswork.
  • Apply the options to the actual dataset object that Keras receives.

Course illustration
Course illustration

All Rights Reserved.