[AutoParallel] inplace version support auto parallel (#58769)

wanghuancoder · web-flow · commit ccb37f0affe8 · 2023-11-08T15:22:40.000+08:00
* inplace version support auto parallel
diff --git a/paddle/fluid/eager/tensor_wrapper.h b/paddle/fluid/eager/tensor_wrapper.h
@@ -43,11 +43,17 @@ class TensorWrapper {
                          bool no_need_buffer = false) {
     // set inplace_version_snapshot_ according to tensor's current inplace
     // version.
-    if (tensor.impl() && phi::DenseTensor::classof(tensor.impl().get())) {
+    if (tensor.initialized() && tensor.is_dense_tensor()) {
       phi::DenseTensor* dense_tensor =
           static_cast<phi::DenseTensor*>(tensor.impl().get());
       auto& inplace_version_counter = dense_tensor->InplaceVersionCounter();
       inplace_version_snapshot_ = inplace_version_counter.CurrentVersion();
+    } else if (tensor.initialized() && tensor.is_dist_tensor()) {
+      phi::DenseTensor* dense_tensor =
+          static_cast<phi::distributed::DistTensor*>(tensor.impl().get())
+              ->unsafe_mutable_value();
+      auto& inplace_version_counter = dense_tensor->InplaceVersionCounter();
+      inplace_version_snapshot_ = inplace_version_counter.CurrentVersion();
     }
 
     /**
@@ -200,10 +206,20 @@ class TensorWrapper {
                  "no_need_buffer_ is true.";
       return;
     }
-    if (intermidiate_tensor_.impl() &&
-        phi::DenseTensor::classof(intermidiate_tensor_.impl().get())) {
-      phi::DenseTensor* dense_tensor =
-          static_cast<phi::DenseTensor*>(intermidiate_tensor_.impl().get());
+    if (intermidiate_tensor_.impl()) {
+      phi::DenseTensor* dense_tensor = nullptr;
+      if (phi::DenseTensor::classof(intermidiate_tensor_.impl().get())) {
+        dense_tensor =
+            static_cast<phi::DenseTensor*>(intermidiate_tensor_.impl().get());
+      } else if (phi::distributed::DistTensor::classof(
+                     intermidiate_tensor_.impl().get())) {
+        dense_tensor = static_cast<phi::distributed::DistTensor*>(
+                           intermidiate_tensor_.impl().get())
+                           ->unsafe_mutable_value();
+      } else {
+        return;
+      }
+
       auto& inplace_version_counter = dense_tensor->InplaceVersionCounter();
 
       uint32_t wrapper_version_snapshot = inplace_version_snapshot_;
diff --git a/test/auto_parallel/CMakeLists.txt b/test/auto_parallel/CMakeLists.txt
@@ -127,6 +127,10 @@ if(WITH_DISTRIBUTE AND WITH_GPU)
                   test_semi_auto_parallel_hybrid_strategy)
   set_tests_properties(test_semi_auto_parallel_hybrid_strategy
                        PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 120)
+  py_test_modules(test_semi_auto_parallel_dygraph_inplace MODULES
+                  test_semi_auto_parallel_dygraph_inplace)
+  set_tests_properties(test_semi_auto_parallel_dygraph_inplace
+                       PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE" TIMEOUT 100)
   py_test_modules(test_gpt_with_pir MODULES test_gpt_with_pir)
   set_tests_properties(test_gpt_with_pir PROPERTIES LABELS "RUN_TYPE=EXCLUSIVE"
                                                     TIMEOUT 100)
diff --git a/test/auto_parallel/semi_auto_parallel_dygraph_inplace.py b/test/auto_parallel/semi_auto_parallel_dygraph_inplace.py
@@ -0,0 +1,57 @@
+# 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 numpy as np
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestInplaceForSemiAutoParallel(unittest.TestCase):
+    def run_test_case(self):
+        mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+
+        x_np = np.random.random(size=[64, 32]).astype(np.float32)
+        y_np = np.random.random(size=[32, 48]).astype(np.float32)
+        x = paddle.to_tensor(x_np)
+        y = paddle.to_tensor(y_np)
+        x.stop_gradient = False
+        y.stop_gradient = False
+
+        x_dist_attr = dist.DistAttr(mesh=mesh, sharding_specs=['x', None])
+        y_dist_attr = dist.DistAttr(mesh=mesh, sharding_specs=[None, None])
+
+        dist_x = dist.shard_tensor(x_np, dist_attr=x_dist_attr)
+        dist_y = dist.shard_tensor(y_np, dist_attr=y_dist_attr)
+        dist_x.stop_gradient = False
+        dist_y.stop_gradient = False
+        dist_x = dist_x.add(dist_x)
+        dist_y = dist_y.add(dist_y)
+        dist_out = paddle.matmul(
+            dist_x, dist_y, transpose_x=False, transpose_y=False
+        )
+        dist_x.add_(dist_x)
+        dist_y.add_(dist_y)
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "received tensor_version:1 != wrapper_version_snapshot:0",
+        ):
+            dist_out.backward()
+
+
+if __name__ == '__main__':
+    TestInplaceForSemiAutoParallel().run_test_case()
diff --git a/test/auto_parallel/test_semi_auto_parallel_dygraph_inplace.py b/test/auto_parallel/test_semi_auto_parallel_dygraph_inplace.py
@@ -0,0 +1,44 @@
+# 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 TestSemiAutoParallelInplace(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(
+            num_of_devices=2,
+            timeout=120,
+        )
+        self._default_envs = {
+            "dtype": "float32",
+            "seed": "2023",
+        }
+        self._changeable_envs = {"backend": ["cpu", "gpu"]}
+
+    def test_simple_net_single_strategy(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_dygraph_inplace.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
