huggingface / huggingface/accelerate
`use_stateful_dataloader` + `num_workers>0` draws epoch-0 permutation during `prepare()` before cross-rank RNG sync — corrupts DDP data partition
- Dominant language
- Python
- Stars
- 9.9k
- Forks
- 1.5k
- Avg merge
- 5d 2h
- Merged PRs (30d)
- 27
Description
### System Info
```Shell
- `accelerate` version: 1.13.0 (and main branch)
- `torch` version: 2.13.0
- `torchdata` version: 0.11.0
- Platform: Linux
- Python version: 3.11
- Distributed environment: 2-process CPU/Gloo (reproducible without GPU)
```
### Information
- [ ] The official example scripts
- [x] My own modified scripts
### Tasks
- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [x] My own task or dataset (give details below)
### Reproduction
#### Summary & Root Cause
With `DataLoaderConfiguration(use_stateful_dataloader=True)`, `accelerator.prepare()` wraps the dataloader as a torchdata `StatefulDataLoader` and immediately invokes `base_dataloader.state_dict()` inside `DataLoaderAdapter.__init__` to record the initial dataloader state.
When `num_workers > 0`:
1. `StatefulDataLoader.state_dict()` calls `self._get_iterator()` when `self._iterator is None` to snapshot iterator state.
2. The multiprocessing iterator's `_reset()` spawns worker processes and primes the prefetch queue (`prefetch_factor * num_workers` items), which iterates `_sampler_iter`.
3. For a vanilla `RandomSampler(generator=None)`, this immediately draws the epoch-0 permutation seed from the process's **local global RNG at `prepare()` time**.
4. Cross-rank dataloader RNG synchronization (`DataLoaderShard.__iter__` -> `synchronize_rng_states`) only happens when iteration actually starts, so any divergence in per-rank RNG before `prepare()` (e.g. no global `torch.manual_seed` set, deliberate per-rank seeding, or unequal RNG consumption across ranks during model/dataset setup) causes ranks to derive **different** epoch-0 permutations.
5. `BatchSamplerShard` then rank-interleaves divergent permutations, which **corrupts the data partition during epoch 0** (samples are duplicated across ranks and others are dropped).
*(Note: This affects uninterrupted training runs during epoch 0, independent of any checkpoint/resume operation).*
#### Minimal Reproducible Example
Run with 2 processes on CPU (no GPU needed):
```bash
# Case A (diverged/unseeded RNG across ranks): partition is corrupted
python -m torch.distributed.run --standalone --nproc_per_node=2 repro_prepare_prefetch.py noseed
# Case B (identical seed explicitly set before prepare): partition is intact
python -m torch.distributed.run --standalone --nproc_per_node=2 repro_prepare_prefetch.py seed
```
`repro_prepare_prefetch.py`:
```python
import sys
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader, TensorDataset
from accelerate import Accelerator, DataLoaderConfiguration
MODE = sys.argv[1] # "seed" or "noseed"
N, BS = 96, 4
if MODE == "seed":
torch.manual_seed(1234) # identical on both ranks before prepare()
acc = Accelerator(cpu=True, dataloader_config=DataLoaderConfiguration(use_stateful_dataloader=True))
loader = acc.prepare(
DataLoader(TensorDataset(torch.arange(N)), batch_size=BS, shuffle=True, num_workers=6)
)
# Iterate through epoch 0
flat = [i for b in loader for i in b[0].tolist()]
gathered = [None] * acc.num_processes
dist.all_gather_object(gathered, flat)
acc.wait_for_everyone()
if acc.is_main_process:
union = [i for rank_items in gathered for i in rank_items]
union_set = set(union)
missing = len(set(range(N)) - union_set)
duplicates = len(union) - len(union_set)
print(
f"mode={MODE}: epoch-0 union={len(union_set)}/{N}, missing={missing}, duplicates={duplicates}"
)
```
#### Observed Output
```
mode=noseed: epoch-0 union=72/96, missing=24, duplicates=24
mode=seed: epoch-0 union=96/96, missing=0, duplicates=0
```
With `num_workers=6` and no prior shared seed synchronization, 24 out of 96 samples are duplicated across ranks and 24 samples are completely omitted from epoch 0.
### Expected behavior
`accelerator.prepare(dataloader)` should not draw from the sampler or consume generator RNG before iteration begins:
1. **Lazy state capture (Preferred)**: `DataLoaderAdapter.__init__` should avoid materializing `_get_iterator()` during initialization, capturing state lazily on the first step/checkpoint instead.
2. **Post-init cleanup / reset**: If `state_dict()` must be called during `prepare_data_loader`, the early iterator should be cleared (`dataloader.base_dataloader._iterator = None` or equivalent) to reap worker pools and defer the sampler draw to iteration time.
3. At minimum, `synchronize_rng_states` should run prior to adapter initial state capture if `use_stateful_dataloader=True` and `num_workers > 0`.
Contributor guide
Research direction
Trace DataLoaderAdapter.__init__ and prepare_data_loader to confirm when state_dict() materializes the iterator, then follow DataLoaderShard.__iter__ and synchronize_rng_states. Run the provided two-process repro with seed and noseed modes. Done means prepare() no longer consumes divergent sampler RNG before iteration and epoch 0 has no missing or duplicated samples.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100