Repository: microsoft/DeepSpeed
File: deepspeed/runtime/engine.py (Python)

Code before:
        return f'{"bf16_" if bf16_mode else ""}zero_pp_rank_{dp_rank}'

    def _get_rank_zero_ckpt_name(self, checkpoints_path, tag, mp_rank, dp_rank, bf16_mode):
        file_prefix = self._get_zero_ckpt_prefix(dp_rank, bf16_mode=bf16_mode)
        zero_ckpt_name = os.path.join(
            checkpoints_path,
            str(tag),
            f"{file_prefix}_mp_rank_{mp_rank:02d}_optim_states.pt",
        )
        return zero_ckpt_name

    def _get_zero_ckpt_name(self, checkpoints_path, tag):
        mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
        pp_rank = dist.get_rank(group=self.optimizer.dp_process_group)
        bf16_mode = self.bfloat16_enabled()
        return self._get_rank_zero_ckpt_name(checkpoints_path, tag, mp_rank, pp_rank, bf16_mode)

    def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None, read_mode=False):
        if mp_placeholder is not None:
            mp_rank_str = mp_placeholder
        else:
            mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
            mp_rank_str = f"{mp_rank:02d}"

        if self.zero_optimization_partition_weights():
            # For stage 3, when loading the checkpoint, the world size may change, non-existent files may be loaded,
            # so load the first ckpt file for every worker. When writing the checkpoint, each process still writes
            # its own ckpt file to avoid conflicts.
            filename = "zero_pp_rank_{}".format(0 if read_mode else dist.get_rank(
                group=self.optimizer.dp_process_group))
            ckpt_name = os.path.join(
                checkpoints_path,
                str(tag),
                f"{filename}_mp_rank_{mp_rank_str}_model_states.pt",
            )
        else:
            ckpt_name = os.path.join(
                checkpoints_path,
                str(tag),
                "mp_rank_" + mp_rank_str + "_model_states.pt",
            )
        return ckpt_name

    def _get_optimizer_ckpt_name(self, checkpoints_path, tag, expp_rank):
        mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
        ckpt_name = os.path.join(checkpoints_path, str(tag),
                                 f'expp_rank_{expp_rank}_mp_rank_{mp_rank:02d}_optim_states.pt')
        return ckpt_name

    @staticmethod
    def _get_expert_ckpt_name(checkpoints_path, layer_id, expert_id, tag, mpu=None):

Review comment (suggestion):
Please clarify this idea as it appears to contradict the design principle that we only support world size changes through universal checkpointing. In other words, if `self.load_universal_checkpoint() == False` there is no world size change. 

Code after:
                    continue
                if param not in self.param_names:
                    raise ValueError(f"failed to find frozen {param} in named params")
                name = self.param_names[param]
                if hasattr(param, 'ds_id'):
                    param.ds_tensor.data.copy_(saved_frozen_params[name].data)
                else:
                    param.data.copy_(saved_frozen_params[name].data)

    def _get_zero_ckpt_prefix(self, dp_rank, bf16_mode):
        return f'{"bf16_" if bf16_mode else ""}zero_pp_rank_{dp_rank}'

    def _get_rank_zero_ckpt_name(self, checkpoints_path, tag, mp_rank, dp_rank, bf16_mode):
        file_prefix = self._get_zero_ckpt_prefix(dp_rank, bf16_mode=bf16_mode)
        zero_ckpt_name = os.path.join(
            checkpoints_path,
            str(tag),
            f"{file_prefix}_mp_rank_{mp_rank:02d}_optim_states.pt",
        )
        return zero_ckpt_name

    def _get_zero_ckpt_name(self, checkpoints_path, tag):
        mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
        pp_rank = dist.get_rank(group=self.optimizer.dp_process_group)
        bf16_mode = self.bfloat16_enabled()
        return self._get_rank_zero_ckpt_name(checkpoints_path, tag, mp_rank, pp_rank, bf16_mode)

    def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None):
        if mp_placeholder is not None:
            mp_rank_str = mp_placeholder
        else:
            mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
            mp_rank_str = f"{mp_rank:02d}"

        if self.zero_optimization_partition_weights():
            filename = "zero_pp_rank_{}".format(dist.get_rank(group=self.optimizer.dp_process_group))
            ckpt_name = os.path.join(
                checkpoints_path,
                str(tag),
                f"{filename}_mp_rank_{mp_rank_str}_model_states.pt",
            )
        else:
            ckpt_name = os.path.join(
                checkpoints_path,
                str(tag),
                "mp_rank_" + mp_rank_str + "_model_states.pt",
            )
        return ckpt_name

    def _get_optimizer_ckpt_name(self, checkpoints_path, tag, expp_rank):
        mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()