github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish colossalai/communication/collective.py (#1262)

pull/1298/head
Zangwei Zheng 2022-07-12 18:07:15 +08:00 committed by Frank Lee
parent f1cafcc73a
commit 197a2c89e2
1 changed files with 2 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
colossalai/communication/collective.py
View File

@ -10,10 +10,7 @@ from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
def all_gather(tensor: Tensor,
dim: int,
parallel_mode: ParallelMode,
async_op: bool = False) -> Tensor:
def all_gather(tensor: Tensor, dim: int, parallel_mode: ParallelMode, async_op: bool = False) -> Tensor:
r"""Gathers all tensors from the parallel group and concatenates them in a
specific dimension.
@ -163,11 +160,7 @@ def broadcast(tensor: Tensor, src: int, parallel_mode: ParallelMode, async_op: b
return out
def reduce(tensor: Tensor,
dst: int,
parallel_mode: ParallelMode,
op: ReduceOp = ReduceOp.SUM,
async_op: bool = False):
def reduce(tensor: Tensor, dst: int, parallel_mode: ParallelMode, op: ReduceOp = ReduceOp.SUM, async_op: bool = False):
r"""Reduce tensors across whole parallel group. Only the process with
rank ``dst`` is going to receive the final result.