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

2e11853d · Zangwei · binmakeswell · 01cc941e · 2e11853d
Commit 2e11853d authored Apr 01, 2022 by Zangwei Committed by binmakeswell Apr 06, 2022
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 6 deletions

colossalai/communication/ring.py colossalai/communication/ring.py +4 -6

No files found.
--- a/colossalai/communication/ring.py
+++ b/colossalai/communication/ring.py
@@ -34,14 +34,12 @@ def ring_forward(tensor_send_next: torch.Tensor, parallel_mode: ParallelMode):
                                   dtype=tensor_send_next.dtype)

    # send to next rank
-    send_next_op = torch.distributed.P2POp(
-        torch.distributed.isend, tensor_send_next,
+    send_next_op = torch.distributed.P2POp(torch.distributed.isend, tensor_send_next,
                                           gpc.get_next_global_rank(parallel_mode))
    ops.append(send_next_op)

    # receive from prev rank
-    recv_prev_op = torch.distributed.P2POp(
-        torch.distributed.irecv, tensor_recv_prev,
+    recv_prev_op = torch.distributed.P2POp(torch.distributed.irecv, tensor_recv_prev,
                                           gpc.get_prev_global_rank(parallel_mode))
    ops.append(recv_prev_op)