Modify codes for pr 36744

PaddlePaddle · Oct 27, 2021 · c819435 · c819435
1 parent 950ed88
commit c819435
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/python/paddle/distributed/auto_parallel/process_mesh.py b/python/paddle/distributed/auto_parallel/process_mesh.py
@@ -120,7 +120,8 @@ def ndim(self):
         return len(self._topology)
 
     def __eq__(self, other):
-        assert other and isinstance(other, ProcessMesh)
+        if not isinstance(other, ProcessMesh):
+            return False
         if self.topology != other.topology or self.processes != other.processes:
             return False
         return True

diff --git a/python/paddle/fluid/tests/unittests/test_auto_parallel_partitioner_gpt.py b/python/paddle/fluid/tests/unittests/test_auto_parallel_partitioner_gpt.py
@@ -67,16 +67,20 @@ def is_valid_completed_program(dist_context, program):
         if op_dist_attrs.process_mesh == None:
             return False
 
-        if None in op_dist_attrs._dims_mapping.values():
-            return False
+        for tensor_dist_attr in op_dist_attrs.inputs_dist_attrs.values():
+            if None == tensor_dist_attr.dims_mapping:
+                return False
+        for tensor_dist_attr in op_dist_attrs.outputs_dist_attrs.values():
+            if None == tensor_dist_attr.dims_mapping:
+                return False
 
     for var in vars_:
         var_dist_attrs = dist_context.get_tensor_dist_attr_for_program(var)
         if var_dist_attrs == None:
             return False
         elif var_dist_attrs.process_mesh == None:
             return False
-        elif var_dist_attrs.get_dims_mapping == None:
+        elif var_dist_attrs.dims_mapping == None:
             return False
 
     return True

diff --git a/python/paddle/fluid/tests/unittests/test_auto_parallel_reshard_dpmppp.py b/python/paddle/fluid/tests/unittests/test_auto_parallel_reshard_dpmppp.py
@@ -168,7 +168,7 @@ def test_mlp_dpmppp(self):
         dist_main_prog, dist_startup_prog = get_dist_prog(
             train_program, startup_program, dist_context, rank_id)
         reshard(dist_main_prog, dist_startup_prog, rank_id, dist_context)
-        print_program_with_dist_attr(dist_main_prog, dist_context)
+        # print_program_with_dist_attr(dist_main_prog, dist_context)
         # check send and recv result
         self.assertTrue(check_send_recv_result(dist_main_prog, rank_id))