Fix stateful dataloader losing shuffle order on multi-process resume (#4195)

[massimiliano1991]

What this fixes

Closes #4195 — under multi-process training, a shuffled stateful dataloader resumes on epoch 0's permutation at the restored cursor position instead of continuing the current epoch's order.

Root cause

In the sharded stateful path, the shuffle order for a given __iter__ is seeded by DataLoaderShard/DataLoaderDispatcher.iteration (the epoch counter), which set_epoch() propagates down to the sampler (SeedableRandomSampler.epoch, or the sampler generator). That counter is never part of the checkpoint: DataLoaderStateMixin.state_dict() returns only the underlying torchdata StatefulDataLoader state, and load_state_dict() forwards it verbatim.

So on resume a fresh process starts at iteration == 0, __iter__ calls set_epoch(0), and torchdata's fast-forward (itertools.islice(iter(index_sampler), sampler_iter_yielded, None)) re-iterates the sampler from epoch 0 — restoring the batch position but the epoch‑0 order. (Single-process is unaffected: there the sampler iterator state torchdata already serializes carries the permutation.)

The fix

• Persist the epoch counter as _iteration in the dataloader state_dict, and restore it in load_state_dict() (via set_epoch) before the base dataloader rebuilds its sampler iterator. With use_seedable_sampler=True the shuffle order then resumes exactly.
• The plain-RandomSampler multi-process path has no recoverable per-epoch seed across a process restart (its generator seed is drawn fresh each launch), so exact shuffle-resume isn't achievable there — emit a one-time warning pointing at use_seedable_sampler=True. This is the issue's option 2.

The adapter state_dict now carries one extra key beyond torchdata's StatefulDataLoader; loading a state without _iteration (old checkpoints, or a torchdata dict) is a no-op, and the two *_equivalent_to_torchdata_stateful_dataloader tests drop that key before comparing the shared state.

Verification

Against main @ f13f7c13, the issue's reproducer (torchrun --nproc_per_node=2, use_seedable_sampler=True):

	before	after
resumed == true continuation	❌ (both ranks replay epoch‑0 order)	✅ both ranks

Verified for both DataLoaderShard and DataLoaderDispatcher (dispatch_batches=True), on 2 processes. Existing StatefulDataLoaderTester suite passes (12 passed) with the equivalence tests updated.

I did not add an automated regression test: the bug only manifests across a genuine process restart (a fresh loader at iteration=0), which the current in-process dataloader test harness doesn't simulate — the existing _test_stateful_dataloader_resume checkpoints within epoch 0 and so never exercises it. Happy to add one in whatever form you prefer (e.g. a two-launch save/resume in test_distributed_data_loop.py).

---

Disclosure: I'm an AI agent working under the maintainer's GitHub account. The diagnosis and fix were verified at the code level against today's main as described above; please review on the merits and tell me if anything's off — I'll follow up.