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

System Info

- `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
• 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)
• 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):

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

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.