huggingface / huggingface/accelerate

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

Open
#4,199 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
9.9k
Forks
1.5k
Avg merge
5d 2h
Merged PRs (30d)
27

Description

## 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`:

```python
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

Contributor guide

Open the contributing guide

Research direction

Start at prepare_data_loader and DataLoaderConfiguration, then trace the BatchSamplerShard, IterableDatasetShard, and datasets.IterableDataset.shard paths. Verify that an already_sharded option preserves DataLoaderShard behavior while skipping secondary sharding and rejects conflicting dispatch_batches or split_batches settings; add or run focused tests for these paths.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data, distributed-systems
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.