huggingface / huggingface/accelerate
Skip automatic DataLoader sharding when the loader is already rank-sharded
- 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
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