[AutoParallel] send/recv_object_list function and serialize method for placement object (#72098)

xuxinyi389 · zty-king · web-flow · commit 3283385fe00e · 2025-05-08T20:50:51.000+08:00
* add___reduce___method

* support send/recv_object_list

* polish api and add tests

* typing-fix

* typing-fix

* fix_tests

---------

Co-authored-by: zty-king &lt;17786324919@163.com&gt;
diff --git a/paddle/fluid/pybind/auto_parallel_py.cc b/paddle/fluid/pybind/auto_parallel_py.cc
@@ -487,6 +487,9 @@ void BindAutoParallel(py::module *m) {
                        py::arg("memo"))
                    .def(py::self == py::self)   // NOLINT
                    .def(py::self != py::self);  // NOLINT
+  Shard.def("__reduce__", [Shard](const phi::distributed::Shard &self) {
+    return py::make_tuple(Shard, py::make_tuple(self.get_dim()));
+  });
 
   auto Replicate =
       py::class_<phi::distributed::Replicate,
@@ -522,6 +525,10 @@ void BindAutoParallel(py::module *m) {
               py::arg("memo"))
           .def(py::self == py::self)   // NOLINT
           .def(py::self != py::self);  // NOLINT
+  Replicate.def("__reduce__",
+                [Replicate](const phi::distributed::Replicate &self) {
+                  return py::make_tuple(Replicate, py::make_tuple());
+                });
 
   auto Partial =
       py::class_<phi::distributed::Partial,
@@ -562,6 +569,9 @@ void BindAutoParallel(py::module *m) {
               py::arg("memo"))
           .def(py::self == py::self)   // NOLINT
           .def(py::self != py::self);  // NOLINT
+  Partial.def("__reduce__", [Partial](const phi::distributed::Partial &self) {
+    return py::make_tuple(Partial, py::make_tuple(self.get_reduce_type()));
+  });
 
   g_placement_shard_pytype = reinterpret_cast<PyTypeObject *>(Shard.ptr());
   g_placement_replicated_pytype =
diff --git a/python/paddle/distributed/__init__.py b/python/paddle/distributed/__init__.py
@@ -94,11 +94,13 @@
     is_initialized,
     isend,
     recv,
+    recv_object_list,
     reduce,
     reduce_scatter,
     scatter,
     scatter_object_list,
     send,
+    send_object_list,
     stream,
     wait,
 )
@@ -167,6 +169,8 @@
     "destroy_process_group",
     "isend",
     "irecv",
+    "send_object_list",
+    "recv_object_list",
     "reduce_scatter",
     "is_available",
     "get_backend",
diff --git a/python/paddle/distributed/communication/__init__.py b/python/paddle/distributed/communication/__init__.py
@@ -25,8 +25,8 @@
     is_initialized,
     wait,
 )
-from .recv import irecv, recv  # noqa: F401
+from .recv import irecv, recv, recv_object_list  # noqa: F401
 from .reduce import ReduceOp, reduce  # noqa: F401
 from .reduce_scatter import reduce_scatter  # noqa: F401
 from .scatter import scatter, scatter_object_list  # noqa: F401
-from .send import isend, send  # noqa: F401
+from .send import isend, send, send_object_list  # noqa: F401
diff --git a/python/paddle/distributed/communication/group.py b/python/paddle/distributed/communication/group.py
@@ -91,6 +91,23 @@ def get_group_rank(self, rank: int) -> int | Literal[-1]:
         else:
             return -1
 
+    def get_global_rank(self, rank: int) -> int | Literal[-1]:
+        """
+        Get the global rank of a process within a group.
+
+        Args:
+            rank (int): The local rank within the group.
+
+        Returns:
+            If the current process is a member of the group, returns the corresponding global rank;
+            otherwise returns -1.
+
+        """
+        if self.is_member():
+            return self.ranks[rank]
+        else:
+            return -1
+
     def __repr__(self) -> str:
         debug_str = (
             f"rank: {self.rank}, nranks: {self.nranks}, id: {self.id}, ranks: "
diff --git a/python/paddle/distributed/communication/recv.py b/python/paddle/distributed/communication/recv.py
@@ -14,9 +14,17 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+import paddle
 from paddle.distributed.communication import stream
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _warn_cur_rank_not_in_group,
+)
+from paddle.distributed.communication.serialization_utils import (
+    convert_tensor_to_object,
+)
 
 if TYPE_CHECKING:
     from paddle import Tensor
