fix (#68156)

AndSonder · web-flow · commit c5f555e86aaa · 2024-09-12T19:45:35.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -697,7 +697,7 @@ def _parallel_pir(self, mode):
         # TODO(JZ-LIANG) regulization pass with pass management.
         dist_program = mix_fw_program.clone()
         apply_mix2dist_pass(dist_program)
-        set_all_ops_op_role(dist_program, OpRole.Forward)
+        set_all_ops_op_role(dist_program.global_block(), OpRole.Forward)
 
         # Step 1.2: pir backward
         if mode == "train" and self._loss and self._optimizer:
@@ -1319,7 +1319,7 @@ def _initialize(self, mode, init_parameters=True):
                 for del_op in del_ops:
                     del_op.erase()
 
-                set_all_ops_op_role(startup_prog, OpRole.Forward)
+                set_all_ops_op_role(startup_prog.global_block(), OpRole.Forward)
                 apply_reshard_pass(startup_prog)
                 for op in changed_ouput_op_list:
                     op.operand_source(0).persistable = True
diff --git a/python/paddle/distributed/auto_parallel/static/utils.py b/python/paddle/distributed/auto_parallel/static/utils.py
@@ -2510,8 +2510,8 @@ def update_grad_var_to_var(program, strategy, grad_var_to_var):
             grad_var_to_var[scale_loss_grad_var_name] = scale_loss_var_name
 
 
-def set_all_ops_op_role(main_program, op_role):
-    all_ops = main_program.global_block().ops
+def set_all_ops_op_role(block, op_role):
+    all_ops = block.ops
     for op in all_ops:
         if op.op_role == -1:
             op.op_role = op_role
diff --git a/test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh.py b/test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh.py
@@ -81,7 +81,7 @@ def create_program(self, input_shape, input_placements, output_placements):
                     dist_input, self._mesh, output_placements
                 )
             dist_program = main_program.clone()
-            set_all_ops_op_role(dist_program, OpRole.Forward)
+            set_all_ops_op_role(dist_program.global_block(), OpRole.Forward)
             apply_reshard_pass(dist_program)
 
         return main_program, dist_program
diff --git a/test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh_cross_mesh.py b/test/auto_parallel/hybrid_strategy/pir_reshard_nd_mesh_cross_mesh.py
@@ -82,7 +82,7 @@ def create_program(self, input_shape, input_placements, output_placements):
                     dist_input, self._mesh1, output_placements
                 )
             dist_program = main_program.clone()
-            set_all_ops_op_role(dist_program, OpRole.Forward)
+            set_all_ops_op_role(dist_program.global_block(), OpRole.Forward)
             apply_reshard_pass(dist_program)
 
         return main_program, dist_program
diff --git a/test/auto_parallel/pir/mlp_demo_3d.py b/test/auto_parallel/pir/mlp_demo_3d.py
@@ -107,7 +107,7 @@ def test_to_static_program(self):
         dist_program = engine._fwd_main_progs["train"]
 
         apply_mix2dist_pass(dist_program)
-        set_all_ops_op_role(dist_program, OpRole.Forward)
+        set_all_ops_op_role(dist_program.global_block(), OpRole.Forward)
         loss = dist_program.get_output_value_by_name(engine._loss_names[0])
         with paddle.static.program_guard(dist_program):
             with auto_complete_op_role(dist_program, OpRole.Backward):
diff --git a/test/auto_parallel/pir/pir_reshard_s_to_r.py b/test/auto_parallel/pir/pir_reshard_s_to_r.py
@@ -68,7 +68,7 @@ def run_pir_test_case(self):
                 reshard_tensor = paddle._C_ops.reshard(
                     input_tensor, self._mesh, [dist.Replicate()]
                 )
-            set_all_ops_op_role(main_program, OpRole.Forward)
+            set_all_ops_op_role(main_program.global_block(), OpRole.Forward)
             apply_reshard_pass(main_program)
         ops = [op.name() for op in main_program.global_block().ops]
         if self._shard == 0:
diff --git a/test/auto_parallel/pir/pir_reshard_s_to_r_cross_mesh.py b/test/auto_parallel/pir/pir_reshard_s_to_r_cross_mesh.py
@@ -62,7 +62,7 @@ def run_pir_test_case(self):
                 reshard_tensor = paddle._pir_ops.reshard(
                     input_tensor, self._out_mesh, [dist.Replicate()]
                 )
-            set_all_ops_op_role(main_program, OpRole.Forward)
+            set_all_ops_op_role(main_program.global_block(), OpRole.Forward)
             apply_reshard_pass(main_program)
 
         ops = [op.name() for op in main_program.global_block().ops]
diff --git a/test/auto_parallel/pir_reshard_r_to_s_cross_mesh.py b/test/auto_parallel/pir_reshard_r_to_s_cross_mesh.py
@@ -56,7 +56,7 @@ def run_test_case(self):
 
             old_ops = [op.name() for op in main_program.global_block().ops]
             assert 'dist_op.reshard' in old_ops
-            set_all_ops_op_role(main_program, OpRole.Forward)
+            set_all_ops_op_role(main_program.global_block(), OpRole.Forward)
             apply_reshard_pass(main_program)
         # np.testing.assert_equal(dist_program.num_ops(), 6)
         new_ops = [op.name() for op in main_program.global_block().ops]
diff --git a/test/auto_parallel/reshard_p_to_r.py b/test/auto_parallel/reshard_p_to_r.py
@@ -90,7 +90,7 @@ def run_pir_static_test_case(self):
                 reshard_tensor = paddle._C_ops.reshard(
                     input_tensor, self._mesh, [dist.Replicate()]
                 )
