modelscope / modelscope/ms-swift

`fix(megatron): align on-policy resume offset and sampler seed with rollout`

Open
#10,099 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
15.7k
Forks
1.7k
Avg merge
1d 16h
Merged PRs (30d)
136

Description

Checklist / 检查清单
  • I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues,确认这是一个新的 bug report。
Bug Description / Bug 描述

Two independent dataloader bugs in on-policy (GRPO/GKD) Megatron training:

  1. Resume offset. Megatron checkpoint state.consumed_train_samples counts
    optimized rollout rows (queries × n). The on-policy dataloader indexes
    unrepeated query rows. _prepare_dataloader fed the raw persisted offset into MegatronPretrainingRandomSampler, so after a resume with n-way generation the sampler skipped n times too far — training silently continued on the wrong part of the dataset.

  2. Sampler seed. MegatronPretrainingRandomSampler seeded the epoch permutation with epoch alone, so every run (regardless of the configured data_seed) shared the same shuffle order, and data_seed had no effect on
    on-policy data ordering.

How to Reproduce / 如何复现
  • _prepare_dataloader (base.py): for rlhf_type in {'grpo', 'gkd'} convert consumed_train_samples back to query rows (//= num_generations), failing closed with RuntimeError if the persisted value is not divisible by
    num_generations.
  • MegatronPretrainingRandomSampler (batch_sampler.py): accept a seed parameter; use seed + epoch for the epoch permutation generator.
  • _prepare_dataloader passes seed=args.data_seed (available from the base SFT arguments).
Additional Information / 补充信息

No response

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in base.py at _prepare_dataloader and then read MegatronPretrainingRandomSampler in batch_sampler.py, focusing on consumed_train_samples, num_generations, and data_seed. Done means on-policy resume offsets are converted to query rows with invalid values rejected, while sampler permutations use seed plus epoch and the configured seed is passed through.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.