[Bugfix] fix static graph exe and make bias optional (#67580)

penPenf28 · web-flow · commit cc5dab49c988 · 2024-08-21T14:15:04.000+08:00
diff --git a/paddle/phi/kernels/fusion/cutlass/fused_moe_kernel.cu b/paddle/phi/kernels/fusion/cutlass/fused_moe_kernel.cu
@@ -50,10 +50,10 @@ void FusedMoeKernel(const Context& ctx,
                     const DenseTensor& gate_weight,
                     const DenseTensor& ffn1_weight,
                     const paddle::optional<DenseTensor>& ffn1_scale,
-                    const DenseTensor& ffn1_bias,
+                    const paddle::optional<DenseTensor>& ffn1_bias,
                     const DenseTensor& ffn2_weight,
                     const paddle::optional<DenseTensor>& ffn2_scale,
-                    const DenseTensor& ffn2_bias,
+                    const paddle::optional<DenseTensor>& ffn2_bias,
                     const std::string& quant_method,
                     const int moe_topk,
                     const bool norm_topk_prob,
@@ -80,10 +80,10 @@ void FusedMoeKernel(const Context& ctx,
                          &gate_weight,
                          &ffn1_weight,
                          ffn1_scale,
-                         &ffn1_bias,
+                         ffn1_bias,
                          &ffn2_weight,
                          ffn2_scale,
-                         &ffn2_bias,
+                         ffn2_bias,
                          nullptr,
                          moe_topk,
                          norm_topk_prob,
diff --git a/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_helper.h b/paddle/phi/kernels/fusion/cutlass/moe/fused_moe_helper.h
@@ -126,19 +126,19 @@ class MoeHelper {
                   const DenseTensor *gate_weight,
                   const DenseTensor *ffn1_weight,
                   const paddle::optional<DenseTensor> &ffn1_scale,
-                  const DenseTensor *ffn1_bias,
+                  const paddle::optional<DenseTensor> &ffn1_bias,
                   const DenseTensor *ffn2_weight,
                   const paddle::optional<DenseTensor> &ffn2_scale,
-                  const DenseTensor *ffn2_bias,
+                  const paddle::optional<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>();
     auto *gating_weights = gate_weight->data<float>();
-    auto *fc1_expert_biases = ffn1_bias->data<T>();
-    auto *fc2_expert_biases = ffn2_bias->data<T>();
+    auto *fc1_expert_biases = ffn1_bias.get_ptr()->data<T>();
+    auto *fc2_expert_biases = ffn2_bias.get_ptr()->data<T>();
 
     auto *output_ = output->data<T>();
     auto stream = ctx.stream();
diff --git a/paddle/phi/ops/yaml/fused_ops.yaml b/paddle/phi/ops/yaml/fused_ops.yaml
@@ -849,5 +849,5 @@
   kernel:
     func: fused_moe
     data_type : x
-  optional : ffn1_scale, ffn2_scale
+  optional : ffn1_bias, ffn1_scale, ffn2_bias, ffn2_scale
   support_dygraph_mode : true
diff --git a/python/paddle/incubate/nn/functional/fused_moe.py b/python/paddle/incubate/nn/functional/fused_moe.py
@@ -23,10 +23,10 @@ def fused_moe(
     x,
     gate_weight,
     ffn1_weight,
-    ffn1_bias,
     ffn2_weight,
-    ffn2_bias,
+    ffn1_bias=None,
     ffn1_scale=None,
+    ffn2_bias=None,
     ffn2_scale=None,
     quant_method="None",
     moe_topk=2,
@@ -40,10 +40,10 @@ def fused_moe(
         x (Tensor): the input Tensor. Its shape is [bsz, seq_len, d_model].
         gate_weight (Tensor): the gate Tensor to choose expert. Its shape is [bsz, seq_len, num_experts].
         ffn1_weight (Tensor): the first batch matrix matmul weight. Its shape is [num_experts, d_model, d_feed_forward*2].
-        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_bias (Tensor, optional): the first batch matrix matmul bias. Its shape is [num_experts, 1, d_feed_forward*2].
         ffn1_scale (Tensor, optional): the input scale Tensor Provided to weight for dequantization. Its shape is [num_experts, d_feed_forward*2].
+        ffn2_bias (Tensor, optional): the second batch matrix matmul bias. Its shape is [num_experts, 1, d_model].
         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 (int): Select the top k experts for each token.
@@ -68,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, True)
+            >>> out = fused_moe(x, gate_weight, ffn1_weight, ffn2_weight, ffn1_bias, None, ffn2_bias, None, "None", moe_topk, True)
             >>> print(out.shape)
             [10, 128, 1024]
 
@@ -91,18 +91,25 @@ def fused_moe(
     else:
         helper = LayerHelper('fused_moe', **locals())
         final_out = helper.create_variable_for_type_inference(dtype=x.dtype)
+
+        inputs = {
+            'x': x,
+            'gate_weight': gate_weight,
+            'ffn1_weight': ffn1_weight,
+            'ffn2_weight': ffn2_weight,
+        }
+        if ffn1_bias is not None:
+            inputs['ffn1_bias'] = ffn1_bias
+        if ffn1_scale is not None:
+            inputs['ffn1_scale'] = ffn1_scale
+        if ffn2_bias is not None:
+            inputs['ffn2_bias'] = ffn2_bias
+        if ffn2_scale is not None:
+            inputs['ffn2_scale'] = ffn2_scale
+
         helper.append_op(
             type='fused_moe',
-            inputs={
-                'x': x,
-                'gate_weight': gate_weight,
-                'ffn1_weight': ffn1_weight,
-                'ffn1_scale': ffn1_scale,
-                'ffn1_bias': ffn1_bias,
-                'ffn2_weight': ffn2_weight,
-                'ffn2_scale': ffn2_scale,
-                'ffn2_bias': ffn2_bias,
-            },
+            inputs=inputs,
             outputs={'out': final_out},
             attrs={
                 'quant_method': quant_method,
diff --git a/test/legacy_test/test_fused_moe_op.py b/test/legacy_test/test_fused_moe_op.py
