Function _load_by_rank

modelscope/utils/megatron_utils.py:191–198

(checkpoint_dir: Union[str, bytes, os.PathLike],
                  rank: int)

def _load_by_rank(checkpoint_dir: Union[str, bytes, os.PathLike],
                  rank: int) -> Dict[str, torch.Tensor]:
    checkpoint_name = _CHECKPOINT_FORMAT.replace('XX', f'{rank:02d}')
    state_dict = torch.load(
        os.path.join(checkpoint_dir, checkpoint_name),
        map_location=lambda storage, loc: storage,
        weights_only=True)
    return state_dict['module'] if 'module' in state_dict else state_dict


def _split_tensor(tensor: torch.Tensor, num_partitions: int,