github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-26 04:03:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

[doc] improved docstring in the communication module (#863)

Browse Source
This commit is contained in:
Frank Lee
2022-04-25 13:41:43 +08:00
committed by GitHub
parent 8af5f7423d
commit 8004c8e938
4 changed files with 49 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/communication/ring.py
View File

@@ -8,13 +8,13 @@ from colossalai.core import global_context as gpc
from colossalai.utils import get_current_device, synchronize
def ring_forward(tensor_send_next: torch.Tensor, parallel_mode: ParallelMode):
def ring_forward(tensor_send_next: torch.Tensor, parallel_mode: ParallelMode) -> torch.Tensor:
"""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.
Args:
tensor_send_next: Tensor sent to next member
parallel_mode: Parallel group mode used in this communication
tensor_send_next (:class:`torch.Tensor`): Tensor sent to next member
parallel_mode (ParallelMode): Parallel group mode used in this communication
Returns:
:class:`torch.Tensor`: The tensor received from the previous.