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

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