[AutoParallel] Add paddle.distributed.dtensor_from_fn api (#56565)

yangxiaoyu14 · web-flow · commit 85be34f630e6 · 2023-09-12T19:13:33.000+08:00
* def dtensor_from_fn first edition

* dtensor_from_fn first edition

* Delete file /home/Paddle/build/test/auto_parallel/test_dist_tensor.py

* polish code format

* fix sample code formatting issues

* change sample codes ' &gt;&gt;&gt;' to '&gt;&gt;&gt; '

* Add static image single measurement

* modify the Indent of Sample Code

* complete the sample code modification according to ZhongKai's suggestion

* modify according to the review

* change fluid.Variable to static.Variable

* modify according to zhongkai's review

* According to Yifan's suggestion, pull the latest code to resolve conflicts

* remove paddle/fluid/ir/dialect/paddle_dialect/ir/generated/pd_ops_backward.parsed.yaml and paddle/fluid/ir/dialect/paddle_dialect/ir/generated/pd_ops.parsed.yaml
diff --git a/python/paddle/distributed/__init__.py b/python/paddle/distributed/__init__.py
@@ -66,6 +66,7 @@
 
 from .auto_parallel import shard_op  # noqa: F401
 from .auto_parallel.api import shard_tensor  # noqa: F401
+from .auto_parallel.api import dtensor_from_fn  # noqa: F401
 
 from .fleet import BoxPSDataset  # noqa: F401
 
@@ -126,4 +127,5 @@
     "ProcessMesh",
     "DistAttr",
     "shard_tensor",
+    "dtensor_from_fn",
 ]
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import paddle
 from paddle.distributed.auto_parallel.interface import (
     shard_tensor as shard_tensor_static,
@@ -68,15 +69,13 @@ def __init__(self, mesh, sharding_specs):
 
         self.process_mesh = mesh
         self.dims_mapping = dims_mapping
-
         self.mark_annotated("process_mesh")
         self.mark_annotated("dims_mapping")
 
     @property
     def sharding_specs(self):
         """
         Get sharding_specs of the dist_attr
-
         Returns:
             list[str]: sharding_specs
         """
@@ -142,3 +141,33 @@ def shard_tensor(
         return shard_tensor_static(
             data, dist_attr.process_mesh, dist_attr.sharding_specs
         )
+
+
+def dtensor_from_fn(fn, dist_attr, *args, **kwargs):
+    """
+    Construct a Distributed Tensor from a function of arguments.
+
+    Args:
+        fn (callable): A callable function that takes arguments of Distributed Tensor and returns tensor.
+        dist_attr (paddle.distributed.DistAttr): Specify how tensors are distributed or sliced on ProcessMesh.
+        *args (tuple): A tuple of arguments to be passed to the ``fn`` function.
+        **kwargs (dict): A dict of arguments to be passed to the ``fn`` function.
+
+    Retruns:
+        Tensor: A Tensor constructed from ``fn`` with distributed attributes.
+
+    Examples:
+
+        .. code-block:: python
+
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> # Create a distributed attribute
+            >>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+            >>> dist_attr = dist.DistAttr(mesh=mesh, sharding_specs=[None])
+            >>> # Call the function dtensor_from_fn with dist_attr parameter
+            >>> d_tensor = dist.dtensor_from_fn(paddle.ones, dist_attr=dist_attr, shape=[1])
+            >>> print(d_tensor)
+    """
+    tensor = fn(*args, **kwargs)
+    return shard_tensor(tensor, dist_attr=dist_attr)
diff --git a/test/auto_parallel/test_dist_tensor.py b/test/auto_parallel/test_dist_tensor.py
@@ -53,6 +53,73 @@ def test_dist_tensor_creation(self):
         self.assertEqual(dist_tensor_with_tensor.dist_attr, dist_attr)
 
 
+class TestDistTensorFromFn(unittest.TestCase):
+    def run_dtensor_from_fn(self):
+        # Create a dist_attr
+        mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+        dist_attr = dist.DistAttr(mesh=mesh, sharding_specs=[None])
+
+        # Call the function dtensor_from_fn with dist_attr parameter
+        result = dist.dtensor_from_fn(
+            paddle.ones, dist_attr=dist_attr, shape=[16]
+        )
+        # Verify the result
+        if paddle.in_dynamic_mode():
+            dist_attr.dynamic_dims = []
+            self.assertIsInstance(result, paddle.Tensor)
+            self.assertEqual(result.shape, [16])
+            self.assertEqual(result.dist_attr, dist_attr)
+        else:
+            dist_attr.dynamic_dims = [0]
+            self.assertIsInstance(result, paddle.static.Variable)
+            self.assertEqual(result.shape, (16,))
+            self.assertEqual(result.dist_attr, dist_attr)
+
+        result_zeros = dist.dtensor_from_fn(
+            paddle.zeros, dist_attr=dist_attr, shape=[16]
+        )
+        if paddle.in_dynamic_mode():
+            dist_attr.dynamic_dims = []
+            self.assertIsInstance(result, paddle.Tensor)
+            self.assertEqual(result.shape, [16])
+            self.assertEqual(result.dist_attr, dist_attr)
+        else:
+            dist_attr.dynamic_dims = [0]
+            self.assertIsInstance(result, paddle.static.Variable)
+            self.assertEqual(result.shape, (16,))
+            self.assertEqual(result.dist_attr, dist_attr)
+
+        result_random = dist.dtensor_from_fn(
+            paddle.rand, dist_attr=dist_attr, shape=[16]
+        )
+        if paddle.in_dynamic_mode():
+            dist_attr.dynamic_dims = []
+            self.assertIsInstance(result, paddle.Tensor)
+            self.assertEqual(result.shape, [16])
+            self.assertEqual(result.dist_attr, dist_attr)
+        else:
+            dist_attr.dynamic_dims = [0]
+            self.assertIsInstance(result, paddle.static.Variable)
+            self.assertEqual(result.shape, (16,))
+            self.assertEqual(result.dist_attr, dist_attr)
+
+        # Test with invalid sharding_specs length
+        with self.assertRaises(AssertionError):
+            invalid_dist_attr = dist.DistAttr(mesh=mesh, sharding_specs=['x'])
+            dist.dtensor_from_fn(
+                paddle.ones, dist_attr=invalid_dist_attr, shape=[2, 3]
+            )
+
+    def test_dynamic_mode(self):
+        self.run_dtensor_from_fn()
+
+    # Test exceptions when running in static mode
+    def test_static_mode(self):
+        paddle.enable_static()
+        self.run_dtensor_from_fn()
+        paddle.disable_static()
+
+
 class TestDistTensorForDygraphAPI(unittest.TestCase):
     def check_tensor_eq(self, a, b):
         np1 = a.numpy()
