[Megatron] Preserve RNG state across checkpoint resume

[taking-lying-flat]

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.

[addsubmuldiv]

The RNG isolation fix looks good. One related follow-up: the resampling iterator currently resumes from the main training stream’s consumed_train_samples, even though the two streams may have consumed different numbers of samples. This can change the final batch after resume. If exact data-order reproducibility across resume is intended, we should track and restore the resampling progress separately.