@@ -102,3 +110,69 @@ def irecv(
             >>> # [7, 8, 9] (2 GPUs)
     """
     return recv(tensor, src, group, sync_op=False)
+
+
+def recv_object_list(
+    object_list: list[Any],
+    src: int | None = None,
+    group: Group | None = None,
+    src_in_group: int | None = None,
+):
+    """
+    Receive a list of Python objects from the sender.
+
+    Args:
+        object_list (list): The list to store received objects. Must be pre-allocated with correct size.
+        src (int, optional): The source rank id. Default: 0.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+        src_in_group (int, optional): The source rank within the group. Cannot be specified together with src. Default: None.
+
+    Returns:
+        This function does not return any value.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = ["hello", {"key": 100}, [1, 2, 3]]
+            ...     dist.send_object_list(data, dst=1)
+            >>> else:
+            ...     data = [None] * 3  # type: ignore
+            ...     dist.recv_object_list(data, src=0)
+            >>> print(data)
+            >>> # ["hello", {"key": 100}, [1, 2, 3]] (2 GPUs)
+    """
+    if object_list is None or len(object_list) == 0:
+        raise ValueError("object_list cannot be None or empty")
+
+    group = _get_global_group() if group is None else group
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if src_in_group is not None:
+        if src is not None:
+            raise ValueError(
+                "Cannot specify both 'src' and 'src_in_group' arguments."
+            )
+        src = group.get_global_rank(src_in_group)
+    else:
+        src = 0 if src is None else src
+
+    object_sizes_tensor = paddle.empty((len(object_list),), dtype='int64')
+    recv(object_sizes_tensor, src=src, group=group)
+
+    total_size = paddle.sum(object_sizes_tensor).item()
+    object_tensor = paddle.empty((total_size,), dtype=paddle.uint8)
+    recv(object_tensor, src=src, group=group)
+
+    offset = 0
+    for i, obj_size in enumerate(object_sizes_tensor):
+        obj_size = obj_size.item()
+        obj_view = object_tensor[offset : offset + obj_size]
+        object_list[i] = convert_tensor_to_object(obj_view, obj_size)
+        offset += obj_size
diff --git a/python/paddle/distributed/communication/send.py b/python/paddle/distributed/communication/send.py
@@ -14,9 +14,17 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+import paddle
 from paddle.distributed.communication import stream
+from paddle.distributed.communication.group import (
+    _get_global_group,
+    _warn_cur_rank_not_in_group,
+)
+from paddle.distributed.communication.serialization_utils import (
+    convert_object_to_tensor,
+)
 
 if TYPE_CHECKING:
     from paddle import Tensor
@@ -101,3 +109,72 @@ def isend(tensor: Tensor, dst: int, group: Group | None = None) -> task | None:
 
     """
     return send(tensor, dst, group, sync_op=False)
+
+
+def send_object_list(
+    object_list: list[Any],
+    dst: int | None = None,
+    group: Group | None = None,
+    dst_in_group: int | None = None,
+):
+    """
+    Send a list of Python objects to the receiver.
+
+    Args:
+        object_list (list): The list of Python objects to send.
+        dst (int, optional): The destination rank id. Default: 0.
+        group (Group, optional): The group instance return by new_group or None for global default group. Default: None.
+        dst_in_group (int, optional): The destination rank within the group. Cannot be specified together with dst. Default: None.
+
+    Returns:
+        This function does not return any value.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env: DISTRIBUTED)
+            >>> import paddle
+            >>> import paddle.distributed as dist
+
+            >>> dist.init_parallel_env()
+            >>> if dist.get_rank() == 0:
+            ...     data = ["hello", {"key": 100}, [1, 2, 3]]
+            ...     dist.send_object_list(data, dst=1)
+            >>> else:
+            ...     data = [None] * 3  # type: ignore
+            ...     dist.recv_object_list(data, src=0)
+            >>> print(data)
+            >>> # ["hello", {"key": 100}, [1, 2, 3]] (2 GPUs)
+    """
+    if object_list is None or len(object_list) == 0:
+        raise ValueError("object_list cannot be None or empty")
+
+    group = _get_global_group() if group is None else group
+    if _warn_cur_rank_not_in_group(group):
+        return
+
+    if dst_in_group is not None:
+        if dst is not None:
+            raise ValueError(
+                "Cannot specify both 'dst' and 'dst_in_group' arguments."
+            )
+        dst = group.get_global_rank(dst_in_group)
+    else:
+        dst = 0 if dst is None else dst
+
+    # Convert objects to tensors and get their sizes
+    tensor_list, size_list = zip(
+        *[convert_object_to_tensor(obj) for obj in object_list]
+    )
+    size_list_values = [size.item() for size in size_list]
+
+    # Send sizes first
+    object_sizes_tensor = paddle.to_tensor(size_list_values, dtype='int64')
+    send(object_sizes_tensor, dst=dst, group=group)
+
+    # Send object data
+    if len(tensor_list) == 1:
+        object_tensor = tensor_list[0]
+    else:
+        object_tensor = paddle.concat(tensor_list)
+    send(object_tensor, dst=dst, group=group)
diff --git a/test/auto_parallel/CMakeLists.txt b/test/auto_parallel/CMakeLists.txt
@@ -140,6 +140,10 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
   py_test_modules(test_moe_utils MODULES test_moe_utils)
   set_tests_properties(test_moe_utils PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE"
                                                  TIMEOUT 30)
+  py_test_modules(test_object_list_communication MODULES
+                  test_object_list_communication)
+  set_tests_properties(test_object_list_communication
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 50)
   # End of unittests WITH multi cards and timeout
 
   # NOTE(zyl): unittests WITH multi cards and WITHOUT timeout
diff --git a/test/auto_parallel/object_list_communication_demo.py b/test/auto_parallel/object_list_communication_demo.py
diff --git a/test/auto_parallel/test_object_list_communication.py b/test/auto_parallel/test_object_list_communication.py
