Resume a mid-epoch checkpoint on the same permutation

[caiotheodoro]

Stacked on #4232, the first commit is that PR, only the second one is new here.

After #4232 an epoch-boundary checkpoint resumes fine, but one taken inside an epoch doesn't, with the default sampler. RandomSampler draws the permutation as soon as iteration starts, so by the time save_state runs mid-epoch, synchronized_generator already holds the state for the next epoch. Restore that, call skip_first_batches(dl, n), and you skip n batches of a permutation you never trained on. The seedable sampler was already fine here, its seed is initial_seed + epoch and doesn't care when you saved.

Then there's a second thing I only noticed while testing this, and it hits both samplers. skip_first_batches builds a fresh DataLoaderShard with the source's iteration (that's #4071), but the source never hears that the epoch finished. Your loop goes back to the original loader for the next epoch, __iter__ calls set_epoch(k) on it again, and you get permutation k twice in a row. examples/by_feature/checkpointing.py has exactly this shape.

What changed: DataLoaderShard.__iter__ records (iteration, generator.get_state()) right before it creates the base iterator, and save_state uses that pair as long as iteration still points at the same epoch, the live state otherwise. One case needed its own rule: saving while handling the last batch of an epoch. iteration only moves after the final yield, so that checkpoint looked like "inside epoch k" and resumed by replaying epoch k, with every later epoch off by one. The example hits it whenever checkpointing_steps divides the batches per epoch, and the checkpoint looks fine. end_of_dataloader is already set before that last batch is handed out, so save_state treats it as the end of the epoch. Test for it in both test files. skip_first_batches sets _source_dataloader on the loader it returns, _finish_epoch writes the incremented iteration back to the source, and _record_epoch_start mirrors the epoch-start state onto it, so crashing a second time inside the resumed epoch works too (I ran that one: save 3 batches in, resume, save 2 batches later, resume again, rest of the epoch and the next one match on both ranks). Nested skip_first_batches calls chain to the original loader. The write-back only happens when the source is still one epoch behind, and the recorded epoch-start state is dropped when an epoch completes and when load_state runs, so a save_state right after load_state writes what was loaded, not what this process happened to iterate before.

Nothing changes for a run that never resumes: _record_epoch_start only calls get_state(), it consumes no RNG, so an uninterrupted run produces the same order as before this PR, byte for byte. Only resumed runs move.

use_stateful_dataloader=True with the default sampler gets the same fix for free. torchdata fast-forwards the sampler by re-iterating it, so restoring the epoch-start generator state makes it redraw the same permutation, and then the rest of the epoch and the next one match on both ranks. On #4232 alone that path is wrong on both counts. That's the config the checkpointing example uses with --use_stateful_dataloader.

Still not exact, and left alone here: the paths where the permutation comes from a generator Accelerate doesn't hold, so the default sampler in a single process (global RNG, or torchdata's own generator with use_stateful_dataloader) and the default sampler with dispatch_batches=True (global RNG on the main process). load_state restores the global RNG as of the checkpoint, so a mid-epoch resume there gets a fresh permutation for the interrupted epoch, same as plain PyTorch. Epoch-boundary resume on those paths is fine after #4232. There are two more things on the dispatcher I'll send separately, both pre-existing: skip_first_batches on a DataLoaderDispatcher wraps the unsharded batch sampler, so it skips n global batches instead of n per-rank ones, and DataLoaderDispatcher.set_epoch only looks at batch_sampler.sampler, which a SkipBatchSampler doesn't have, so the seedable sampler's epoch never reaches the skipping loader. Either one makes the rest of a resumed epoch wrong under dispatch_batches=True (the following epoch is right).

Checked on CPU, torch 2.14.0, 2 processes over gloo: save after 3 batches of epoch 2, resume with skip_first_batches(dl, 3), compare the rest of epoch 2 and all of epoch 3 against the uninterrupted run.

                                        rest of epoch 2   epoch 3
main + #4232, default                   no                no
main + #4232, seedable                  yes               no
main + #4232, default + stateful        no                no
this PR, default and seedable           yes               yes
this PR, default + stateful             yes               yes
this PR, seedable + dispatch            no (skip bugs)    yes
this PR, default + dispatch             no (global RNG)   no
this PR, save on the last batch         next epoch        yes

pytest tests/test_state_checkpointing.py -q     # 32 passed; the mid-epoch and last-batch cases fail on #4232 alone
python test_script.py                            # single process, passes (the mid-epoch part is gated on a private generator)
torchrun --nproc_per_node 2 test_script.py       # check_dataloader_resume_order walks the mid-epoch path on both ranks, stateful too
make quality                                      # clean on ruff 0.13.1

One design question rather than a change: the single-process default path could get the same private generator prepare_data_loader already attaches in the multi-process branch. Every non-dispatch path would then resume exactly and the gate in the tests would go away. The cost is that a single-process run's epoch-0 order would change versus earlier releases, since the seed would be drawn at prepare instead of the first __iter__. I didn't do it here because of that; happy to if you'd rather have one code path.

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