-            set_all_ops_op_role(main_program, OpRole.Forward)
+            set_all_ops_op_role(main_program.global_block(), OpRole.Forward)
             apply_reshard_pass(main_program)
         np.testing.assert_equal(main_program.num_ops(), 4)
         ops = main_program.global_block().ops
diff --git a/test/auto_parallel/reshard_p_to_r_cross_mesh.py b/test/auto_parallel/reshard_p_to_r_cross_mesh.py
@@ -87,7 +87,7 @@ def run_pir_static_test_case(self):
                 reshard_tensor = paddle._C_ops.reshard(
                     input_tensor, self._out_mesh, [dist.Replicate()]
                 )
-            set_all_ops_op_role(main_program, OpRole.Forward)
+            set_all_ops_op_role(main_program.global_block(), OpRole.Forward)
             apply_reshard_pass(main_program)
 
         ops = [op.name() for op in main_program.global_block().ops]
diff --git a/test/auto_parallel/reshard_r_to_s.py b/test/auto_parallel/reshard_r_to_s.py
@@ -96,7 +96,7 @@ def run_pir_test_case(self):
                     input_tensor, self._mesh, [dist.Shard(self._shard)]
                 )
             dist_program = main_program.clone()
-            set_all_ops_op_role(dist_program, OpRole.Forward)
+            set_all_ops_op_role(dist_program.global_block(), OpRole.Forward)
             apply_reshard_pass(dist_program)
             np.testing.assert_equal(dist_program.num_ops(), 6)
             old_ops = [op.name() for op in main_program.global_block().ops]

Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ def create_program(self, input_shape, input_placements, output_placements):`
`81`	`81`	`dist_input, self._mesh, output_placements`
`82`	`82`	`)`
`83`	`83`	`dist_program = main_program.clone()`
`84`		`- set_all_ops_op_role(dist_program, OpRole.Forward)`
	`84`	`+ set_all_ops_op_role(dist_program.global_block(), OpRole.Forward)`
`85`	`85`	`apply_reshard_pass(dist_program)`
`86`	`86`
`87`	`87`	`return main_program, dist_program`
Original file line number	Diff line number	Diff line change
`@@ -82,7 +82,7 @@ def create_program(self, input_shape, input_placements, output_placements):`
`82`	`82`	`dist_input, self._mesh1, output_placements`
`83`	`83`	`)`
`84`	`84`	`dist_program = main_program.clone()`
`85`		`- set_all_ops_op_role(dist_program, OpRole.Forward)`
	`85`	`+ set_all_ops_op_role(dist_program.global_block(), OpRole.Forward)`
`86`	`86`	`apply_reshard_pass(dist_program)`
`87`	`87`
`88`	`88`	`return main_program, dist_program`
Original file line number	Diff line number	Diff line change
`@@ -68,7 +68,7 @@ def run_pir_test_case(self):`
`68`	`68`	`reshard_tensor = paddle._C_ops.reshard(`
`69`	`69`	`input_tensor, self._mesh, [dist.Replicate()]`
`70`	`70`	`)`
`71`		`- set_all_ops_op_role(main_program, OpRole.Forward)`
	`71`	`+ set_all_ops_op_role(main_program.global_block(), OpRole.Forward)`
`72`	`72`	`apply_reshard_pass(main_program)`
`73`	`73`	`ops = [op.name() for op in main_program.global_block().ops]`
`74`	`74`	`if self._shard == 0:`
Original file line number	Diff line number	Diff line change
`@@ -62,7 +62,7 @@ def run_pir_test_case(self):`
`62`	`62`	`reshard_tensor = paddle._pir_ops.reshard(`
`63`	`63`	`input_tensor, self._out_mesh, [dist.Replicate()]`
`64`	`64`	`)`
`65`		`- set_all_ops_op_role(main_program, OpRole.Forward)`
	`65`	`+ set_all_ops_op_role(main_program.global_block(), OpRole.Forward)`
`66`	`66`	`apply_reshard_pass(main_program)`
`67`	`67`
`68`	`68`	`ops = [op.name() for op in main_program.global_block().ops]`
Original file line number	Diff line number	Diff line change
`@@ -90,7 +90,7 @@ def run_pir_static_test_case(self):`
`90`	`90`	`reshard_tensor = paddle._C_ops.reshard(`
`91`	`91`	`input_tensor, self._mesh, [dist.Replicate()]`
`92`	`92`	`)`
`93`		`- set_all_ops_op_role(main_program, OpRole.Forward)`
	`93`	`+ set_all_ops_op_role(main_program.global_block(), OpRole.Forward)`
`94`	`94`	`apply_reshard_pass(main_program)`
`95`	`95`	`np.testing.assert_equal(main_program.num_ops(), 4)`
`96`	`96`	`ops = main_program.global_block().ops`
Original file line number	Diff line number	Diff line change
`@@ -87,7 +87,7 @@ def run_pir_static_test_case(self):`
`87`	`87`	`reshard_tensor = paddle._C_ops.reshard(`
`88`	`88`	`input_tensor, self._out_mesh, [dist.Replicate()]`
`89`	`89`	`)`
`90`		`- set_all_ops_op_role(main_program, OpRole.Forward)`
	`90`	`+ set_all_ops_op_role(main_program.global_block(), OpRole.Forward)`
`91`	`91`	`apply_reshard_pass(main_program)`
`92`	`92`
`93`	`93`	`ops = [op.name() for op in main_program.global_block().ops]`