[akshitkhokhani]

Changes

Fixes #1561.

A StatefulDataLoader and its random sampler can share a generator. Replaying a checkpoint from the sampler's initial RNG state omits intervening draws, such as the loader's worker seed. The current permutation can appear correct while later permutations—and therefore resumed training samples—diverge.

Save the RNG state immediately before the current permutation, its cursor, and the live generator state at checkpoint time. Loading regenerates only that cached permutation and restores the live RNG position. This adds constant-size RNG metadata without storing the permutation or changing ordinary sampling order. Existing initial-generator/yielded fields and the legacy loading path remain supported; old snapshots cannot recover interleaved draws they never recorded.

Tests cover replacement/nonreplacement boundaries, shared/separate generators, external RNG draws, exhausted and initial checkpoints, later epochs, legacy snapshots, and two-worker prefetch compatibility. The original implementation fails the shared-generator regressions; the prefetch case is compatibility coverage and passes the baseline.

Validation

• Original source: 12 regression cases fail; separate-generator, legacy, and prefetch compatibility controls pass.
• Full sampler test file: 14 tests and 26 subtests passed.
• Existing multi-epoch and end-of-epoch restoration classes: 8 passed, including two-worker cases.
• All applicable pre-commit hooks pass for both changed files; git diff --check passes.
• An independent run confirms the exact issue reproducer and repeated checkpoint/restore cycles preserve both samples and later RNG draws.
• Tested on macOS arm64, Python 3.12.13, PyTorch 2.14.1, using synthetic CPU data and isolated temporary/cache directories. No full-suite or GPU result is claimed.

[liangqianxing]

Independently verified the #1561 reproducer: it fails on the base and passes at eb4b7e1. The full sampler file passes (14 tests, 26 subtests), as do the existing multi-epoch and end-of-epoch classes (8 tests), on macOS arm64 with Python 3.12.14 / PyTorch 2.14.1, CPU only.

One pre-existing edge case remains in the legacy fallback: loading into an already-used iterator retains its old yielded and perm_index during replay. A partially consumed target returns the wrong sequence; an exhausted target raises StopIteration. This is also reproducible on the base.

import torch
from torchdata.stateful_dataloader.sampler import RandomSampler

sampler = RandomSampler(range(10), num_samples=70,
                        generator=torch.Generator().manual_seed(42))
iterator = iter(sampler)
for _ in range(35):
    next(iterator)
state = iterator.state_dict()
legacy = {key: state[key] for key in ("generator", "yielded")}
expected = list(iterator)
iterator.load_state_dict(legacy)  # StopIteration during replay
assert list(iterator) == expected

Resetting the counters before rebuilding the permutation fixes it:

if self.next_yielded is not None:
    self.yielded = 0
    self.perm_index = 0
    self.perm = self._get_perm()

A regression covering both replacement modes, partially consumed / exhausted targets, and restoring the same legacy snapshot twice fails in four subtests before these two resets and passes afterward. It checks the remaining samples and subsequent RNG draws. With the regression and fix applied, the sampler and the two epoch-behavior classes pass together (23 tests, 30 subtests); the applicable pre-commit hooks pass too.

[akshitkhokhani]

Thanks for checking the reproducer and the legacy path. I also reproduced the legacy restore issue on both the base and this PR's head: a partially consumed target produces different samples, and an exhausted target raises StopIteration, in both replacement modes. The existing legacy test only restores into a fresh iterator.

This legacy failure predates this PR. I'll keep this PR focused on preserving shared-generator progress. Since you've already implemented and tested the legacy fix, would you like to carry it as a follow-up? I'll avoid duplicating that work.

[liangqianxing]

Sounds good — keeping this PR scoped to #1561 makes sense. I've opened #1566 to address the legacy restore issue on used iterators separately.

The shared-generator changes here look good to me!