Function _merge_checkpoint

modelscope/utils/megatron_utils.py:153–173

(model: nn.Module, checkpoint_dir: Union[str, bytes,
                                                              os.PathLike],
                      num_partitions: int)

def _merge_checkpoint(model: nn.Module, checkpoint_dir: Union[str, bytes,
                                                              os.PathLike],
                      num_partitions: int) -> Dict[str, torch.Tensor]:
    target_rank = int(os.getenv('RANK'))
    origin_rank_list = [
        target_rank * num_partitions + i for i in range(num_partitions)
    ]
    state_dict_list = [
        _load_by_rank(checkpoint_dir, i) for i in origin_rank_list
    ]

    target_state_dict = {}
    for name, parameter in model.named_parameters():
        dim = _get_diff_dim(parameter, state_dict_list[0][name])
        if dim == -1:
            target_state_dict[name] = state_dict_list[0][name]
            continue
        target_state_dict[name] = torch.cat(
            [state_dict[name] for state_dict in state_dict_list],
            dim=dim).clone()
    return target_state_dict


def _save_converted_checkpoint(