huggingface / huggingface/accelerate

`use_stateful_dataloader` + `num_workers>0` draws epoch-0 permutation during `prepare()` before cross-rank RNG sync — corrupts DDP data partition

Open
#4,204 3 comments 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

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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.