Fix stateful dataloader checkpointing across processes

[jrsmartin]

What does this PR do?

This PR fixes stateful dataloader checkpointing during distributed training by saving each process's state separately. A stateful dataloader can hold process-specific cursor, dataset, worker, and RNG state. Currently, every process writes its state to the same dl_state_dict.bin filename in a shared checkpoint directory, so the final file contains whichever process wrote last, and other processes can restore the wrong data position.

Distributed checkpoints now use rank-qualified filenames such as dl_state_dict_rank0.bin. Loading prefers the current process's file and falls back to the legacy unqualified filename, which preserves compatibility with existing checkpoints. Single-process checkpoint filenames remain unchanged.

Related context

Discussed in #3080, where manually saving and loading the dataloader state on every rank is described as the current workaround. This PR addresses that per-rank checkpoint persistence gap.

Tests

• Added a two-process regression using distinct per-rank dataset streams and a shared checkpoint directory. It verifies exact batch replay after restore and asserts that both rank-qualified state files exist.
• Added a CPU/Gloo entry point so the distributed regression runs in the normal pull-request test suite as well as the existing multi-device suite. It skips when the optional torchdata dependency is unavailable.
• Added unit coverage for single-process naming, multiple dataloaders, process-specific load precedence, and legacy fallback.
• Verified the regression failed before the implementation and passes afterward.
• make quality
• pytest -q tests/test_state_checkpointing.py -k "not map_location" — 18 passed
• pytest -q tests/test_accelerator.py -k stateful_dataloader — 17 passed
• CPU/Gloo two-process regression — passed
• Two-GPU distributed data-loop suite — passed

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue or the forum? Discussed in #3080, where per-rank dataloader state saving/loading is described as the required workaround. This PR addresses that per-rank checkpoint persistence gap
• Did you make sure to update the documentation with your changes? Here are the documentation guidelines, and here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed.

[jrsmartin]

@SunMarc would you mind reviewing this PR when you have a chance - thank you!

[jrsmartin]

@SunMarc would you mind reviewing this PR when you have a chance - thank you!

Friendly bump @SunMarc, and tagging @muellerzr since this addresses the per-rank StatefulDataLoader checkpoint workaround discussed in #3080. Would either of you be able to review? Cheers

[GoldenStain]

Thanks for this PR — the per-rank filenames are a real fix for the shared-file collision.

One gap worth flagging during review: the new regression
(test_stateful_dataloader_save_state_per_process) builds
DataLoader(dataset, batch_size=1, num_workers=0), which uses the default
SequentialSampler — so its "exact batch replay" assertion passes even though
only cursor counters carry values (sequential order is derivable from the
cursor alone). With shuffle=True the saved state's sampler slots
(_index_sampler_state, _sampler_iter_state, _shared_seed) are all None:
BatchSamplerShard doesn't implement Stateful, so torchdata's collection
skips the wrapped chain entirely, and resume restores the position but not
the order.

Full analysis + reproducer: #4195. A shuffle=True variant of this test
would fail today and would pin the remaining gap in the same suite — happy to
help with that case if useful.