support customized base for fused_rope (#63160)

deepllz · web-flow · commit 528a0d75709a · 2024-04-03T15:46:58.000+08:00
* support customized base for fused_rope

* fix float compare
diff --git a/paddle/phi/api/yaml/fused_backward.yaml b/paddle/phi/api/yaml/fused_backward.yaml
@@ -41,8 +41,8 @@
   support_dygraph_mode : true
 
 - backward_op : fused_rotary_position_embedding_grad
-  forward: fused_rotary_position_embedding (Tensor q, Tensor k, Tensor v, Tensor sin, Tensor cos, Tensor position_ids, bool use_neox_rotary_style, bool time_major) -> Tensor(out_q), Tensor(out_k), Tensor(out_v)
-  args : (Tensor sin, Tensor cos, Tensor position_ids, Tensor out_q_grad, Tensor out_k_grad,Tensor out_v_grad, bool use_neox_rotary_style, bool time_major)
+  forward: fused_rotary_position_embedding (Tensor q, Tensor k, Tensor v, Tensor sin, Tensor cos, Tensor position_ids, bool use_neox_rotary_style, bool time_major, float rotary_emb_base) -> Tensor(out_q), Tensor(out_k), Tensor(out_v)
+  args : (Tensor sin, Tensor cos, Tensor position_ids, Tensor out_q_grad, Tensor out_k_grad,Tensor out_v_grad, bool use_neox_rotary_style, bool time_major, float rotary_emb_base)
   output : Tensor(q_grad), Tensor(k_grad), Tensor(v_grad)
   optional :  sin, cos, position_ids, out_k_grad, out_v_grad, k_grad, v_grad
   infer_meta :
diff --git a/paddle/phi/api/yaml/fused_ops.yaml b/paddle/phi/api/yaml/fused_ops.yaml
@@ -273,7 +273,7 @@
   optional : cache_kv, pre_caches, rotary_pos_emb, time_step, seq_lengths, src_mask, gather_index
 
 - op : fused_rotary_position_embedding
-  args : (Tensor q, Tensor k, Tensor v, Tensor sin, Tensor cos, Tensor position_ids, bool use_neox_rotary_style = true, bool time_major = false)
+  args : (Tensor q, Tensor k, Tensor v, Tensor sin, Tensor cos, Tensor position_ids, bool use_neox_rotary_style = true, bool time_major = false, float rotary_emb_base = 10000.0)
   output : Tensor(out_q), Tensor(out_k), Tensor(out_v)
   infer_meta :
     func : FusedRopeInferMeta
diff --git a/paddle/phi/infermeta/backward.cc b/paddle/phi/infermeta/backward.cc
@@ -1405,6 +1405,7 @@ void FusedRopeGradInferMeta(const MetaTensor& sin,
                             const MetaTensor& dout_v,
                             bool use_neox_rotary_style,
                             bool time_major,
+                            float rotary_emb_base,
                             MetaTensor* dq,
                             MetaTensor* dk,
                             MetaTensor* dv) {
diff --git a/paddle/phi/infermeta/backward.h b/paddle/phi/infermeta/backward.h
@@ -210,6 +210,7 @@ void FusedRopeGradInferMeta(const MetaTensor& sin,
                             const MetaTensor& dout_v,
                             bool use_neox_rotary_style,
                             bool time_major,
+                            float rotary_emb_base,
                             MetaTensor* dq,
                             MetaTensor* dk,
                             MetaTensor* dv);
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -4563,6 +4563,7 @@ void FusedRopeInferMeta(const MetaTensor& q,
                         const MetaTensor& position_ids,
                         bool use_neox_rotary_style,
                         bool time_major,
+                        float rotary_emb_base,
                         MetaTensor* out_q,
                         MetaTensor* out_k,
                         MetaTensor* out_v) {
diff --git a/paddle/phi/infermeta/multiary.h b/paddle/phi/infermeta/multiary.h
@@ -904,6 +904,7 @@ void FusedRopeInferMeta(const MetaTensor& q,
                         const MetaTensor& position_ids,
                         bool use_neox_rotary_style,
                         bool time_major,
+                        float rotary_emb_base,
                         MetaTensor* out_q,
                         MetaTensor* out_k,
                         MetaTensor* out_v);
diff --git a/paddle/phi/infermeta/spmd_rules/fused_rope.cc b/paddle/phi/infermeta/spmd_rules/fused_rope.cc
@@ -241,7 +241,8 @@ SpmdInfo FusedRopeInferSpmd(const DistMetaTensor& q,
                             const DistMetaTensor& cos,
                             const DistMetaTensor& position_ids,
                             bool use_neox_rotary_style,
-                            bool time_major) {
+                            bool time_major,
+                            float rotary_emb_base) {
   check_q(q);
 
   std::vector<std::pair<std::string, std::vector<int64_t>>>
@@ -392,7 +393,8 @@ SpmdInfo FusedRopeInferSpmdReverse(const DistMetaTensor& q,
                                    const DistMetaTensor& out_k,
                                    const DistMetaTensor& out_v,
                                    bool use_neox_rotary_style,
-                                   bool time_major) {
+                                   bool time_major,
+                                   float rotary_emb_base) {
   check_q(out_q);
   std::vector<std::pair<std::string, std::vector<int64_t>>>
       outputs_sharding_info;
@@ -548,7 +550,8 @@ SpmdInfo FusedRopeGradInferSpmd(const DistMetaTensor& sin,
                                 const DistMetaTensor& out_k_grad,
                                 const DistMetaTensor& out_v_grad,
                                 bool use_neox_rotary_style,
-                                bool time_major) {
+                                bool time_major,
+                                float rotary_emb_base) {
   // NOTE(zhonghui): The forward and backward kernels of fuse rope are same, so
   // the spmd rules can be shared.
   SpmdInfo spmd_info = FusedRopeInferSpmd(out_q_grad,
diff --git a/paddle/phi/infermeta/spmd_rules/fused_rope.h b/paddle/phi/infermeta/spmd_rules/fused_rope.h
@@ -30,7 +30,8 @@ SpmdInfo FusedRopeInferSpmd(const DistMetaTensor& q,
                             const DistMetaTensor& cos,
                             const DistMetaTensor& position_ids,
                             bool use_neox_rotary_style = true,
-                            bool time_major = false);
+                            bool time_major = false,
+                            float rotary_emb_base = 10000.f);
 
 SpmdInfo FusedRopeInferSpmdReverse(const DistMetaTensor& q,
                                    const DistMetaTensor& k,
@@ -42,7 +43,8 @@ SpmdInfo FusedRopeInferSpmdReverse(const DistMetaTensor& q,
                                    const DistMetaTensor& out_k,
                                    const DistMetaTensor& out_v,
                                    bool use_neox_rotary_style = true,
-                                   bool time_major = false);
+                                   bool time_major = false,
+                                   float rotary_emb_base = 10000.f);
 
 SpmdInfo FusedRopeGradInferSpmd(const DistMetaTensor& sin,
                                 const DistMetaTensor& cos,
@@ -51,7 +53,8 @@ SpmdInfo FusedRopeGradInferSpmd(const DistMetaTensor& sin,
                                 const DistMetaTensor& out_k_grad,
                                 const DistMetaTensor& out_v_grad,
                                 bool use_neox_rotary_style = true,
-                                bool time_major = false);
+                                bool time_major = false,
+                                float rotary_emb_base = 10000.f);
 
 }  // namespace distributed
 }  // namespace phi
diff --git a/paddle/phi/kernels/fusion/gpu/fused_rope_grad_kernel.cu b/paddle/phi/kernels/fusion/gpu/fused_rope_grad_kernel.cu
@@ -33,6 +33,7 @@ void FusedRopeGradKernel(const Context& dev_ctx,
                          const paddle::optional<DenseTensor>& dout_v,
                          bool use_neox_rotary_style,
                          bool time_major,
+                         float rotary_emb_base,
                          DenseTensor* dq,
                          DenseTensor* dk,
                          DenseTensor* dv) {
@@ -132,9 +133,10 @@ void FusedRopeGradKernel(const Context& dev_ctx,
                                             head_dim,
                                             batch_stride,
                                             seq_stride,
-                                            outs_data,
                                             num_inputs,
-                                            div_c);
+                                            div_c,
+                                            rotary_emb_base,
+                                            outs_data);
 
   } else {
     // rotary position embedding Q
@@ -153,9 +155,10 @@ void FusedRopeGradKernel(const Context& dev_ctx,
                                             head_dim,
                                             batch_stride_q,
                                             seq_stride_q,
-                                            outs_data,
                                             1,
-                                            div_c);
+                                            div_c,
+                                            rotary_emb_base,
+                                            outs_data);
 
     // rotary position embedding K,V
     int64_t batch_stride_kv = time_major
@@ -178,9 +181,10 @@ void FusedRopeGradKernel(const Context& dev_ctx,
                                             head_dim,
                                             batch_stride_kv,
                                             seq_stride_kv,
-                                            out_kv,
                                             num_inputs - 1,
-                                            div_c);
+                                            div_c,
+                                            rotary_emb_base,
+                                            out_kv);
   }
 }
 
