Skip to content

fix: isolate dataset shuffle random state - #2395

Open
LZC-BELIEVER wants to merge 1 commit into
THUDM:mainfrom
LZC-BELIEVER:shuffle
Open

LZC-BELIEVER wants to merge 1 commit into
THUDM:mainfrom
LZC-BELIEVER:shuffle

Conversation

@LZC-BELIEVER

Copy link
Copy Markdown

Summary

  • Use a dataset-local random.Random instance for deterministic prompt shuffling.
  • Preserve the existing seed + epoch permutation behavior without mutating Python's process-global RNG state.
  • Add CPU regression tests for permutation compatibility, reproducibility, and RNG isolation.
  • Register the new test in the regular CPU CI matrix.

Problem

Dataset.shuffle() currently reseeds Python's module-level random number generator whenever the dataset changes epoch:

random.seed(self.seed + new_epoch_id)
random.shuffle(permutation)

RolloutDataSource invokes this method during initialization and at epoch boundaries. As a result, shuffling the prompt dataset also resets the random stream observed by unrelated code in the same long-lived rollout process, including custom generation or reward functions that use Python's random module.

This creates hidden coupling between dataset ordering and otherwise independent random operations. Adding or moving a shuffle call can therefore change behavior outside the dataset.

Fix

Create a dedicated RNG for each deterministic shuffle:

rng = random.Random(self.seed + new_epoch_id)
rng.shuffle(permutation)

The local RNG uses the same seed and standard-library shuffle implementation as before, so the dataset permutation remains deterministic for each epoch. Its state is confined to Dataset.shuffle(), leaving the process-global RNG unchanged.

No command-line arguments, dataset APIs, sampling parameters, or GPU code paths are changed.

Testing

The new CPU tests verify that:

  • the seeded permutation is preserved;
  • Dataset.shuffle() does not modify the process-global RNG state;
  • independent datasets produce the same permutation for the same seed and epoch.

Relevant CPU regression suite:

PYTHONPATH=. python -m pytest \
  tests/test_dataset_shuffle.py \
  tests/test_filter_long_prompt.py \
  tests/test_read_file_slicing.py \
  tests/test_docs_consistency.py

Result: 22 passed.

Repository checks:

pre-commit run --files \
  slime/utils/data.py \
  tests/test_dataset_shuffle.py \
  .github/workflows/pr-test.yml.j2 \
  .github/workflows/pr-test.yml \
  --show-diff-on-failure --color=never

git diff --check

All pre-commit hooks passed, and git diff --check reported no errors.

This branch has not been deployed

No deployments
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.

1 participant