modelscope / modelscope/ms-swift
`fix(megatron): align on-policy resume offset and sampler seed with rollout`
Nobody has claimed this yet.
- 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:
-
Resume offset. Megatron checkpoint
state.consumed_train_samplescounts
optimized rollout rows (queries × n). The on-policy dataloader indexes
unrepeated query rows._prepare_dataloaderfed the raw persisted offset intoMegatronPretrainingRandomSampler, 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. -
Sampler seed.
MegatronPretrainingRandomSamplerseeded the epoch permutation withepochalone, so every run (regardless of the configureddata_seed) shared the same shuffle order, anddata_seedhad no effect on
on-policy data ordering.
How to Reproduce / 如何复现
_prepare_dataloader(base.py): forrlhf_type in {'grpo', 'gkd'}convertconsumed_train_samplesback to query rows (//= num_generations), failing closed withRuntimeErrorif the persisted value is not divisible by
num_generations.MegatronPretrainingRandomSampler(batch_sampler.py): accept aseedparameter; useseed + epochfor the epoch permutation generator._prepare_dataloaderpassesseed=args.data_seed(available from the base SFT arguments).
Additional Information / 补充信息
No response
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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