Checklist / 检查清单
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).
Additional Information / 补充信息
No response
Checklist / 检查清单
Bug Description / Bug 描述
Two independent dataloader bugs in on-policy (GRPO/GKD) Megatron training:
Resume offset. Megatron checkpoint
state.consumed_train_samplescountsoptimized 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 onon-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 bynum_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