[Feature] add non-norm_topk in fused moe (#67425)

penPenf28 · web-flow · commit f7c5585fc115 · 2024-08-15T00:11:31.000-07:00
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -5912,6 +5912,7 @@ void FusedMoeInferMeta(const MetaTensor& X,
                        const MetaTensor& ffn2_bias,
                        const std::string& quant_method,
                        const int moe_topk,
+                       const bool norm_topk_prob,
                        MetaTensor* out) {
   out->set_dims(X.dims());
   out->share_lod(X);
diff --git a/paddle/phi/infermeta/multiary.h b/paddle/phi/infermeta/multiary.h
@@ -1176,6 +1176,7 @@ void FusedMoeInferMeta(const MetaTensor& X,
                        const MetaTensor& ffn2_bias,
                        const std::string& quant_method,
                        const int moe_topk,
+                       const bool norm_topk_prob,
                        MetaTensor* out);
 
 void FusedMultiHeadAttentionInferMeta(const MetaTensor& query,
diff --git a/paddle/phi/kernels/fusion/cutlass/fused_moe_kernel.cu b/paddle/phi/kernels/fusion/cutlass/fused_moe_kernel.cu
@@ -56,6 +56,7 @@ void FusedMoeKernel(const Context& ctx,
                     const DenseTensor& ffn2_bias,
                     const std::string& quant_method,
                     const int moe_topk,
+                    const bool norm_topk_prob,
                     DenseTensor* out) {
   out->Resize(X.dims());
   auto* output_data = ctx.template Alloc<T>(out);
@@ -85,6 +86,7 @@ void FusedMoeKernel(const Context& ctx,
                          &ffn2_bias,
                          nullptr,
                          moe_topk,
+                         norm_topk_prob,
                          "ffn",
                          out);
 }
diff --git a/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_helper.h b/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_helper.h
@@ -132,6 +132,7 @@ class MoeHelper {
                   const DenseTensor *ffn2_bias,
                   const DenseTensor *moe_token_type_ids,
                   const int moe_topk,
+                  const bool norm_topk_prob,
                   const std::string moe_type,
                   DenseTensor *output) {
     auto *input_activations = X->data<T>();
@@ -413,6 +414,7 @@ class MoeHelper {
           hidden_size,
           k,
           static_cast<int>(1),
+          norm_topk_prob,
           ctx.stream());
     } else {
       finalize_moe_routing_kernelLauncher(
@@ -427,6 +429,7 @@ class MoeHelper {
           inter_size,
           k,
           static_cast<int>(0),
+          norm_topk_prob,
           ctx.stream());
     }
     VLOG(4) << " Finished EXPERT \n";
diff --git a/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_op.h b/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_op.h
@@ -723,7 +723,8 @@ __global__ void finalize_moe_routing_kernel(
     const int* expert_for_source_row,
     const int cols,
     const int k,
-    const int compute_bias) {
+    const int compute_bias,
+    const bool norm_topk_prob) {
   const int original_row = blockIdx.x;
   const int num_rows = gridDim.x;
   T* reduced_row_ptr = reduced_unpermuted_output + original_row * cols;
@@ -753,7 +754,8 @@ __global__ void finalize_moe_routing_kernel(
                           bias_ptr[tid] *
                               static_cast<T>(static_cast<float>(compute_bias)));
     }
-    thread_output = static_cast<float>(thread_output) / row_rescale;
+    thread_output = static_cast<float>(thread_output) /
+                    (norm_topk_prob ? row_rescale : 1.0f);
     reduced_row_ptr[tid] = thread_output;
   }
 }
@@ -770,6 +772,7 @@ void finalize_moe_routing_kernelLauncher(
     const int cols,
     const int k,
     const int compute_bias,
+    const bool norm_topk_prob,
     cudaStream_t stream) {
   const int blocks = num_rows;
   const int threads = std::min(cols, 1024);
@@ -783,7 +786,8 @@ void finalize_moe_routing_kernelLauncher(
                                        expert_for_source_row,
                                        cols,
                                        k,
-                                       compute_bias);
+                                       compute_bias,
+                                       norm_topk_prob);
 }
 
 // ========================= TopK Softmax specializations
@@ -863,6 +867,7 @@ template void finalize_moe_routing_kernelLauncher(const float*,
                                                   const int,
                                                   const int,
                                                   const int,
+                                                  const bool,
                                                   cudaStream_t);
 template void finalize_moe_routing_kernelLauncher(const half*,
                                                   half*,
@@ -874,6 +879,7 @@ template void finalize_moe_routing_kernelLauncher(const half*,
                                                   const int,
                                                   const int,
                                                   const int,
+                                                  const bool,
                                                   cudaStream_t);
 #ifdef PADDLE_CUDA_BF16
 template void finalize_moe_routing_kernelLauncher(const __nv_bfloat16*,
@@ -886,6 +892,7 @@ template void finalize_moe_routing_kernelLauncher(const __nv_bfloat16*,
                                                   const int,
                                                   const int,
                                                   const int,
+                                                  const bool,
                                                   cudaStream_t);
 #endif
 template void compute_total_rows_before_expert(
diff --git a/paddle/phi/ops/yaml/fused_ops.yaml b/paddle/phi/ops/yaml/fused_ops.yaml
@@ -842,7 +842,7 @@
 
 - op: fused_moe
   args: (Tensor x, Tensor gate_weight, Tensor ffn1_weight, Tensor ffn1_scale, Tensor ffn1_bias, Tensor ffn2_weight, Tensor ffn2_scale, Tensor ffn2_bias,
-    str quant_method = "None", int moe_topk = 2)
+    str quant_method = "None", int moe_topk = 2, bool norm_topk_prob = true)
   output: Tensor (out)
   infer_meta:
     func: FusedMoeInferMeta
diff --git a/python/paddle/incubate/nn/functional/fused_moe.py b/python/paddle/incubate/nn/functional/fused_moe.py
@@ -30,6 +30,7 @@ def fused_moe(
     ffn2_scale=None,
     quant_method="None",
     moe_topk=2,
+    norm_topk_prob=True,
 ):
     """
     Applies fused moe kernel.
@@ -45,7 +46,8 @@ def fused_moe(
         ffn1_scale (Tensor, optional): the input scale Tensor Provided to weight for dequantization. Its shape is [num_experts, d_feed_forward*2].
         ffn2_scale (Tensor, optional): the input scale Tensor Provided to weight for dequantization. Its shape is [num_experts, d_model].
         quant_method (string): Currently not supported.
-        moe_topk: Select the top k experts for each token.
+        moe_topk (int): Select the top k experts for each token.
+        norm_topk_prob (bool): Whether to normalize the topk probabilities.
 
     Returns:
         Tensor: the output Tensor.
@@ -66,7 +68,7 @@ def fused_moe(
             >>> ffn2_weight = paddle.randn([8, 2048, 1024])
             >>> ffn2_bias = paddle.randn([8, 1, 1024])
             >>> moe_topk = 2
-            >>> out = fused_moe(x, gate_weight, ffn1_weight, ffn1_bias, ffn2_weight, ffn2_bias, None, None,"None", moe_topk)
+            >>> out = fused_moe(x, gate_weight, ffn1_weight, ffn1_bias, ffn2_weight, ffn2_bias, None, None,"None", moe_topk, True)
             >>> print(out.shape)
             [10, 128, 1024]
 
@@ -83,6 +85,7 @@ def fused_moe(
             ffn2_bias,
             quant_method,
             moe_topk,
+            norm_topk_prob,
         )
         return final_out
     else:
@@ -104,6 +107,7 @@ def fused_moe(
             attrs={
                 'quant_method': quant_method,
                 'moe_topk': moe_topk,
+                'norm_topk_prob': norm_topk_prob,
             },
         )
         return final_out
diff --git a/test/legacy_test/test_fused_moe_op.py b/test/legacy_test/test_fused_moe_op.py
@@ -119,6 +119,7 @@ def config(self):
         self.d_feedforward = 3072
         self.top_k = 2
         self.quant_method = "None"
+        self.norm_topk_prob = True
 
     def GetWintData(self):
         if self.quant_method == "None":
@@ -168,6 +169,7 @@ def GetFusedMoeOut(self, tensor_x):
             None if self.quant_method == "None" else self.scale1,
             self.quant_method,
             self.top_k,
+            self.norm_topk_prob,
         )
 
         return fused_out
@@ -182,7 +184,11 @@ def GetBaselineOut(self, hidden_states):
         routing_weights, selected_experts = paddle.topk(
             routing_weights, self.top_k, axis=-1
         )
-        routing_weights /= paddle.sum(routing_weights, axis=-1, keepdim=True)
+        # mixtral true, qwen_moe false
+        if self.norm_topk_prob:
+            routing_weights /= paddle.sum(
+                routing_weights, axis=-1, keepdim=True
+            )
         # we cast back to the input dtype
         routing_weights = routing_weights.cast(np.float32)
 
@@ -258,5 +264,19 @@ def config(self):
         self.quant_method = "weight_only_int8"
 
 
+@unittest.skipIf(
+    not paddle.is_compiled_with_cuda()
+    or get_cuda_version() < 11030
+    or paddle.device.cuda.get_device_capability()[0] < 8,
+    "FusedMoe requires CUDA >= 11.2 and CUDA_ARCH >= 8",
+)
+class TestFusedMoEOpNonNorm(TestFusedMoEOp):
+    def config(self):
+        super().config()
+        self.rtol = 1e-2
+        self.atol = 1e-2
+        self.norm_topk_prob = False
+
+
 if __name__ == "__main__":
     unittest.main()
