Function _split_checkpoint

modelscope/utils/megatron_utils.py:134–150

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

def _split_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 = target_rank // num_partitions
    state_dict = _load_by_rank(checkpoint_dir, origin_rank)

    target_state_dict = {}
    for name, parameter in model.named_parameters():
        dim = _get_diff_dim(parameter, state_dict[name])
        if dim == -1:
            target_state_dict[name] = state_dict[name]
            continue
        partitions_list = _split_tensor(state_dict[name], num_partitions, dim)
        target_state_dict[name] = partitions_list[target_rank
                                                  % num_partitions].clone()
    return target_state_dict


def _merge_checkpoint(model: nn.Module, checkpoint_dir: Union[str, bytes,