colossalai.communication.utils
- colossalai.communication.utils.send_tensor_meta(tensor, need_meta=True, next_rank=None)[source]
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 (torch.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)[source]
Receives tensor meta information before receiving a specific tensor. Since the recipient must know the shape of the tensor in p2p communications, meta information of the tensor should be received before communications. This function synchronizes with
send_tensor_meta().- Parameters
tensor_shape (torch.Size) – The shape of the tensor to be received.
prev_rank (int) – The rank of the source of the tensor.
- Returns
The shape of the tensor to be received.
- Return type
torch.Size
- colossalai.communication.utils.split_tensor_into_1d_equal_chunks(tensor, new_buffer=False)[source]
Break a tensor into equal 1D chunks.
- Parameters
tensor (torch.Tensor) – Tensor to be split before communication.
new_buffer (bool, optional) – Whether to use a new buffer to store sliced tensor.
- Returns
The split tensor
- Return type
torch.Size