Skip to content

Fix epoch iteration after skipping or taking a streaming subset - #8707

Open
galleonli wants to merge 1 commit into
huggingface:mainfrom
galleonli:fix/iterable-epoch-frozen-shards
Open

galleonli wants to merge 1 commit into
huggingface:mainfrom
galleonli:fix/iterable-epoch-frozen-shards

Conversation

@galleonli

Copy link
Copy Markdown

Fixes #8706.

After skip() or take(), iterating an IterableDataset at a nonzero epoch currently raises DataSourcesShufflingDisallowed. Preserve the frozen shard order during epoch preparation, while still updating shuffle-buffer RNGs.

Shift RNGs on shallow copies of the iterable wrappers so the stored pipeline remains reusable. Otherwise, handling the exception alone makes repeated iterations of the same epoch produce different orders.

Validation on main at fa995bdaea171d50d086743364f23e9db6b557d8:

  • All 12 new regression cases fail on the unmodified source and pass with the fix. They cover Python and Arrow sources, skip/take, fixed shard order, and repeatable epoch transitions through a nested map wrapper.
  • Full tests/test_iterable_dataset.py: 516 passed, 31 skipped, 2 Hub-loading failures in offline mode. Both failures passed when rerun online (2 passed).
  • CPU checks with two persistent PyTorch DataLoader workers preserve sample counts and repeatable orders across epochs 0, 1, 1, 2, 0.
  • make quality and git diff --check pass. The released 5.0.1 also reproduces the original exception.

AI assistance: Codex assisted with the investigation, patch, tests, and PR text. The checks above were run locally.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

IterableDataset skip/take fails when iterating a nonzero epoch

1 participant