megatron.checkpointing.get_rng_state#

megatron.checkpointing.get_rng_state()#

collect rng state across data parallel ranks