colossalai.communication.ring

colossalai.communication.ring.ring_forward(tensor_send_next, parallel_mode)

Sends a tensor to the next member and recieves a tensor from the previous member. This function returns the recieved tensor from the previous member.

Parameters

• tensor_send_next (torch.Tensor) – Tensor sent to next member

• parallel_mode (colossalai.context.ParallelMode) – Parallel group mode used in this communication

Returns

The tensor recieved from the previous

Return type

torch.Tensor