Skip to content

[Megatron] Preserve RNG state across checkpoint resume - #9935

Open
taking-lying-flat wants to merge 1 commit into
modelscope:mainfrom
taking-lying-flat:agent/fix-megatron-rng-checkpoint-resume
Open

[Megatron] Preserve RNG state across checkpoint resume#9935
taking-lying-flat wants to merge 1 commit into
modelscope:mainfrom
taking-lying-flat:agent/fix-megatron-rng-checkpoint-resume

Conversation

@taking-lying-flat

@taking-lying-flat taking-lying-flat commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Summary

  • preserve checkpoint-restored RNG state when GRPO/GKD initialize resampling iterators
  • give resampling batch samplers and DataLoaders private generators seeded independently from training
  • save RNG state when save_safetensors=true, no_save_optim=true, and no_save_rng=false

Root cause

After checkpoint loading restored Python, NumPy, Torch, CUDA, and Megatron CUDA-tracker state, GRPO and GKD resampling initialization called set_random_seed(args.seed + 1) and then set_random_seed(args.seed). This replaced the restored training RNG position with the initial seed.

For safetensors-only checkpoints, BaseMegatronTrainer intentionally passes an empty model list to the MCore checkpoint writer when optimizer state is omitted. RNG collection was conditional on that list being non-empty, so RNG was silently omitted even when no_save_rng=false.

Fix

Resampling now receives its private seed through the batch sampler and DataLoader generators without modifying global or model RNG state. The default sampler seed remains unchanged for existing training DataLoaders.

RNG collection now follows no_save_rng instead of the model list. A model-less checkpoint containing sharded RNG state uses the distributed checkpoint path so the existing load path can restore it.

Impact

GRPO/GKD resume no longer resets restored RNG state while creating dynamic/deletion resampling iterators. Safetensors checkpoints without optimizer state retain RNG whenever the user has not disabled RNG saving.

@taking-lying-flat
taking-lying-flat marked this pull request as ready for review August 17, 2026 12:10
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