megatron.checkpointing.find_checkpoint_rank_0

megatron.checkpointing.find_checkpoint_rank_0(checkpoints_path, iteration, use_distributed_optimizer, release=False)

Finds the checkpoint for rank 0 without knowing if we are using pipeline parallelism or not.

Since the checkpoint naming scheme changes if pipeline parallelism is present, we need to look for both naming schemes if we don’t know if the checkpoint has pipeline parallelism.