Persist dataloader iteration and shuffle generator in save_state/load_state

[caiotheodoro]

Fixes #3996

Resuming from save_state / load_state replays the epoch-0 shuffle order. I hit it with the reporter's scripts on main @ 0f7e35f, then with a 2-process gloo run on CPU, same result both ways.

The sampler branch in checkpointing.py:137 (and :273 on load) only fires when the dataset is an IterableDatasetShard and the sampler is a SeedableRandomSampler. Those two never happen together, an iterable shard sits on _InfiniteConstantSampler, so nothing has been written for dataloaders since #2080. What's missing is two small things:

• with use_seedable_sampler=True, the seed is initial_seed + epoch (data_loader.py:99-101) and __iter__ sets that epoch from self.iteration (:582, :874 for the dispatcher). iteration is never saved, so a resumed loader starts at epoch 0 again. That's also why fix(checkpointing): save/restore SeedableRandomSampler for map-style datasets #4019 couldn't fix it by saving the sampler object, the next __iter__ overwrites the epoch anyway.
• with the default sampler and num_processes > 1, prepare_data_loader hands the sampler a private generator (data_loader.py:1247-1252) and keeps it as synchronized_generator. Its state is never saved, so prepare() recreates the same sequence on resume.

So save_state now writes {"iteration", "generator_state"} per prepared dataloader into the sampler{_i}.bin slot that already existed, and load_state puts both back when the file is there, through set_epoch so the sampler and dataset hear about it too. The generator is whichever one the sampler draws from: the private one prepare attached, or a generator= the user passed to the DataLoader, which covers the plain PyTorch reproducibility recipe in a single process as well (that one replayed too, generator state was never written). One file from the main process is enough, the generator state is the same on every rank after synchronize_rng_states, same as the optimizer and scheduler files. Old checkpoints don't have the file and load as before. With use_stateful_dataloader=True the file is written and read as well, next to torchdata's dl_state_dict.bin; the torchdata branch itself is untouched. I unwrap XLA's MpDeviceLoaderWrapper too but couldn't run XLA here.

One thing this doesn't cover: a checkpoint taken inside an epoch still resumes that epoch on the next permutation. #4233 stacks on this one and handles that.

Before, the resumed order equals epoch 0's in both configurations. After, it equals the recorded continuation on every rank. torch 2.14.0, Python 3.13.12, CPU. #3242 is the same report from 2024 on 2 GPUs.

I also ran examples/by_feature/checkpointing.py end to end (bert-base-cased, MRPC cut to 256 train / 64 eval so it finishes on CPU, 2 processes, --checkpointing_steps epoch, the example's default sampler). Uninterrupted run, rank-0 loss sum per epoch: 5.152098, 4.962221, 4.925249. Resume from epoch_0 with this branch: 4.962221, 4.925249, same numbers. Resume from the same checkpoint on main: 5.068580, 4.844967, so it trained epochs 1 and 2 on a different order.

pytest tests/test_state_checkpointing.py -q     # 24 passed, the 2 seedable cases and the user-generator case fail on main
torchrun --nproc_per_node 2 test_script.py     # check_dataloader_resume_order passes, on main it fails with
                                               #   "replayed the epoch-0 shuffle order on processes [True, True]"
pytest tests/test_data_loader.py tests/test_accelerator.py tests/test_utils.py tests/test_big_modeling.py -q   # 154 passed, 24 skipped
make quality                                   # clean on ruff 0.13.1

Drafted with Claude Opus / Fable 5.1. Reviewed by Muse Spark 1.3 and GLM 5.3 as judges before submission.

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? Please add a link
  to it if that's the case.
• 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?

@SunMarc