Add distributed checkpoint support for non-colocated MiMo#4020

yashaswikarnati merged 1 commit intoNVIDIA:mainNVIDIA/Megatron-LM:mainfrom

yashaswikarnati:yash/mimo-checkpoint-pryashaswikarnati/Megatron-LM:yash/mimo-checkpoint-prCopy head branch name to clipboard