【Comm】fix bug in pir pass (#67386)

liym27 · web-flow · commit 31d23dacbb3f · 2024-08-14T20:35:01.000+08:00
diff --git a/paddle/phi/core/distributed/check/static_check.cc b/paddle/phi/core/distributed/check/static_check.cc
@@ -81,7 +81,12 @@ void CommStaticCheck::CheckShape(const phi::DenseTensor& out_tensor,
       out_size * out_size_factor,
       in_size * in_size_factor,
       common::errors::InvalidArgument(
-          "Input and output tensors should have matching sizes."));
+          "Input and output tensors should have matching sizes. "
+          "out_size=%ld, out_size_factor=%d, in_size=%ld, in_size_factor=%d",
+          out_size,
+          out_size_factor,
+          in_size,
+          in_size_factor));
 }
 
 void CommStaticCheck::CheckShape(const phi::DenseTensor& out_tensor,
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -367,7 +367,7 @@ def remove_other_rank_input_output_pass(dist_program):
     "pd_op.c_allreduce_sum",
     "pd_op.all_gather",
     "pd_op.c_allreduce_max",
-    "pd_op.c_reducescatter",
+    "pd_op.reduce_scatter",
 ]
 
 
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/p_to_s_reshard_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/p_to_s_reshard_func.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import paddle
-from paddle.distributed.passes.pass_utils import AutoParallelStreamType
+from paddle.distributed.utils.stream_utils import ExecutionStreamType
 
 from ..process_group import new_process_group
 from .base_reshard_func import (
@@ -81,7 +81,7 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):
             src_value, group.id, num_of_process
         )
         dst_value.get_defining_op().set_execution_stream(
-            AutoParallelStreamType.CALC_STREAM.value
+            ExecutionStreamType.DefaultStream.value
         )
 
         # set dist type and dist attr
diff --git a/python/paddle/distributed/passes/auto_parallel_sequence_parallel_optimization.py b/python/paddle/distributed/passes/auto_parallel_sequence_parallel_optimization.py
@@ -18,7 +18,7 @@
     naive_set_dist_op_attr_for_program_by_mesh,
 )
 from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY
-from paddle.distributed.passes.pass_utils import AutoParallelStreamType
+from paddle.distributed.utils.stream_utils import ExecutionStreamType
 from paddle.static import default_main_program
 
 from .auto_parallel_sharding import _is_reshard_op
@@ -135,7 +135,7 @@ def is_valid_split_op(idx, block):
                 },
             )
             new_op.dist_attr.execution_stream = (
-                AutoParallelStreamType.CALC_STREAM.value
+                ExecutionStreamType.DefaultStream.value
             )
             block._remove_op(i, False)
             block._remove_op(i - 1, False)

Original file line number	Diff line number	Diff line change
`@@ -367,7 +367,7 @@ def remove_other_rank_input_output_pass(dist_program):`
`367`	`367`	`"pd_op.c_allreduce_sum",`
`368`	`368`	`"pd_op.all_gather",`
`369`	`369`	`"pd_op.c_allreduce_max",`
`370`		`- "pd_op.c_reducescatter",`
	`370`	`+ "pd_op.reduce_scatter",`
`371`	`371`	`]`
`372`	`372`
`373`	`373`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@`
`13`	`13`	`# limitations under the License.`
`14`	`14`
`15`	`15`	`import paddle`
`16`		`-from paddle.distributed.passes.pass_utils import AutoParallelStreamType`
	`16`	`+from paddle.distributed.utils.stream_utils import ExecutionStreamType`
`17`	`17`
`18`	`18`	`from ..process_group import new_process_group`
`19`	`19`	`from .base_reshard_func import (`
`@@ -81,7 +81,7 @@ def reshard(self, src_dist_attr, dst_dist_attr, src_value, dst_type):`
`81`	`81`	`src_value, group.id, num_of_process`
`82`	`82`	`)`
`83`	`83`	`dst_value.get_defining_op().set_execution_stream(`
`84`		`- AutoParallelStreamType.CALC_STREAM.value`
	`84`	`+ ExecutionStreamType.DefaultStream.value`
`85`	`85`	`)`
`86`	`86`
`87`	`87`	`# set dist type and dist attr`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@`
`18`	`18`	`naive_set_dist_op_attr_for_program_by_mesh,`
`19`	`19`	`)`
`20`	`20`	`from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY`
`21`		`-from paddle.distributed.passes.pass_utils import AutoParallelStreamType`
	`21`	`+from paddle.distributed.utils.stream_utils import ExecutionStreamType`
`22`	`22`	`from paddle.static import default_main_program`
`23`	`23`
`24`	`24`	`from .auto_parallel_sharding import _is_reshard_op`
`@@ -135,7 +135,7 @@ def is_valid_split_op(idx, block):`
`135`	`135`	`},`
`136`	`136`	`)`
`137`	`137`	`new_op.dist_attr.execution_stream = (`
`138`		`- AutoParallelStreamType.CALC_STREAM.value`
	`138`	`+ ExecutionStreamType.DefaultStream.value`
`139`	`139`	`)`
`140`	`140`	`block._remove_op(i, False)`
`141`	`141`	`block._remove_op(i - 1, False)`