fix the unit test: test_semi_auto_llama_acc_align (#66955)

winter-wang · web-flow · commit fa40a78d3b3e · 2024-08-05T10:26:46.000+08:00
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1566,7 +1566,7 @@ void BindInsertionPoint(pybind11::module *m) {
                   "The insertion point is already at the begin and can't call "
                   "prev()."));
             }
-            return *(self.value.second--);
+            return *(--self.value.second);
           },
           return_value_policy::reference)
       .def(
diff --git a/paddle/phi/infermeta/spmd_rules/reshape.cc b/paddle/phi/infermeta/spmd_rules/reshape.cc
@@ -350,11 +350,11 @@ SpmdInfo ReshapeGradInferSpmd(const DistMetaTensor& x_shape,
   return {{out_grad_dist_dst}, {x_shape_dist_dst}};
 }
 
-SpmdInfo StaticReshapeGradInferSpmd(const DistMetaTensor& x_shape,
+SpmdInfo StaticReshapeGradInferSpmd(const DistMetaTensor& x,
                                     const DistMetaTensor& out_grad) {
-  auto spmd_info = ReshapeGradInferSpmd(x_shape, out_grad);
-  spmd_info.first.insert(spmd_info.first.begin(), x_shape.dist_attr());
-  return spmd_info;
+  std::vector<int64_t> out_grad_shape = common::vectorize(out_grad.dims());
+  auto tmp = ReshapeInferSpmd(x, out_grad_shape);
+  return {{tmp.first[0], tmp.second[0]}, {tmp.first[0]}};
 }
 
 }  // namespace phi::distributed
diff --git a/paddle/phi/infermeta/spmd_rules/reshape.h b/paddle/phi/infermeta/spmd_rules/reshape.h
@@ -35,7 +35,7 @@ SpmdInfo ReshapeInferSpmdDynamic(const DistMetaTensor& x,
 SpmdInfo ReshapeGradInferSpmd(const DistMetaTensor& x_shape,
                               const DistMetaTensor& out_grad);
 
-SpmdInfo StaticReshapeGradInferSpmd(const DistMetaTensor& x_shape,
+SpmdInfo StaticReshapeGradInferSpmd(const DistMetaTensor& x,
                                     const DistMetaTensor& out_grad);
 
 }  // namespace distributed
diff --git a/paddle/phi/ops/yaml/inconsistent/static_backward.yaml b/paddle/phi/ops/yaml/inconsistent/static_backward.yaml
@@ -383,7 +383,7 @@
   infer_meta :
     func : ReshapeGradInferMeta
     param : [x, out_grad]
-    spmd_rule: ReshapeGradInferSpmd
+    spmd_rule: StaticReshapeGradInferSpmd
   kernel :
     func : reshape_grad
     param : [out_grad]
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -233,6 +233,8 @@ def _remove_other_rank_params_grads(dist_params_grads):
     cur_rank = paddle.distributed.get_rank()
     need_remove_idx = []
     for idx, (_, grad) in enumerate(dist_params_grads):
+        if grad is None:
+            continue
         if cur_rank not in grad.dist_attr().process_mesh.process_ids:
             need_remove_idx.append(idx)
     for idx in need_remove_idx[::-1]:
diff --git a/python/paddle/distributed/auto_parallel/static/reshard_funcs/s_to_r_reshard_func.py b/python/paddle/distributed/auto_parallel/static/reshard_funcs/s_to_r_reshard_func.py
@@ -14,7 +14,6 @@
 
 
 import paddle
-from paddle.distributed.passes.pass_utils import AutoParallelStreamType
 
 from ..process_group import new_process_group
 from .base_reshard_func import ReshardFunction, is_replicated, is_shard
@@ -147,9 +146,6 @@ def reshard_s_to_r_with_padding(
                 src_mesh, [src_dist_attr], [new_dist_attr]
             )
         )
-        allgather_value.get_defining_op().set_execution_stream(
-            AutoParallelStreamType.CALC_STREAM.value
-        )
 
         if split_axis != 0 or padding_num != 0:
             allgather_op = allgather_value.get_defining_op()
diff --git a/test/auto_parallel/hybrid_strategy/test_semi_auto_llama_acc_align.py b/test/auto_parallel/hybrid_strategy/test_semi_auto_llama_acc_align.py
@@ -27,6 +27,7 @@ def setUp(self):
             "acc_step": "1",
             "FLAGS_embedding_deterministic": "1",
             "FLAGS_cudnn_deterministic": "1",
+            "FLAGS_enable_pir_api": "1",
         }
         self._changeable_envs = {
             "backend": ["gpu"],

Original file line number	Diff line number	Diff line change
`@@ -1566,7 +1566,7 @@ void BindInsertionPoint(pybind11::module *m) {`
`1566`	`1566`	`"The insertion point is already at the begin and can't call "`
`1567`	`1567`	`"prev()."));`
`1568`	`1568`	`}`
`1569`		`- return *(self.value.second--);`
	`1569`	`+ return *(--self.value.second);`
`1570`	`1570`	`},`
`1571`	`1571`	`return_value_policy::reference)`
`1572`	`1572`	`.def(`
Original file line number	Diff line number	Diff line change
`@@ -350,11 +350,11 @@ SpmdInfo ReshapeGradInferSpmd(const DistMetaTensor& x_shape,`
`350`	`350`	`return {{out_grad_dist_dst}, {x_shape_dist_dst}};`
`351`	`351`	`}`
`352`	`352`
`353`		`-SpmdInfo StaticReshapeGradInferSpmd(const DistMetaTensor& x_shape,`
	`353`	`+SpmdInfo StaticReshapeGradInferSpmd(const DistMetaTensor& x,`
`354`	`354`	`const DistMetaTensor& out_grad) {`
`355`		`- auto spmd_info = ReshapeGradInferSpmd(x_shape, out_grad);`
`356`		`- spmd_info.first.insert(spmd_info.first.begin(), x_shape.dist_attr());`
`357`		`- return spmd_info;`
	`355`	`+ std::vector<int64_t> out_grad_shape = common::vectorize(out_grad.dims());`
	`356`	`+ auto tmp = ReshapeInferSpmd(x, out_grad_shape);`
	`357`	`+ return {{tmp.first[0], tmp.second[0]}, {tmp.first[0]}};`
`358`	`358`	`}`
`359`	`359`
`360`	`360`	`} // namespace phi::distributed`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,6 @@`
`14`	`14`
`15`	`15`
`16`	`16`	`import paddle`
`17`		`-from paddle.distributed.passes.pass_utils import AutoParallelStreamType`
`18`	`17`
`19`	`18`	`from ..process_group import new_process_group`
`20`	`19`	`from .base_reshard_func import ReshardFunction, is_replicated, is_shard`
`@@ -147,9 +146,6 @@ def reshard_s_to_r_with_padding(`
`147`	`146`	`src_mesh, [src_dist_attr], [new_dist_attr]`
`148`	`147`	`)`
`149`	`148`	`)`
`150`		`- allgather_value.get_defining_op().set_execution_stream(`
`151`		`- AutoParallelStreamType.CALC_STREAM.value`
`152`		`- )`
`153`	`149`
`154`	`150`	`if split_axis != 0 or padding_num != 0:`
`155`	`151`	`allgather_op = allgather_value.get_defining_op()`
Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,7 @@ def setUp(self):`
`27`	`27`	`"acc_step": "1",`
`28`	`28`	`"FLAGS_embedding_deterministic": "1",`
`29`	`29`	`"FLAGS_cudnn_deterministic": "1",`
	`30`	`+ "FLAGS_enable_pir_api": "1",`
`30`	`31`	`}`
`31`	`32`	`self._changeable_envs = {`
`32`	`33`	`"backend": ["gpu"],`