[AutoParallel] Revise GlobalToSub Reshard Rule (#68468)

JZ-LIANG · winter-wang · web-flow · commit 591d021d4991 · 2024-09-29T14:26:17.000+08:00
* update unitest

* align param_grad order

* bugfix for optimizer cases

* unitest

* bugfix

* fixed unitest

* fixed

* update engine

* pir unshard tensor

* trigger CI

* remove print

* bugfix

* bugfix

* update unitest cmake

* update unitest cmake

* update reshard func

* enable unitest

---------

Co-authored-by: winter-wang &lt;1030748926@qq.com&gt;
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/global_to_sub_mesh_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/global_to_sub_mesh_func.py
@@ -14,12 +14,19 @@
 
 import paddle
 
-from .base_reshard_func import ReshardFunction
+from .base_reshard_func import (
+    ReshardFunction,
+    is_replicated,
+)
+from .nd_mesh_reshard_func import NdMeshReshardFunction
 
 
 class GlobaleToSubMeshFunction(ReshardFunction):
     def is_suitable(self, src_dist_attr, dst_dist_attr):
-        if 0 in src_dist_attr.dims_mapping or 0 in src_dist_attr.partial_status:
+
+        # NOTE we could allow the src_dist_attr is not replicated and reshard it as replicated before go through the global_to_sub logic
+        # but the dst_dist_attr should be replicated otherwise there will be un-defined result when change the mesh.
+        if not is_replicated(dst_dist_attr):
             return False
         in_mesh = src_dist_attr.process_mesh
         out_mesh = dst_dist_attr.process_mesh
@@ -32,6 +39,29 @@ def is_suitable(self, src_dist_attr, dst_dist_attr):
             return out_mesh in sub_meshes
 
     def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
+
+        # reshard operand as replicated before change the mesh.
+        if not is_replicated(src_dist_attr):
+            tmp_dist_attr = (
+                paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+                    src_dist_attr.process_mesh,
+                    [-1] * len(src_dist_attr.dims_mapping),
+                    {},
+                )
+            )
+            tmp_dst_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
+                src_value.type(), tmp_dist_attr
+            )
+
+            pre_reshard_func = NdMeshReshardFunction()
+            src_value = pre_reshard_func.reshard(
+                src_dist_attr,
+                tmp_dist_attr,
+                src_value,
+                tmp_dst_type,
+            )
+            src_dist_attr = tmp_dist_attr
+
         if src_value.has_one_use():
             src_value.update_dist_attr(dst_dist_attr)
             prev_op = src_value.get_defining_op()
diff --git a/test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_llama_model_amp.py b/test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_llama_model_amp.py
@@ -12,13 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import unittest
 
 import collective.test_communication_api_base as test_base
 
-os.environ['FLAGS_enable_pir_api'] = '0'
-
 
 class TestSemiAutoParallelLlama3DAMPTest(test_base.CommunicationTestDistBase):
     def setUp(self):
