[Bugfix] fix fused moe wint4 (#67717)

penPenf28 · web-flow · commit 98e45971d5d4 · 2024-08-26T20:30:36.000-07:00
* [Bugfix] fix fused moe wint4

* [Refactor] revert skip

* [Bugfix] fix paddlenlp wint4 static export
diff --git a/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_helper.h b/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_helper.h
@@ -161,6 +161,10 @@ class MoeHelper {
       inter_dim = ffn1_dims[2];
     }
 
+    if (gemm_method_ == "weight_only_int4") {
+      inter_dim = inter_dim * 2;
+    }
+
     const int inter_size = inter_dim;
     const int num_experts = ffn1_dims[0];
     const int k = moe_topk;
@@ -376,7 +380,7 @@ class MoeHelper {
             inter_size / 2,
             num_experts,
             ctx.stream());
-      } else if (gemm_method_ == "weight_onlyint4") {
+      } else if (gemm_method_ == "weight_only_int4") {
         int4_moe_gemm_runner_->moe_gemm(
             reinterpret_cast<NvType *>(act_out),
             reinterpret_cast<const cutlass::uint4b_t *>(
diff --git a/python/paddle/base/data_feeder.py b/python/paddle/base/data_feeder.py
@@ -136,6 +136,7 @@ def convert_dtype(dtype: DTypeLike) -> _DTypeLiteral:
             'uint16',
             'float32',
             'float64',
+            'int4',
             'int8',
             'int16',
             'int32',
diff --git a/test/legacy_test/test_fused_moe_op.py b/test/legacy_test/test_fused_moe_op.py
@@ -136,6 +136,8 @@ def GetWintData(self):
             fc0_expert_weights_for_ref_list.append(
                 fc0_expert_weights_for_ref_i.reshape(
                     [self.d_model, self.d_feedforward * 2]
+                    if self.quant_method == "weight_only_int8"
+                    else [self.d_model, self.d_feedforward]
                 )
             )
             scale0.append(fc0_expert_weights_scale_for_ref_i)
@@ -151,6 +153,8 @@ def GetWintData(self):
             fc1_expert_weights_for_ref_list.append(
                 fc1_expert_weights_for_ref_i.reshape(
                     [self.d_feedforward, self.d_model]
+                    if self.quant_method == "weight_only_int8"
+                    else [self.d_feedforward, self.d_model // 2]
                 )
             )
             scale1.append(fc1_expert_weights_scale_for_ref_i)
@@ -267,6 +271,20 @@ 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 TestFusedMoEOpWint4(TestFusedMoEOp):
+    def config(self):
+        super().config()
+        self.rtol = 1e-2
+        self.atol = 1e-2
+        self.quant_method = "weight_only_int4"
+
+
 @unittest.skipIf(
     not paddle.is_compiled_with_cuda()
     or get_cuda_version() < 11030
