github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-03 10:36:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish colossalai/communication/ring.py code style (#630)

Browse Source
This commit is contained in:
Zangwei 2022-04-01 21:38:05 +08:00 committed by binmakeswell
parent 01cc941e1d
commit 2e11853d04
1 changed files with 4 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/communication/ring.py
View File

@ -34,14 +34,12 @@ def ring_forward(tensor_send_next: torch.Tensor, parallel_mode: ParallelMode):
dtype=tensor_send_next.dtype) dtype=tensor_send_next.dtype)
# send to next rank # send to next rank
send_next_op = torch.distributed.P2POp( send_next_op = torch.distributed.P2POp(torch.distributed.isend, tensor_send_next,
torch.distributed.isend, tensor_send_next,
gpc.get_next_global_rank(parallel_mode)) gpc.get_next_global_rank(parallel_mode))
ops.append(send_next_op) ops.append(send_next_op)
# receive from prev rank # receive from prev rank
recv_prev_op = torch.distributed.P2POp( recv_prev_op = torch.distributed.P2POp(torch.distributed.irecv, tensor_recv_prev,
torch.distributed.irecv, tensor_recv_prev,
gpc.get_prev_global_rank(parallel_mode)) gpc.get_prev_global_rank(parallel_mode))
ops.append(recv_prev_op) ops.append(recv_prev_op)