[Bugfix] fix fused moe wint gemm bug (#67396)

penPenf28 · web-flow · commit 8e3be110eead · 2024-08-14T15:51:17.000+08:00
diff --git a/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h b/paddle/phi/kernels/fusion/cutlass/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h
@@ -677,6 +677,48 @@ void dispatch_moe_gemm_to_cutlass(const T* A,
           stream,
           occupancy);
       break;
+    case CutlassTileConfig::CtaShape64x128x64_WarpShape64x64x64:
+      dispatch_gemm_config<T,
+                           WeightType,
+                           arch,
+                           EpilogueTag,
+                           cutlass::gemm::GemmShape<64, 128, 64>,
+                           cutlass::gemm::GemmShape<64, 64, 64>>(
+          A,
+          B,
+          weight_scales,
+          biases,
+          C,
+          total_rows_before_expert,
+          gemm_n,
+          gemm_k,
+          num_experts,
+          gemm_config,
+          multi_processor_count,
+          stream,
+          occupancy);
+      break;
+    case CutlassTileConfig::CtaShape128x128x64_WarpShape64x64x64:
+      dispatch_gemm_config<T,
+                           WeightType,
+                           arch,
+                           EpilogueTag,
+                           cutlass::gemm::GemmShape<128, 128, 64>,
+                           cutlass::gemm::GemmShape<64, 64, 64>>(
+          A,
+          B,
+          weight_scales,
+          biases,
+          C,
+          total_rows_before_expert,
+          gemm_n,
+          gemm_k,
+          num_experts,
+          gemm_config,
+          multi_processor_count,
+          stream,
+          occupancy);
+      break;
     case CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64:
       dispatch_gemm_config<T,
                            WeightType,
diff --git a/python/paddle/incubate/nn/functional/fused_moe.py b/python/paddle/incubate/nn/functional/fused_moe.py
@@ -42,8 +42,8 @@ def fused_moe(
         ffn1_bias (Tensor): the first batch matrix matmul bias. Its shape is [num_experts, 1, d_feed_forward*2].
         ffn2_weight (Tensor): the second batch matrix matmul weight. Its shape is [num_experts, d_feed_forward, d_model].
         ffn2_bias (Tensor): the second batch matrix matmul bias. Its shape is [num_experts, 1, d_model].
-        ffn1_scale (Tensor, optional): the input scale Tensor Provided to weight for dequantization. Its shape is [TODO: xinhw].
-        ffn2_scale (Tensor, optional): the input scale Tensor Provided to weight for dequantization. Its shape is [TODO: xinhw].
+        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.
 
diff --git a/test/legacy_test/test_fused_moe_op.py b/test/legacy_test/test_fused_moe_op.py
@@ -23,6 +23,7 @@
 from paddle import nn
 from paddle.incubate.nn.functional import fused_moe, swiglu
 from paddle.nn.layer.common import Linear
+from paddle.nn.quant import weight_quantize
 
 paddle.seed(42)
 
@@ -119,7 +120,42 @@ def config(self):
         self.top_k = 2
         self.quant_method = "None"
 
+    def GetWintData(self):
+        if self.quant_method == "None":
+            return
+        fc0_expert_weights_for_ref_list = []
+        scale0 = []
+        for i in range(self.num_expert):
+            fc0_expert_weights_for_ref_i, fc0_expert_weights_scale_for_ref_i = (
+                weight_quantize(self.bmm_w0[i], algo=self.quant_method)
+            )
+            fc0_expert_weights_for_ref_list.append(
+                fc0_expert_weights_for_ref_i.reshape(
+                    [self.d_model, self.d_feedforward * 2]
+                )
+            )
+            scale0.append(fc0_expert_weights_scale_for_ref_i)
+        self.bmm_w0 = paddle.to_tensor(fc0_expert_weights_for_ref_list)
+        self.scale0 = paddle.to_tensor(scale0)
+
+        fc1_expert_weights_for_ref_list = []
+        scale1 = []
+        for i in range(self.num_expert):
+            fc1_expert_weights_for_ref_i, fc1_expert_weights_scale_for_ref_i = (
+                weight_quantize(self.bmm_w1[i], algo=self.quant_method)
+            )
+            fc1_expert_weights_for_ref_list.append(
+                fc1_expert_weights_for_ref_i.reshape(
+                    [self.d_feedforward, self.d_model]
+                )
+            )
+            scale1.append(fc1_expert_weights_scale_for_ref_i)
+        self.bmm_w1 = paddle.to_tensor(fc1_expert_weights_for_ref_list)
+        self.scale1 = paddle.to_tensor(scale1)
+
     def GetFusedMoeOut(self, tensor_x):
+        if self.quant_method != "None":
+            self.GetWintData()
         paddle.disable_static(place=paddle.CUDAPlace(0))
         fused_out = fused_moe(
             tensor_x,
@@ -194,6 +230,12 @@ def test_fused_moe_op_new(self):
         )
 
 
+@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 TestFusedMoEOpBf16(TestFusedMoEOp):
     def config(self):
         super().config()
@@ -202,5 +244,19 @@ def config(self):
         self.atol = 1e-2
 
 
+@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 TestFusedMoEOpWint8(TestFusedMoEOp):
+    def config(self):
+        super().config()
+        self.rtol = 1e-2
+        self.atol = 1e-2
+        self.quant_method = "weight_only_int8"
+
+
 if __name__ == "__main__":
     unittest.main()
