#4195·accelerate

`use_stateful_dataloader` under multi-process training restores the cursor, not the shuffle order — sampler permutation is never serialized

Author: GoldenStainCreated Aug 30, 2026Updated Sep 9, 2026

System Info

bash
- accelerate `1.13.0`, torch `2.13.0`, torchdata `0.11.0`, Linux
- Reproduced on multi-GPU DDP and on 2-process CPU/Gloo (reproducer below needs no 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

bash
python -m torch.distributed.run --standalone --nproc_per_node=2 issue_repro.py save   /tmp/sfd_ckpt
python -m torch.distributed.run --standalone --nproc_per_node=2 issue_repro.py resume /tmp/sfd_ckpt
# on a GPU machine, prefix with: CUDA_VISIBLE_DEVICES=""
python
# issue_repro.py — DDP + use_stateful_dataloader + shuffle=True:
# resume restores the right POSITION but the wrong ORDER.
import os
import sys

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator, DataLoaderConfiguration

MODE, OUT = sys.argv[1], sys.argv[2]
RANK = int(os.environ["RANK"])
SEED, N, BS, SKIP = 1234, 64, 4, 3  # 16 batches/epoch; checkpoint after 3 batches of epoch 1

torch.manual_seed(SEED)  # identical in both phases -> reconstruction is deterministic
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=0))


def as_list(batch):
    return batch[0].tolist()  # TensorDataset yields a 1-tuple


if MODE == "save":
    os.makedirs(OUT, exist_ok=True)
    epoch0 = [as_list(b) for b in loader]          # consume full epoch 0
    it = iter(loader)                              # epoch 1
    consumed = [as_list(next(it)) for _ in range(SKIP)]
    state = loader.state_dict()                    # per-rank save, as in #3080/#4165
    torch.save(state, f"{OUT}/dl_rank{RANK}.bin")
    torch.save(
        {"epoch0": epoch0, "consumed": consumed, "cont": [as_list(b) for b in it]},
        f"{OUT}/ref_rank{RANK}.pt",
    )
    print(f"[rank {RANK}] saved dl state keys: {list(state.keys())}")
else:
    ref = torch.load(f"{OUT}/ref_rank{RANK}.pt", weights_only=False)
    loader.load_state_dict(torch.load(f"{OUT}/dl_rank{RANK}.bin", weights_only=False))
    resumed = [as_list(b) for b in loader]
    print(f"[rank {RANK}] batches consumed before checkpoint: {ref['consumed']}")
    print(f"[rank {RANK}] true epoch-1 remainder            : {ref['cont']}")
    print(f"[rank {RANK}] resumed stream                     : {resumed}")
    print(f"[rank {RANK}] epoch-0 remainder (for reference)  : {ref['epoch0'][SKIP:]}")
    assert resumed == ref["epoch0"][SKIP:], "not even the deterministic epoch-0 replay"
    assert resumed != ref["cont"], "resumed matched the true continuation -- bug not reproduced"
    print(f"[rank {RANK}] POSITION restored, ORDER lost: epoch-0's order at epoch-1's position")

Expected behavior

Either:

  • exact resume under multi-process as well (preferred — matching torchdata's native semantics: the restored stream equals the recorded continuation), or
  • an explicit warning at save/load time that only the cursor is being restored when num_processes > 1 and the sampler chain is shuffled.

Observed output (rank 0; rank 1 analogous):

# save
[rank 0] saved dl state keys: ['_index_sampler_state', '_sampler_iter_state', '_sampler_iter_yielded', '_num_yielded', '_IterableDataset_len_called', '_shared_seed', 'fetcher_state', 'dataset_state', '_iterator_finished']

# resume
[rank 0] batches consumed before checkpoint: [[39, 18, 10, 13], [63, 27, 2, 62], [25, 6, 37, 50]]
[rank 0] true epoch-1 remainder            : [[55, 16, 59, 0], [22, 46, 17, 57], [23, 45, 19, 51], [21, 26, 40, 36], [20, 5, 41, 32]]
[rank 0] resumed stream                     : [[15, 41, 52, 61], [5, 50, 18, 14], [31, 53, 36, 47], [35, 9, 23, 58], [22, 7, 51, 55]]
[rank 0] epoch-0 remainder (for reference)  : [[15, 41, 52, 61], [5, 50, 18, 14], [31, 53, 36, 47], [35, 9, 23, 58], [22, 7, 51, 55]]
[rank 0] POSITION restored, ORDER lost: epoch-0's order at epoch-1's position