diff --git a/paddle/phi/kernels/fusion/gpu/fused_rope_kernel.cu b/paddle/phi/kernels/fusion/gpu/fused_rope_kernel.cu
@@ -33,6 +33,7 @@ void FusedRopeKernel(const Context& dev_ctx,
                      const paddle::optional<DenseTensor>& position_ids,
                      bool use_neox_rotary_style,
                      bool time_major,
+                     float rotary_emb_base,
                      DenseTensor* out_q,
                      DenseTensor* out_k,
                      DenseTensor* out_v) {
@@ -202,9 +203,10 @@ void FusedRopeKernel(const Context& dev_ctx,
                                             head_dim,
                                             batch_stride,
                                             seq_stride,
-                                            outs_data,
                                             num_inputs,
-                                            div_c);
+                                            div_c,
+                                            rotary_emb_base,
+                                            outs_data);
   } else {
     // Multi Query Attention (MQA) or Group Query Attention (GQA)
     PADDLE_ENFORCE_EQ(
@@ -245,9 +247,10 @@ void FusedRopeKernel(const Context& dev_ctx,
                                             head_dim,
                                             batch_stride_q,
                                             seq_stride_q,
-                                            outs_data,
                                             1,
-                                            div_c);
+                                            div_c,
+                                            rotary_emb_base,
+                                            outs_data);
 
     // rotary position embedding K,V
     phi::Array<const T*, 3> input_kv{ins_data[1], ins_data[2], nullptr};
@@ -270,9 +273,10 @@ void FusedRopeKernel(const Context& dev_ctx,
                                             head_dim,
                                             batch_stride_kv,
                                             seq_stride_kv,
-                                            out_kv,
                                             num_inputs - 1,
-                                            div_c);
+                                            div_c,
+                                            rotary_emb_base,
+                                            out_kv);
   }
 }
 }  // namespace fusion
diff --git a/paddle/phi/kernels/fusion/gpu/fused_rope_utils.h b/paddle/phi/kernels/fusion/gpu/fused_rope_utils.h
diff --git a/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py b/python/paddle/incubate/nn/functional/fused_rotary_position_embedding.py
