colossalai.communication.p2p

colossalai.communication.p2p.recv_forward(input_tensor_shape, prev_rank=None, dtype=torch.float32, scatter_gather_tensors=False)

Copy the forward output from the previous stage in pipeline as the input tensor of this stage.

Parameters

• input_tensor_shape (torch.Size) – The shape of the tensor to be received.

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

Returns

The input tensor.

Return type

torch.Tensor

colossalai.communication.p2p.recv_backward(output_grad_shape, next_rank=None, dtype=torch.float32, scatter_gather_tensors=False)

Copy the gradient tensor from the next stage in pipeline as the input gradient of this stage.

Parameters

• output_grad_shape (torch.Size) – The shape of the tensor to be received.

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

Returns

The input gradient tensor.

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 stage in pipeline.

Parameters

• output_tensor (torch.Tensor) – Tensor to be sent.

• next_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 gradient tensor to the previous stage in pipeline.

Parameters

• input_tensor_grad (torch.Tensor) – Tensor to be sent

• prev_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 stage in pipeline, while receives the gradient tensor from the next stage in pipeline as the input gradient tensor of this stage.

Parameters

• output_tensor (torch.Tensor) – Tensor to be sent.

• output_grad_shape (torch.Size) – The shape of the tensor to be received.

Returns

The input gradient tensor.

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 gradient tensor to the previous stage in pipeline, while receives the output tensor from the previous stage in pipeline as the input of this stage.

Parameters

• input_tensor_grad (torch.Tensor) – Tensor to be sent.

• input_tensor_shape (torch.Size) – The shape of the tensor to be received.

Returns

The input tensor.

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 stage in pipeline, while receives the output tensor from the previous stage in pipeline as the input of this stage.

Parameters

• output_tensor (torch.Tensor) – Tensor to be sent.

• input_tensor_shape (torch.Size) – The shape of the tensor to be received.

Returns

The input tensor.

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 gradient tensor to the previous stage in pipeline, while receives the gradient tensor from the next member in pipeline as the input of this stage.

Parameters

• input_tensor_grad (torch.Tensor) – Tensor to be sent.

• output_grad_shape (torch.Size) – The shape of the tensor to be received.

Returns

The input gradient tensor.

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 stage in pipeline and the gradient tensor to the previous stage, while receives the input gradient tensor from the next stage and the input tensor from the previous stage.

Parameters

• output_tensor (torch.Tensor) – Tensor sent to the next.

• input_tensor_grad (torch.Tensor) – Tensor sent to the previous.

• input_tensor_shape (torch.Size) – The shape of the tensor received from the previous.

• output_grad_shape (torch.Size) – The shape of the tensor received from the next.

Returns

(the input tensor, the input gradient tensor)

Return type

Tuple(Tensor, Tensor)