colossalai.communication.p2p
- colossalai.communication.p2p.recv_forward(input_tensor_shape, prev_rank=None, dtype=torch.float32, scatter_gather_tensors=False)
Receives the input tensor from the previous member in pipeline.
- Parameters
input_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 input tensor in forward step
- Return type
torch.Tensor
- colossalai.communication.p2p.recv_backward(output_grad_shape, next_rank=None, dtype=torch.float32, scatter_gather_tensors=False)
Receives the grad tensor from the next member in pipeline.
- Parameters
output_grad_shape (torch.Size) – The shape of the tensor to be recieved
next_rank (int, optional) – The rank of the source of the tensor
- Returns
The grad of output tensor in forward step
- Return type
torch.Tensor
- colossalai.communication.p2p.send_forward(output_tensor, next_rank=None, scatter_gather_tensors=False)
Sends the input tensor to the next member in pipeline.
- Parameters
output_tensor (
torch.Tensor) – Tensor to be sentnext_rank (int, optional) – The rank of the recipient of the tensor
- colossalai.communication.p2p.send_backward(input_tensor_grad, prev_rank=None, scatter_gather_tensors=False)
Sends the grad tensor to the previous member in pipeline.
- Parameters
input_tensor_grad (
torch.Tensor) – Tensor to be sentprev_rank (int, optional) – The rank of the recipient of the tensor
- colossalai.communication.p2p.send_forward_recv_backward(output_tensor, output_grad_shape, recv_next=True, next_rank=None, dtype=torch.float32, scatter_gather_tensors=False)
Batched communication operation. Sends the input tensor to the next member in pipeline, while recieves the grad tensor from the next member in pipeline.
- Parameters
output_tensor (
torch.Tensor) – Tensor to be sentoutput_grad_shape (
torch.Size) – The shape of the tensor to be recieved
- Returns
The grad of output tensor in forward step
- Return type
torch.Tensor
- colossalai.communication.p2p.send_backward_recv_forward(input_tensor_grad, input_tensor_shape, recv_prev=True, prev_rank=None, dtype=torch.float32, scatter_gather_tensors=False)
Batched communication operation. Sends the grad tensor to the previous member in pipeline, while recieves the input tensor from the previous member in pipeline.
- Parameters
input_tensor_grad (
torch.Tensor) – Tensor to be sentinput_tensor_shape (
torch.Size) – The shape of the tensor to be recieved
- Returns
The input tensor in forward step
- Return type
torch.Tensor
- colossalai.communication.p2p.send_forward_recv_forward(output_tensor, input_tensor_shape, recv_prev=True, prev_rank=None, next_rank=None, dtype=torch.float32, scatter_gather_tensors=False)
Batched communication operation. Sends the input tensor to the next member in pipeline, while recieves the input tensor from the previous member in pipeline.
- Parameters
output_tensor (
torch.Tensor) – Tensor to be sentinput_tensor_shape (
torch.Size) – The shape of the tensor to be recieved
- Returns
The input tensor in forward step
- Return type
torch.Tensor
- colossalai.communication.p2p.send_backward_recv_backward(input_tensor_grad, output_grad_shape, recv_next=True, prev_rank=None, next_rank=None, dtype=torch.float32, scatter_gather_tensors=False)
Batched communication operation. Sends the grad tensor to the previous member in pipeline, while recieves the grad tensor from the next member in pipeline.
- Parameters
input_tensor_grad (
torch.Tensor) – Tensor to be sentoutput_grad_shape (
torch.Size) – The shape of the tensor to be recieved
- Returns
The grad of output tensor in forward step
- Return type
torch.Tensor
- colossalai.communication.p2p.send_forward_backward_recv_forward_backward(output_tensor, input_tensor_grad, input_tensor_shape, output_grad_shape, recv_prev=True, recv_next=True, prev_rank=None, next_rank=None, dtype=torch.float32, scatter_gather_tensors=False)
Batched communication operation. Sends the input tensor to the next and the grad tensor to the previous, while recieves the grad tensor from the next and the input tensor from the previous.
- Parameters
output_tensor (
torch.Tensor) – Tensor sent to the nextinput_tensor_grad (
torch.Tensor) – Tensor sent to the previousinput_tensor_shape (
torch.Size) – The shape of the tensor recieved from the previousoutput_grad_shape (
torch.Size) – The shape of the tensor recieved from the next
- Returns
(the input tensor in forward step, the grad of output tensor in forward step)
- Return type
(Tensor, Tensor)