Restore RNG state before fetching resumed training batches
Author: KirtiRamchandani
Pull request: yes
Labels: name: Code agent slop

## Problem

On checkpoint resume inside an epoch, `Trainer` fetched the next micro-batches with `get_batch_samples()` before calling `_load_rng_state()`. Any randomness in the dataloader (random augmentations, noisy datasets, etc.) therefore ran on the live RNG stream instead of the restored checkpoint state.

## Root cause (updated after @chikoneuru's verification)

Moving `_load_rng_state()` before `get_batch_samples()` is not enough. `skip_first_batches` still materializes discarded batches inside the first `next(epoch_iterator)` (via `SkipBatchSampler`), which advances the global RNG **between** restore and the first real training fetch.

## Fix

For mid-epoch resume (`rng_to_sync`):
1. Stop using `skip_first_batches` (it couples skip + fetch in one iterator step).
2. Fast-forward `epoch_iterator` through `steps_trained_in_current_epoch` discarded micro-batches.
3. Call `_load_rng_state(resume_from_checkpoint)` **after** the skip, **before** the first `get_batch_samples()` fetch.

This mirrors what the epoch-boundary path already does (restore before any real fetch), without the hidden RNG advance inside `skip_first_batches`.

## Testing

- `repro_39215.py` from #39215 — mid-epoch resume now bit-reproducible (`first divergent item = None`)
- `tests/trainer/test_trainer_checkpointing.py::test_mid_epoch_resume_restores_rng_before_first_fetch`

Fixes #39215

Comments:
@chikoneuru you were right — the swap alone did not close it. Thanks for running the step-keyed verifier; that decomposition (correct indices, wrong augmentation draws, third RNG fingerprint) was the missing piece.

**What v1 got wrong**

\`skip_first_batches\` is not index-only. \`SkipBatchSampler.__iter__\` still walks the underlying batch sampler for every skipped index and materializes those batches (so \`__getitem__\` runs and burns the global torch stream). That consumption is bundled into the *first* \`next(epoch_iterator)\` inside \`get_batch_samples()\`.

So shuffling \`_load_rng_state()\` around \`_run_epoch\` L1710–1719 could never land in the right place:
- **stock** (restore *after* fetch): first resumed item sees fresh \`set_seed\` state
- **v1** (restore *before* fetch): first resumed item sees checkpoint state **advanced by the skip** — neither fresh nor checkpoint

**What v2 does** (\`f932dd2\`, \`trainer.py\` L1686–1705)

1. Drop \`skip_first_batches\` on the mid-epoch path (L1686–1689 — only set \`rng_to_sync\`, no wrapper).
2. After \`epoch_iterator = iter(train_dataloader)\` (L1695), explicitly fast-forward \`steps_trained_in_current_epoch\` micro-batches (L1701–1703).
3. Restore checkpoint RNG **after** that skip, **before** the first real \`get_batch_samples()\` call (L1704–1705).

Epoch-boundary resume is unchanged: still restores at L1690–1691 before any fetch.

**Evidence**

\`repro_39215.py\` (your harness) on \`5.13.0.dev0\` + torch \`2.12.0\`:

\`\`\`
mid-epoch resume    : first divergent post-checkpoint item = None -> reproducible
epoch-boundary ctrl : first divergent post-checkpoint item = None -> reproducible
\`\`\`

Regression test added:
\`tests/trainer/test_trainer_checkpointing.py::TrainerResumeTrainingTest::test_mid_epoch_resume_restores_rng_before_first_fetch\` — compares per-item \`torch.rand\` draws from item 8 onward; passes in 0.23s on CPU.

@SunMarc — could you take a look when you have a minute? This is the mid-epoch path for #39215; happy to adjust if you would rather bracket the skip with an RNG snapshot instead of dropping \`skip_first_batches\`.

**CI Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46628)

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46628&sha=67fec5