[Auto Parallel] Add unshard_dtensor api (#60272)

pkuzyc · web-flow · commit ac2de38a80dd · 2023-12-27T14:33:43.000+08:00
* add dynamic part of unshard

* add unshard_dtensor api

* handle Parameter type in unshard_dtensor
diff --git a/python/paddle/distributed/__init__.py b/python/paddle/distributed/__init__.py
@@ -87,6 +87,7 @@
     to_static,
     Strategy,
     DistModel,
+    unshard_dtensor,
 )
 
 from .fleet import BoxPSDataset  # noqa: F401
@@ -169,4 +170,5 @@
     "to_static",
     "Strategy",
     "DistModel",
+    "unshard_dtensor",
 ]
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -1433,3 +1433,73 @@ def to_static(
     dist_loader = dist_model.dist_loader
 
     return dist_model, dist_loader
+
+
+def unshard_dtensor(dist_tensor):
+    """
+    Converts a distributed tensor to a dense tensor. ``unshard_dtensor``
+    first make the ``dist_tensor`` be ``Replicate`` state on all processes and
+    then converts it to a dense ``paddle.Tensor``. It can be treated as a
+    reverse operation of ``shard_tensor``.
+
+    Args:
+        dist_tensor (paddle.Tensor): The distributed tensor which is constructed
+            from a dense tensor with ``shard_tensor``.
+
+    Returns:
+        paddle.Tensor: The original dense tensor of the input ``dist_tensor``.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> from paddle.distributed import Replicate, Shard
+
+            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
+            >>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+            >>> original_tensor = paddle.rand([4, 1024, 512])
+            >>> dist_tensor = dist.shard_tensor(original_tensor, mesh, [Shard(0)])
+            >>> # dense_tensor's shape is the same as original_tensor
+            >>> dense_tensor = dist.unshard_dtensor(dist_tensor)
+    """
+    if paddle.in_dynamic_mode():
+        # if the input is not a distributed
+        # tensor, return it directly
+        if dist_tensor.is_dist() is False:
+            raise ValueError("The input should be a distributed tensor.")
+
+        mesh = dist_tensor.process_mesh
+        placements = dist_tensor.placements
+        replicate_placements = [dist.Replicate()] * len(placements)
+        r_dist_tensor = reshard(dist_tensor, mesh, replicate_placements)
+
+        if isinstance(dist_tensor, EagerParamBase):
+            return EagerParamBase.from_tensor(
+                r_dist_tensor._local_value(),
+                **dist_tensor.__dict__,
+            )
+        else:
+            return paddle.Tensor(r_dist_tensor._local_value())
+
+    else:
+        assert isinstance(
+            dist_tensor, Variable
+        ), "the input type of 'unshard_dtensor' should be Variable, but got [{}]".format(
+            dist_tensor
+        )
+        # in static mode, 'distributed tensor' and 'dense tensor' are all
+        # Varialble type, the distributed attribute is a property of the Varibale.
+        # So, it's no need to convert the distributed tensor to a dense tensor.
+        # We only need to modify its distributed attribute.
+        empty_dist_attr = (
+            dist.auto_parallel.static.dist_attribute.TensorDistAttr()
+        )
+        dist_tensor.dist_attr = empty_dist_attr
+
+        # remove the distributed tensor from dist_context
+        default_dist_ctx = get_default_distributed_context()
+        serial_tensor_id = dist_tensor.desc.original_id()
+        default_dist_ctx._dist_tensors_for_program.pop(serial_tensor_id, None)
+
+        return dist_tensor
diff --git a/test/auto_parallel/CMakeLists.txt b/test/auto_parallel/CMakeLists.txt
@@ -190,6 +190,10 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
   py_test_modules(test_gpt_with_prim MODULES test_gpt_with_prim)
   set_tests_properties(test_gpt_with_prim
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 200)
+  py_test_modules(test_semi_auto_parallel_unshard_dtensor MODULES
+                  test_semi_auto_parallel_unshard_dtensor)
+  set_tests_properties(test_semi_auto_parallel_unshard_dtensor
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
   # End of unittests WITH multi cards and timeout
 
   # NOTE(zyl): unittests WITH multi cards and WITHOUT timeout
diff --git a/test/auto_parallel/semi_auto_parallel_unshard_dtensor_api.py b/test/auto_parallel/semi_auto_parallel_unshard_dtensor_api.py
@@ -0,0 +1,99 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+import paddle.distributed as dist
+from paddle.base.dygraph.base import switch_to_static_graph
+from paddle.distributed import Replicate, Shard
+from paddle.distributed.auto_parallel.static.dist_context import (
+    get_default_distributed_context,
+)
+
+
+class TestUnshardDTensor(unittest.TestCase):
+    def __init__(self):
+        self.mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+    def run_dynamic(self):
+        ori_tensor = paddle.rand([4, 1024, 512])
+        d_tensor = dist.shard_tensor(ori_tensor, self.mesh, [Shard(0)])
+        dense_tensor = dist.unshard_dtensor(d_tensor)
+        self.assertListEqual(dense_tensor.shape, ori_tensor.shape)
+        self.assertFalse(dense_tensor.is_dist())
+
+        ori_parameter = paddle.create_parameter([1024, 512], dtype='float32')
+        d_tensor = dist.shard_tensor(ori_parameter, self.mesh, [Shard(0)])
+        dense_parameter = dist.unshard_dtensor(d_tensor)
+        self.assertListEqual(dense_parameter.shape, ori_parameter.shape)
+        self.assertFalse(dense_parameter.is_dist())
+        self.assertTrue(
+            isinstance(dense_parameter, paddle.base.framework.EagerParamBase)
+        )
+
+    @switch_to_static_graph
+    def run_static(self):
+        ori_tensor = paddle.static.data(
+            name="input",
+            shape=[4, 1024, 512],
+            dtype='float32',
+        )
+        self.assertIsNone(ori_tensor.dist_attr.process_mesh)
+        d_tensor = dist.shard_tensor(ori_tensor, self.mesh, [Shard(0)])
+
+        default_dist_context = get_default_distributed_context()
+        dist_input = default_dist_context.get_dist_tensor_for_program(
+            ori_tensor
+        )
+        self.assertEqual(dist_input.dist_attr.process_mesh, self.mesh)
+
+        dense_tensor = dist.unshard_dtensor(d_tensor)
+        dist_input = default_dist_context.get_dist_tensor_for_program(
+            ori_tensor
+        )
+        self.assertTupleEqual(dense_tensor.shape, ori_tensor.shape)
+        self.assertIsNone(dense_tensor.dist_attr.process_mesh)
+        self.assertIsNone(dist_input)
+
+    def run_dy2static(self):
+        @paddle.jit.to_static(full_graph=True)
+        def unshard_func():
+            mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+            input = paddle.rand([4, 1024, 512])
+            d_tensor = dist.shard_tensor(input, mesh, [Replicate()])
+            dense_tensor = dist.unshard_dtensor(d_tensor)
+            return input, dense_tensor
+
+        dy_ori_tensor, dy_dense_tensor = unshard_func()
+        st_ori_tensor = unshard_func.outputs[0]
+        st_dense_tensor = unshard_func.outputs[1]
+        self.assertListEqual(dy_dense_tensor.shape, dy_ori_tensor.shape)
+        self.assertFalse(dy_dense_tensor.is_dist())
+
+        default_dist_context = get_default_distributed_context()
+        dist_input = default_dist_context.get_dist_tensor_for_program(
+            st_ori_tensor
+        )
+        self.assertIsNone(st_dense_tensor.dist_attr.process_mesh)
+        self.assertIsNone(dist_input)
+
+    def run_test_cases(self):
+        self.run_dynamic()
+        self.run_static()
+        self.run_dy2static()
+
+
+if __name__ == "__main__":
+    TestUnshardDTensor().run_test_cases()
diff --git a/test/auto_parallel/test_semi_auto_parallel_unshard_dtensor.py b/test/auto_parallel/test_semi_auto_parallel_unshard_dtensor.py
@@ -0,0 +1,40 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestSemiAutoParallelUnshardDTensor(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=2,
+            timeout=100,
+        )
+        self._default_envs = {"dtype": "float32", "seed": "2023"}
+        self._changeable_envs = {"backend": ["gpu"]}
+
+    def test_api_function(self):
+        envs_list = test_base.gen_product_envs_list(
+            {"dtype": "float32", "seed": "2023"}, {"backend": ["gpu"]}
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_unshard_dtensor_api.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -87,6 +87,7 @@`
`87`	`87`	`to_static,`
`88`	`88`	`Strategy,`
`89`	`89`	`DistModel,`
	`90`	`+ unshard_dtensor,`
`90`	`91`	`)`
`91`	`92`
`92`	`93`	`from .fleet import BoxPSDataset # noqa: F401`
`@@ -169,4 +170,5 @@`
`169`	`170`	`"to_static",`
`170`	`171`	`"Strategy",`
`171`	`172`	`"DistModel",`
	`173`	`+ "unshard_dtensor",`
`172`	`174`	`]`