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

Open
  • Is something wrong?
#4,204 5 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

@JohnnyWilson16 is already working on this.

Since Oct 19, 2026.

  • #4225 by @FinalSunFlower — closed without merging
  • #4234 by @jayzuccarelli — closed without merging

Assessment

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

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.

Description

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.
Dominant language
Python
Stars
9.9k
Forks
1.6k
Avg merge
20h 23m
Merged PRs (30d)
12

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from huggingface/accelerate

All issues in huggingface/accelerate

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.