Fix fused_rope dist op by adding time_major attr (#62180)

MarioLulab · web-flow · commit 1b38a067d2ea · 2024-02-28T19:17:05.000+08:00
* fix

* fix
diff --git a/paddle/phi/infermeta/spmd_rules/fused_rope.h b/paddle/phi/infermeta/spmd_rules/fused_rope.h
@@ -29,8 +29,8 @@ SpmdInfo FusedRopeInferSpmd(const DistMetaTensor& q,
                             const DistMetaTensor& sin,
                             const DistMetaTensor& cos,
                             const DistMetaTensor& position_ids,
-                            bool use_neox_rotary_style,
-                            bool time_major);
+                            bool use_neox_rotary_style = true,
+                            bool time_major = false);
 
 SpmdInfo FusedRopeInferSpmdReverse(const DistMetaTensor& q,
                                    const DistMetaTensor& k,
@@ -41,17 +41,17 @@ SpmdInfo FusedRopeInferSpmdReverse(const DistMetaTensor& q,
                                    const DistMetaTensor& out_q,
                                    const DistMetaTensor& out_k,
                                    const DistMetaTensor& out_v,
-                                   bool use_neox_rotary_style,
-                                   bool time_major);
+                                   bool use_neox_rotary_style = true,
+                                   bool time_major = false);
 
 SpmdInfo FusedRopeGradInferSpmd(const DistMetaTensor& sin,
                                 const DistMetaTensor& cos,
                                 const DistMetaTensor& position_ids,
                                 const DistMetaTensor& out_q_grad,
                                 const DistMetaTensor& out_k_grad,
                                 const DistMetaTensor& out_v_grad,
-                                bool use_neox_rotary_style,
-                                bool time_major);
+                                bool use_neox_rotary_style = true,
+                                bool time_major = false);
 
 }  // namespace distributed
 }  // namespace phi
diff --git a/python/paddle/distributed/auto_parallel/static/operators/dist_fused_rope.py b/python/paddle/distributed/auto_parallel/static/operators/dist_fused_rope.py
@@ -100,6 +100,7 @@ def update_dims_mapping(dist_op):
         )
 
         use_neox_rotary_style = op_desc.attr("use_neox_rotary_style")
+        time_major = op_desc.attr("time_major")
 
         # step2: infer spmd
         rule = get_phi_spmd_rule("fused_rotary_position_embedding")
@@ -112,6 +113,7 @@ def update_dims_mapping(dist_op):
             cos_spec,
             position_ids_spec,
             use_neox_rotary_style,
+            time_major,
         )
         bw_results = rule.infer_backward(
             q_spec,
@@ -124,6 +126,7 @@ def update_dims_mapping(dist_op):
             out_k_spec,
             out_v_spec,
             use_neox_rotary_style,
+            time_major,
         )
 
         # remove optional args in spmd results

Original file line number	Diff line number	Diff line change
`@@ -100,6 +100,7 @@ def update_dims_mapping(dist_op):`
`100`	`100`	`)`
`101`	`101`
`102`	`102`	`use_neox_rotary_style = op_desc.attr("use_neox_rotary_style")`
	`103`	`+ time_major = op_desc.attr("time_major")`
`103`	`104`
`104`	`105`	`# step2: infer spmd`
`105`	`106`	`rule = get_phi_spmd_rule("fused_rotary_position_embedding")`
`@@ -112,6 +113,7 @@ def update_dims_mapping(dist_op):`
`112`	`113`	`cos_spec,`
`113`	`114`	`position_ids_spec,`
`114`	`115`	`use_neox_rotary_style,`
	`116`	`+ time_major,`
`115`	`117`	`)`
`116`	`118`	`bw_results = rule.infer_backward(`
`117`	`119`	`q_spec,`
`@@ -124,6 +126,7 @@ def update_dims_mapping(dist_op):`
`124`	`126`	`out_k_spec,`
`125`	`127`	`out_v_spec,`
`126`	`128`	`use_neox_rotary_style,`
	`129`	`+ time_major,`
`127`	`130`	`)`
`128`	`131`
`129`	`132`	`# remove optional args in spmd results`