Skip to content

fix(data): fill sample requests across multiple epochs - #2365

Open
feng1201 wants to merge 1 commit into
THUDM:mainfrom
feng1201:fix/data-source-multiple-epochs
Open

fix(data): fill sample requests across multiple epochs#2365
feng1201 wants to merge 1 commit into
THUDM:mainfrom
feng1201:fix/data-source-multiple-epochs

Conversation

@feng1201

@feng1201 feng1201 commented Sep 6, 2026

Copy link
Copy Markdown

When a request spans more than one dataset boundary, RolloutDataSource.get_samples only wraps once. With three prompts, get_samples(8) returns six groups and stores sample_offset=5; subsequent reads also diverge from the dataset stream.

Read across as many epochs as needed, preserving epoch shuffling, group/sample indices, buffered samples, and checkpoint resume. Reject positive requests from an empty dataset so the loop cannot stall. The regression uses actual data-source classes, a JSONL file, and a locally saved Transformers tokenizer; no model download or mocked data-source implementation is needed. The test is included in the existing CPU CI matrix.

Validation:

  • Original source: 12 failed, 13 passed in the new regression module.
  • Fixed source: python -m pytest tests/test_data_source.py tests/test_read_file_slicing.py tests/test_filter_long_prompt.py tests/test_sample.py -q — 52 passed.
  • All applicable pre-commit hooks passed on the four changed files.
  • Python 3.11.15, PyTorch 2.13.0 CPU, Transformers 5.15.0, Ray 2.55.0 on macOS arm64. Full GPU training was not run.

AI disclosure: implementation and local verification were performed with OpenAI Codex.

Co-authored-by: OpenAI Codex <noreply@openai.com>
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