在多进程训练中,`use_stateful_dataloader` 会恢复 Cursor,而不是随机排序 — 采样器排序永远不会序列化

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,