Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/colossalai.git" did not exist on "19e1a5cf16ead982eb8818cd69e41b06a5d23b20"
Commit 2e11853d authored by Zangwei's avatar Zangwei Committed by binmakeswell
Browse files

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

parent 01cc941e
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 6 deletions
colossalai/communication/ring.py
View file @ 2e11853d
...@@ -34,15 +34,13 @@ def ring_forward(tensor_send_next: torch.Tensor, parallel_mode: ParallelMode): ...@@ -34,15 +34,13 @@ 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)
if current_rank % 2 == 0: if current_rank % 2 == 0:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment