colossalai.communication.ring

colossalai.communication.ring.ring_forward(tensor_send_next, parallel_mode)[source]

Sends a tensor to the next member and receives a tensor from the previous member. This function returns the received tensor from the previous member.

Parameters

  • tensor_send_next – Tensor sent to next member

  • parallel_mode – Parallel group mode used in this communication

Returns

The tensor received from the previous.

Return type

torch.Tensor

Note

The parallel_mode should be concluded in ParallelMode. More details about ParallelMode could be found in parallel_mode.