You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
fix(megatron): align on-policy resume offset and sampler seed with rollout #10099
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_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.
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).
Activity