Implement support for passing dictionary arguments in Pipeline Parallel

Author: Hz188

python/paddle/distributed/fleet/meta_parallel/pipeline_parallel.py

@@ -757,27 +757,37 @@ def forward_backward_pipeline(
                 schedule += f"f{step_id};"
                 logger.info(f"forward step for micro step {step_id}")
                 continue
+
             input_tensor = self._p2p_helper.recv_forward(
                 self.is_pipeline_first_stage(),
                 batch_p2p_comm=self._use_batch_p2p_comm,
             )
 
+            input_tensor_dict, use_dict = tuple_to_dict_helper(input_tensor)
+
             self._record_stamp("F", step_id, '"B"', self._forward_color)
             output_tensor, _, _ = self._forward_step(
-                input_tensor, micro_dataset, step_id=step_id
+                input_tensor=input_tensor_dict if use_dict else input_tensor,
+                micro_dataset=micro_dataset,
+                step_id=step_id,
             )
+
+            # convert dict to tuple whose tensor element has a key attribution
+            output_tensor_tuple = dict_to_tuple_helper(output_tensor)
+
             self._record_stamp("F", step_id, '"E"', self._forward_color)
+            # fwd output dict -> send tuple
             self._p2p_helper.send_forward(
-                output_tensor,
-                self.is_pipeline_last_stage(),
+                output_tensor=output_tensor_tuple,
+                pp_last_stage=self.is_pipeline_last_stage(),
                 batch_p2p_comm=self._use_batch_p2p_comm,
             )
 
             input_buffers.append(input_tensor)
-            output_buffers.append(output_tensor)
+            output_buffers.append(output_tensor_tuple)
 
             if not self.is_pipeline_last_stage():
-                self._release_output(output_tensor)
+                self._release_output(output_tensor_tuple)
 
         if steady_steps > 0 and not static_scheduler:
             input_tensor = self._p2p_helper.recv_forward(
@@ -794,27 +804,33 @@ def forward_backward_pipeline(
                 continue
             last_iter = i == (steady_steps - 1)
 
+            input_tensor_dict, use_dict = tuple_to_dict_helper(input_tensor)
+
             self._record_stamp(
                 "F", startup_steps + i, '"B"', self._forward_color
             )
             output_tensor, _, _ = self._forward_step(
-                input_tensor, micro_dataset, step_id=startup_steps + i
+                input_tensor=input_tensor_dict if use_dict else input_tensor,
+                micro_dataset=micro_dataset,
+                step_id=startup_steps + i,
             )
             self._record_stamp(
                 "F", startup_steps + i, '"E"', self._forward_color
             )
 
+            output_tensor_tuple = dict_to_tuple_helper(output_tensor)
+
             output_tensor_grad = self._p2p_helper.send_forward_recv_backward(
-                output_tensor,
+                output_tensor_tuple,
                 self.is_pipeline_last_stage(),
                 batch_p2p_comm=self._use_batch_p2p_comm,
             )
 
             input_buffers.append(input_tensor)
-            output_buffers.append(output_tensor)
+            output_buffers.append(output_tensor_tuple)
 
             if not self.is_pipeline_last_stage():
-                self._release_output(output_tensor)
+                self._release_output(output_tensor_tuple)
 
             input_tensor, output_tensor = input_buffers.pop(
                 0
@@ -1692,18 +1708,22 @@ def _forward_step_helper(
 
         input_tensor = self._get_forward_input(virtual_pp_rank)
 
+        input_tensor_dict, use_dict = tuple_to_dict_helper(input_tensor)
+
         output_tensor, schedule_chunk, loss_fn_node = self._forward_step(
-            input_tensor,
+            input_tensor_dict if use_dict else input_tensor,
             micro_dataset,
-            virtual_pp_rank,
+            virtual_pp_rank,  # chunk_id
             step_id=micro_step,
             overlap_schedule_mode=overlap_schedule_mode,
         )
 
+        output_tensor_tuple = dict_to_tuple_helper(output_tensor)
+
         self._store_forward_outputs(
-            virtual_pp_rank, output_tensor, schedule_chunk, loss_fn_node
+            virtual_pp_rank, output_tensor_tuple, schedule_chunk, loss_fn_node
         )
-        return output_tensor
+        return output_tensor_tuple
 
     def _overlap_comm_grads(self):
         if self._comm_overlap:
@@ -2953,7 +2973,6 @@ def forward_backward_pipeline(
             )
         )
 
-        # run startup steps
         for micro_step in range(num_steps):
             output_tensor = self._forward_step_helper(micro_dataset, micro_step)
             # determine whether recv forward tensor or not
@@ -3433,3 +3452,41 @@ def forward_backward_pipeline(
         self.processed_steps += 1
         self._check_user_hooks_status_at_step_end()
         return train_loss
+
+
+def tuple_to_dict_helper(input_tensor):
+    # recv tuple -> fwd input dict
+    use_dict = False
+    if isinstance(input_tensor, tuple):
+        use_dict = hasattr(input_tensor[0], "key")
+    else:  # single tensor
+        use_dict = hasattr(input_tensor, "key")
+    if use_dict:
+        input_tensor = convert_tensor_tuple_to_dict(input_tensor)
+    return input_tensor, use_dict
+
+
+def dict_to_tuple_helper(output_tensor):
+    if isinstance(output_tensor, dict):
+        output_tensor_tuple = convert_tensor_dict_to_tuple(
+            output_tensor_dict=output_tensor
+        )
+    else:  # single tensor or tensor tuple
+        output_tensor_tuple = output_tensor
+    return output_tensor_tuple
+
+
+def convert_tensor_dict_to_tuple(output_tensor_dict):
+    for key, tensor in output_tensor_dict.items():
+        tensor.key = key
+
+    return tuple(output_tensor_dict.values())
+
+
+def convert_tensor_tuple_to_dict(input_tensor_tuple):
+    input_tensor_dict = {}
+    for tensor in input_tensor_tuple:
+        key = tensor.key
+        input_tensor_dict[key] = tensor
+        delattr(tensor, "key")
+    return input_tensor_dict

python/paddle/distributed/fleet/meta_parallel/pp_utils/p2p_communication.py

@@ -24,6 +24,10 @@
     _get_global_group,
     _warn_cur_rank_not_in_group,
 )
+from paddle.distributed.communication.serialization_utils import (
+    convert_object_to_tensor,
+    convert_tensor_to_object,
+)
 from paddle.framework.recall_error import check_naninf
 from paddle.utils import strtobool
 
@@ -58,10 +62,12 @@ def __init__(self):
     def init_or_erase_meta(self):
         self.send_shape_message = None
         self.send_dtype_message = None
+        self.send_key_message = None
 
         self.recv_shape_message = None
         self.recv_dtype_message = None
         self.recv_stop_gradient = None
+        self.recv_key_message = None
 
         self.has_send_meta = False
         self.has_recv_meta = False
@@ -99,17 +105,31 @@ def recv_meta(self, group, reverse=False, broadcast=False):
         shapes = []
         dtypes = []
         stop_grads = []
+        keys = []
 
         for _ in range(tensor_num):
             shape_len = data.pop(0)
             shape = data[:shape_len]
             data = data[shape_len:]
             dtype_number = data.pop(0)
             stop_gradient = bool(data.pop(0))
+            # ------------------tensor key meta send-------------
+            key_len = data.pop(0)
+            key_data = data[:key_len]
+            if key_len > 0:
+                key = convert_tensor_to_object(
+                    paddle.to_tensor(key_data).astype("uint8"),
+                    paddle.to_tensor(key_len),
+                )
+            else:
+                key = None
+            data = data[key_len:]
+            # ------------------tensor key meta send-------------
 
             shapes.append(shape)
             dtypes.append(dtype_number)
             stop_grads.append(stop_gradient)
+            keys.append(key)
 
         assert (
             len(data) == 0
@@ -119,10 +139,12 @@ def recv_meta(self, group, reverse=False, broadcast=False):
             self.recv_shape_message = shapes[0]
             self.recv_dtype_message = dtypes[0]
             self.recv_stop_gradient = stop_grads[0]
+            self.recv_key_message = keys[0]
         else:
             self.recv_shape_message = tuple(shapes)
             self.recv_dtype_message = tuple(dtypes)
             self.recv_stop_gradient = tuple(stop_grads)
+            self.recv_key_message = tuple(keys)
 
     def send_meta(self, tensor, group, reverse=False, broadcast=False):
         if reverse:
@@ -152,12 +174,24 @@ def send_meta(self, tensor, group, reverse=False, broadcast=False):
 
         for t in tensors_to_send:
             assert isinstance(t, paddle.Tensor)
+            # ------------------tensor key meta send-------------
+            if hasattr(t, "key"):
+                current_tensor_name = t.key
+                key_data_tensor, _ = convert_object_to_tensor(
+                    current_tensor_name
+                )
+                key_data = key_data_tensor.numpy().tolist()
+            else:
+                key_data = []
+            # ------------------tensor key meta send-------------
             data.extend(
                 [
                     len(t.shape),
                     *t.shape,
                     paddle_2_number(t.dtype),
                     int(t.stop_gradient),
+                    len(key_data),
+                    *key_data,
                 ]
             )
 
@@ -184,35 +218,44 @@ def send_meta(self, tensor, group, reverse=False, broadcast=False):
 
     def _obtain_send_message(self, tensor):
         if isinstance(tensor, paddle.Tensor):
-            return tensor.shape, paddle_2_number(tensor.dtype)
+            key = tensor.key if hasattr(tensor, "key") else None
+            return tensor.shape, paddle_2_number(tensor.dtype), key
         else:
             shapes = []
             dtypes = []
+            keys = []
             for d in tensor:
                 assert isinstance(d, paddle.Tensor)
                 if d.stop_gradient:
                     continue
-                shape, dtype = self._obtain_send_message(d)
+                shape, dtype, key = self._obtain_send_message(d)
                 shapes.append(shape)
                 dtypes.append(dtype)
-            return tuple(shapes), tuple(dtypes)
+                keys.append(key)
+            return tuple(shapes), tuple(dtypes), tuple(keys)
 
     def set_send_message(self, tensor):
         (
             self.send_shape_message,
             self.send_dtype_message,
+            self.send_key_message,  # (key1_str, key2_str, key3_str ... )
         ) = self._obtain_send_message(tensor)
 
     def check_send_message(self, tensor):
         if self.send_shape_message is None or self.send_dtype_message is None:
             return
-        actual_shape, actual_dtype = self._obtain_send_message(tensor)
+        actual_shape, actual_dtype, actual_key = self._obtain_send_message(
+            tensor
+        )
         assert (
             self.send_shape_message == actual_shape
         ), f"send_shape_message: {self.send_shape_message}, actual_shape: {actual_shape}"
         assert (
             self.send_dtype_message == actual_dtype
         ), f"send_dtype_message: {self.send_dtype_message}, actual_dtype: {actual_dtype}"
+        assert (
+            self.send_key_message == actual_key
+        ), f"send_key_message: {self.send_key_message}, actual_key: {actual_key}"
 
     def __repr__(self):
         return f"send_shape_message: {self.send_shape_message}, send_dtype_message: {self.send_dtype_message}, recv_shape_message: {self.recv_shape_message}, recv_dtype_message: {self.recv_dtype_message}, recv_stop_gradient: {self.recv_stop_gradient}"
@@ -619,9 +662,11 @@ def _p2p_helper(
     recv_shape_msg = send_recv_meta.recv_shape_message
     recv_dtype_msg = send_recv_meta.recv_dtype_message
     recv_stop_gradient = send_recv_meta.recv_stop_gradient
+    recv_key_msg = send_recv_meta.recv_key_message
 
     send_shape_msg = send_recv_meta.send_shape_message
     send_dtype_msg = send_recv_meta.send_dtype_message
+    # backward has no key meta message
 
     # model parallel message
     mp_group = _hcg.get_model_parallel_group()
@@ -636,13 +681,17 @@ def _p2p_helper(
                     shape=shape, dtype=number_2_dtype(recv_dtype_msg[idx])
                 )
                 tmp.stop_gradient = recv_stop_gradient[idx]
+                if recv_key_msg[idx] is not None:
+                    tmp.key = recv_key_msg[idx]
                 tensor_recv_prev.append(tmp)
             tensor_recv_prev = tuple(tensor_recv_prev)
         else:
             tensor_recv_prev = paddle.empty(
                 shape=recv_shape_msg, dtype=number_2_dtype(recv_dtype_msg)
             )
             tensor_recv_prev.stop_gradient = recv_stop_gradient
+            if recv_key_msg is not None:
+                tensor_recv_prev.key = recv_key_msg
 
     if recv_next:
         if dynamic_shape:

test/collective/fleet/CMakeLists.txt

@@ -836,3 +836,17 @@ if((WITH_GPU) AND LOCAL_ALL_PLAT)
   )
   set_tests_properties(test_shutdown_process_group PROPERTIES TIMEOUT "200")
 endif()
+if((WITH_GPU) AND LOCAL_ALL_PLAT)
+  bash_test_modules(
+    test_pp_send_recv_dict
+    START_BASH
+    ../../legacy_test/dist_test.sh
+    TIMEOUT
+    "500"
+    LABELS
+    "RUN_TYPE=DIST"
+    ENVS
+    "PADDLE_DIST_UT_PORT=21282;http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python"
+  )
+  set_tests_properties(test_pp_send_recv_dict PROPERTIES TIMEOUT "500")
+endif()