Function gather_split_1d_tensor

internlm/core/communication/utils.py:111–125

Opposite of above function, gather values from model parallel ranks. Args: tensor (:class:`torch.Tensor`): Tensor to be gathered after communication. Returns: :class:`torch.Tensor`: The gathered tensor.

(tensor: torch.Tensor)

def gather_split_1d_tensor(tensor: torch.Tensor) -> torch.Tensor:
    """Opposite of above function, gather values from model parallel ranks.

    Args:
        tensor (:class:`torch.Tensor`): Tensor to be gathered after communication.
    Returns:
        :class:`torch.Tensor`: The gathered tensor.
    """
    world_size = gpc.get_world_size(ParallelMode.TENSOR)
    numel = torch.numel(tensor)
    numel_gathered = world_size * numel
    gathered = torch.empty(numel_gathered, dtype=tensor.dtype, device=torch.cuda.current_device(), requires_grad=False)
    chunks = [gathered[i * numel : (i + 1) * numel] for i in range(world_size)]
    dist.all_gather(chunks, tensor, group=gpc.get_group(ParallelMode.TENSOR))
    return gathered