Function _split_tensor

modelscope/utils/megatron_utils.py:201–208

(tensor: torch.Tensor, num_partitions: int,
                  partition_dim: int)

def _split_tensor(tensor: torch.Tensor, num_partitions: int,
                  partition_dim: int) -> List[torch.Tensor]:
    from megatron_util import mpu
    per_partition_size = mpu.utils.divide(
        tensor.size(partition_dim), num_partitions)
    partitions_list = torch.split(
        tensor, per_partition_size, dim=partition_dim)
    return partitions_list