[auto parallel] Dist tensor set value (#59706)

FeixLiu · web-flow · commit fa5c95797160 · 2023-12-06T16:01:12.000+08:00
diff --git a/paddle/fluid/pybind/tensor.cc b/paddle/fluid/pybind/tensor.cc
@@ -1065,6 +1065,11 @@ void BindTensor(pybind11::module &m) {  // NOLINT
           py::return_value_policy::reference)
       .def("numel",
            [](DistTensor &self) -> int64_t { return self.value().numel(); })
+      .def("set",
+           [](DistTensor &self, const DistTensor &src) {
+             self.unsafe_mutable_value()->ShareDataWith(src.value());
+             return self;
+           })
       .def("_share_data_with", [](DistTensor &self, const DistTensor &src) {
         self.unsafe_set_dims(src.dims());
         self.unsafe_set_dist_attr(src.dist_attr());
diff --git a/python/paddle/base/dygraph/tensor_patch_methods.py b/python/paddle/base/dygraph/tensor_patch_methods.py
@@ -200,6 +200,10 @@ def set_value(self, value):
         assert isinstance(
             value, (np.ndarray, base_tensor, dict, str)
         ), "Variable set_value function, arguments type only support Variable, numpy, Tensor, dict, string."
+        if self.is_dist():
+            assert isinstance(
+                value, (np.ndarray, base_tensor)
+            ), "For set_value function of dist tensor, arguments type only support numpy or Tensor."
 
         if isinstance(value, (dict, str)):
             assert len(self) == len(
@@ -232,6 +236,13 @@ def set_value(self, value):
             # NOTE(wuweilong): self could be Tensor, the subsequent behavior are defined in different files
             # if self is Tensor, method value() return self that defined in this file, get_tensor() defined in eager_method.cc
             # this Interface behavior will be unifed in the future.
+            if self.is_dist():
+                # calling set method bound for DistTensor
+                value = paddle.distributed.shard_tensor(
+                    value, self.value().process_mesh, self.value().placements
+                )
+                self.value().get_tensor().set(value.get_tensor())
+                return
             self.value().get_tensor().set(
                 value, framework._current_expected_place()
             )
diff --git a/test/auto_parallel/semi_auto_parallel_set_value.py b/test/auto_parallel/semi_auto_parallel_set_value.py
@@ -0,0 +1,49 @@
+# 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 os
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestSemiAutoParallelSetValue:
+    def __init__(self):
+        self._backend = os.getenv("backend")
+        self._seed = eval(os.getenv("seed"))
+        self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+    def test_set_value_from_numpy(self):
+        p = paddle.rand(shape=[10, 10])
+        p = dist.shard_tensor(p, self._mesh, [dist.Shard(0)])
+        q = np.arange(0, 100).reshape(10, 10).astype("float32")
+        p.set_value(q)
+        np.testing.assert_equal(p.numpy(), q)
+
+    def test_set_value_from_tensor(self):
+        p = paddle.rand(shape=[10, 10])
+        p = dist.shard_tensor(p, self._mesh, [dist.Shard(0)])
+        q = paddle.arange(0, 100, dtype="float32").reshape(shape=[10, 10])
+        p.set_value(q)
+        np.testing.assert_equal(p.numpy(), q.numpy())
+
+    def run_test_case(self):
+        self.test_set_value_from_tensor()
+        self.test_set_value_from_numpy()
+
+
+if __name__ == '__main__':
+    TestSemiAutoParallelSetValue().run_test_case()
diff --git a/test/auto_parallel/test_semi_auto_parallel_single_strategy.py b/test/auto_parallel/test_semi_auto_parallel_single_strategy.py
@@ -95,6 +95,16 @@ def test_shard_optimizer_api(self):
                 user_defined_envs=envs,
             )
 
+    def test_set_value(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_set_value.py",
+                user_defined_envs=envs,
+            )
+
 
 if __name__ == "__main__":
     unittest.main()
