megatron.checkpointing.get_rng_state# megatron.checkpointing.get_rng_state()# collect rng state across data parallel ranks