feat(zero-bubble): reorder comm-nodes for batch-p2p (#257)

Author: ChengYao-amd

primus/backends/megatron/core/optimizer/zbpp_optimizer.py

@@ -25,10 +25,10 @@ class ZeroBubblePPChainedOptimizer(ChainedOptimizer):
     def __init__(self, chained_optimizers: List[MegatronOptimizer]):
         super().__init__(chained_optimizers)
 
-        self.partial_reduced_total_norm = torch.FloatTensor([0])
+        self.partial_reduced_total_norm = torch.zeros([0], dtype=torch.float, device="cuda")
         self.local_total_norm = None
-        self.dummy_overflow_buf = torch.cuda.IntTensor([0])
-        self.zero_float_tensor = torch.cuda.FloatTensor([0])
+        self.dummy_overflow_buf = torch.zeros([0], dtype=torch.int, device="cuda")
+        self.zero_float_tensor = torch.zeros([0], dtype=torch.float, device="cuda")
         self.parameters_backup = None
         self.do_prev_step = False
         self.do_this_step = False

primus/backends/megatron/core/pipeline_parallel/zerobubble/runtime.py

@@ -1327,8 +1327,8 @@ def fused_pipeline_ops(
         ops.append(recv_next_op)
     if len(ops) > 0:
         reqs = torch.distributed.batch_isend_irecv(ops)
+
         # batch_isend_irecv only returns 1 handle
-        assert len(reqs) == 1
         r = reqs[0]
         # Keep the returned value consistent with p2p_pipeline_ops
         sp_reqs = [r] * len(tensor_send_prev)
@@ -1341,6 +1341,12 @@ def fused_pipeline_ops(
     return reqs, (sp_reqs, rp_reqs, sn_reqs, rn_reqs)
 
 
+class HackReq:
+    """Class to hack async p2p request because the async p2p performance bad"""
+
+    def wait(): ...
+
+
 def multi_pipeline_ops(
     tensor_send_prev: List[torch.Tensor],
     tensor_recv_prev: List[torch.Tensor],
@@ -1353,14 +1359,33 @@ def multi_pipeline_ops(
         p2p_func = fused_pipeline_ops
     else:
         p2p_func = p2p_pipeline_ops
-    return p2p_func(
+
+    reqs = p2p_func(
         tensor_send_prev=tensor_send_prev,
         tensor_recv_prev=tensor_recv_prev,
         tensor_send_next=tensor_send_next,
         tensor_recv_next=tensor_recv_next,
         group=group,
     )
 
+    if batch:
+        hack_req = HackReq()
+        hack_reqs = []
+
+        real_reqs, all_tensor_reqs = reqs
+        for req in real_reqs:
+            req.wait()
+            hack_reqs.append(hack_req)
+
+        torch.cuda.synchronize()
+        for tensor_reqs in all_tensor_reqs:
+            for req in tensor_reqs:
+                req = hack_req
+
+        reqs = (hack_reqs, all_tensor_reqs)
+
+    return reqs
+
 
 def bootstrap_and_profile_p2p_communication(config, send_tensor_shapes, recv_tensor_shapes, p2p_communicator):
     # When we fuse some send-recv communication ops in a device and can't fuse on other devices
@@ -1435,8 +1460,10 @@ def bootstrap_and_profile_p2p_communication(config, send_tensor_shapes, recv_ten
             if not parallel_state.is_pipeline_first_stage(ignore_virtual=True):
                 p2p_communicator.send_backward(recv_data, False)
         t.stop()
-        per_communication = torch.cuda.FloatTensor(
-            [t.elapsed() / (parallel_state.get_pipeline_model_parallel_world_size() - 1) / 2 / 10]
+        per_communication = torch.tensor(
+            [t.elapsed() / (parallel_state.get_pipeline_model_parallel_world_size() - 1) / 2 / 10],
+            dtype=torch.float,
+            device="cuda",
         )
         torch.distributed.all_reduce(per_communication, torch.distributed.ReduceOp.MAX)
         ScheduleTimers.comm_time = per_communication.item()

primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/communication.py

@@ -12,6 +12,8 @@
 import math
 from typing import List, Tuple
 
+from megatron.training.global_vars import get_args
+
 from primus.modules.module_utils import log_rank_all
 
 from .graph import BW, B, CommDirection, F, FuncType, GraphConfig, ScheduledNode
@@ -225,20 +227,67 @@ def add_post_validation_nodes_before_deadline(
     return local_order, comm_pairs
 
 
+def reorder_communication_nodes(local_order: List[List[ScheduledNode]]):
+    """reorder communication nodes to combine them with batch"""
+    recordered_w_list = []
+
+    def ismatch(recv_node, node):
+        return (recv_node.type == FuncType.RECV_FORWARD and node.type == FuncType.F) or (
+            recv_node.type == FuncType.RECV_BACKWARD and node.type == FuncType.B
+        )
+
+    for stage in local_order:
+        stage_list = []
+        w_list = []
+        recv_list = []
+        for node in stage:
+            if node.type == FuncType.W:
+                w_list.append(node)
+            elif node.type in (F, B, BW):
+                for i in range(len(recv_list)):
+                    if (
+                        recv_list[i].microbatch == node.microbatch
+                        and recv_list[i].chunk == node.chunk
+                        and ismatch(recv_list[i], node)
+                    ):
+                        recv_i = recv_list[i]
+                        stage_list.append(recv_i)
+                stage_list.extend(w_list)
+                w_list = []
+                stage_list.append(node)
+
+            elif node.type in (FuncType.RECV_FORWARD, FuncType.RECV_BACKWARD):
+                recv_list.append(node)
+            else:  # communication nodes
+                stage_list.append(node)
+
+        stage_list.extend(w_list)
+        recordered_w_list.append(stage_list)
+
+    return recordered_w_list
+
+
 def add_communication_nodes_without_sorting(
     config: GraphConfig,
     local_order: List[List[ScheduledNode]],
     post_validation: bool,
 ) -> List[List[ScheduledNode]]:
+
     local_order, comm_pairs = insert_send_nodes(config, local_order)
+
     if post_validation:
         local_order, post_validation_comm_pairs = add_post_validation_nodes_before_deadline(
             config, local_order
         )
         comm_pairs.extend(post_validation_comm_pairs)
     local_order = insert_recv_nodes(config, local_order, comm_pairs)
+
     if post_validation:
         local_order = tag_rollback_communication(config, local_order)
+
+    if get_args().num_virtual_stages_per_pipeline_rank is None:
+        local_order = reorder_communication_nodes(local_order)
+
     return local_order