Skip to content

feat(rollout): add --rollout-shuffle to permute prompts per epoch - #238

Draft
Rockdu wants to merge 1 commit into
mainfrom
kangrui/rollout-shuffle
Draft

feat(rollout): add --rollout-shuffle to permute prompts per epoch#238
Rockdu wants to merge 1 commit into
mainfrom
kangrui/rollout-shuffle

Conversation

@Rockdu

@Rockdu Rockdu commented Sep 7, 2026

Copy link
Copy Markdown
Collaborator

What

--rollout-shuffle (off by default): RolloutDataSource calls Dataset.shuffle(epoch_id) at construction, at every epoch wrap, and after load(), exactly where radixark/miles does it. Dataset.shuffle already existed in miles/utils/diffusion_data.py (seeded by --rollout-seed + epoch_id) but the RL data source never invoked it, so prompts were consumed in jsonl file order.

Validation

tests/fast/rollout/test_rollout_shuffle.py (stage-a-cpu): off keeps file order; on yields a deterministic permutation of every prompt; the epoch wrap re-permutes and load() resumes the same order.

Checklist

  • pre-commit passes on the touched files
  • Test added
  • New flag in the CLI reference

@Rockdu
Rockdu force-pushed the kangrui/rollout-shuffle branch from 17f9fad to 036f9f9 Compare September 7, 2026 06:01
@Rockdu
Rockdu force-pushed the kangrui/rollout-shuffle branch from 036f9f9 to 869272d Compare September 7, 2026 06:13
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