colossalai.communication.utils

colossalai.communication.utils.send_tensor_meta(tensor, need_meta=True, next_rank=None)

Sends tensor meta information before sending a specific tensor. Since the recipient must know the shape of the tensor in p2p communications, meta information of the tensor should be sent before communications. This function synchronizes with recv_tensor_meta().

Parameters

• tensor (Tensor) – Tensor to be sent

• need_meta (bool, optional) – If False, meta information won’t be sent

• next_rank (int) – The rank of the next member in pipeline parallel group

Returns

False

Return type

bool

colossalai.communication.utils.recv_tensor_meta(tensor_shape, prev_rank=None)

Recieves tensor meta information before recieving a specific tensor. Since the recipient must know the shape of the tensor in p2p communications, meta information of the tensor should be recieved before communications. This function synchronizes with send_tensor_meta().

Parameters

• tensor_shape (torch.Size) – The shape of the tensor to be recieved

• prev_rank (int, optional) – The rank of the source of the tensor

Returns

The shape of the tensor to be recieved

Return type

torch.Size

colossalai.communication.utils.split_tensor_into_1d_equal_chunks(tensor, new_buffer=False)

Break a tensor into equal 1D chunks.

Parameters

• tensor (torch.Tensor) – Tensor to be splitted before communication

• new_buffer (bool, optional) – Whether uses a new buffer to store sliced tensor

Return splitted_tensor

The splitted tensor

Rtype splitted_tensor

torch.Tensor

colossalai.communication.utils.gather_split_1d_tensor(tensor)

Opposite of above function, gather values from model parallel ranks.

Parameters

tensor (torch.Tensor) – Tensor to be gathered after communication

Return gathered

The gathered tensor

Rtype gathered

torch.Tensor