Skip automatic DataLoader sharding when the loader is already rank-sharded

Summary

Accelerator.prepare(dataloader) always assumes the incoming loader is a global view of the dataset. In a distributed run it then wraps the sampler/dataset again (BatchSamplerShard, IterableDatasetShard, or datasets.IterableDataset.shard()). If the user already partitioned by rank, that second cut is silent data loss: each process keeps 1 / num_processes of the shard it was given, the loss curve still looks normal, and there is no warning.

This is the same request as #4075 (closed stale) and the same idea as #4087 (closed stale, no review). #4144 documents the iterable path only. The code path is unchanged.

What happens today

In prepare_data_loader, when num_processes != 1 and dispatch_batches is false:

1. Map-style. The batch_sampler is wrapped in BatchSamplerShard. A rank-aware DistributedSampler (or any custom per-rank sampler) is sharded a second time.
2. Torch IterableDataset. The dataset is wrapped in IterableDatasetShard, which iterates the entire underlying stream on every rank and keeps 1 / N. A dataset that already yields only the local shard is cut again. Every rank also rereads the full source (num_processes× read amplification on remote/streaming data).
3. HF datasets.IterableDataset. If n_shards >= num_processes, Accelerate calls .shard(...) itself. A dataset the user already sharded is sharded again.

The usual workaround is to skip prepare on the dataloader. That also drops device placement, set_epoch forwarding, and dataloader state tracking.

Proposal

Add an explicit opt-in on DataLoaderConfiguration:

accelerator = Accelerator(
    dataloader_config=DataLoaderConfiguration(already_sharded=True)
)
loader = accelerator.prepare(loader)

When already_sharded=True:

• do not wrap with BatchSamplerShard / IterableDatasetShard
• do not call datasets.IterableDataset.shard(...)
• still return DataLoaderShard so device placement, set_epoch, and state tracking keep working
• the user is responsible for equal step counts across ranks

Reject combinations that conflict with a per-rank loader:

• dispatch_batches=True (main process iterates and broadcasts)
• split_batches=True (would split the already-local batch)

Default remains already_sharded=False. Existing scripts are unchanged.

dispatch_batches=True is not a substitute: it changes the I/O model so one process reads and the others wait. That is the opposite of a pre-sharded streaming or memmap pipeline.

Why not only docs

#4144 is useful, but an opt-in flag is what stops the silent failure. The failure mode is not an exception; it is a quieter training run on a subset of the data.

Related

• #4075, #4087 (stale)
• #3520, #4062, #3547, #3124
• huggingface/datasets#6594