change api

Author: zhangbo9674

paddlenlp/transformers/fp8_utils.py

@@ -27,7 +27,6 @@ def swiglu(x, y=None):
 
 try:
     import deep_gemm
-    import FusedQuantOps as FQO
     import kitchen
     import kitchen.quantization_subchannel_block_hybrid
     from kitchen.quantization import QParams, ScalingType
@@ -343,7 +342,7 @@ def backward(ctx, do3):
         o1 = paddle.empty([x_fp8.shape[0], w1_fp8.shape[0]], dtype=do3.dtype)
         deep_gemm.gemm_fp8_fp8_bf16_nt((x_fp8, x_scale.T), (w1_fp8, w1_sacle), o1)
 
-        x_dequant_fp16 = FQO.fused_act_dequant(x_fp8, x_scale.T.contiguous())
+        x_dequant_fp16 = paddle.incubate.nn.functional.fused_act_dequant(x_fp8, x_scale.T.contiguous())
         x_dequant_fp16 = padding(x_dequant_fp16, 0)
 
         _, _, x_t_fp8, x_t_scale = kitchen_quant(
@@ -468,7 +467,7 @@ def fwd_gate_up(self, x_bf16, expert_w1, num_expert, tokens_per_expert):
         self.tokens_per_expert = tokens_per_expert
         self.m_indices = gen_m_indices(tokens_per_expert)
         # concat w1, shape is [num_groups, n, k]
-        w1_t_quant, w1_t_scale = FQO.fused_stack_transpose_quant(expert_w1)
+        w1_t_quant, w1_t_scale = paddle.incubate.nn.functional.fused_stack_transpose_quant(expert_w1, transpose=True)
         w1_t_quant = w1_t_quant.reshape([num_expert, -1, w1_t_quant.shape[-1]])
         w1_t_scale = w1_t_scale.reshape([num_expert, -1, w1_t_scale.shape[-1]])
 
@@ -504,12 +503,14 @@ def fwd_down(self, o1, unzipped_probs, expert_w2, num_expert):
         [m_sum, k] = [m_sum, n] * [num_groups, n, k]
         """
         # concat and transpose w2
-        w2_quant, w2_sacle = FQO.fused_stack_transpose_quant(expert_w2)
+        w2_quant, w2_sacle = paddle.incubate.nn.functional.fused_stack_transpose_quant(expert_w2, transpose=True)
         w2_quant = w2_quant.reshape([num_expert, -1, w2_quant.shape[-1]])
         w2_sacle = w2_sacle.reshape([num_expert, -1, w2_sacle.shape[-1]])
 
         # quant o2
-        o2_fp8, o2_scale = FQO.fused_spaq(o1, unzipped_probs, using_pow2_scaling=True)
+        o2_fp8, o2_scale = paddle.incubate.nn.functional.fused_weighted_swiglu_act_quant(
+            o1, unzipped_probs, using_pow2_scaling=True
+        )
         o2_scale = paddle.transpose(paddle.transpose(o2_scale, [1, 0]).contiguous(), [1, 0])
         self.unzipped_probs = unzipped_probs
 
@@ -527,7 +528,9 @@ def bwd_dowm_input(self, expert_w2, unzipped_grad, o1):
         [m_sum, n] = [m_sum, k] * [num_groups, k, n]
         """
         # recompute concated_w2_2d
-        bw_w2_quant, bw_w2_scale = FQO.fused_stack_quant(expert_w2)
+        bw_w2_quant, bw_w2_scale = paddle.incubate.nn.functional.fused_stack_transpose_quant(
+            expert_w2, transpose=False
+        )
         bw_w2_quant = bw_w2_quant.reshape([len(expert_w2), -1, bw_w2_quant.shape[-1]])
         bw_w2_scale = bw_w2_scale.reshape([len(expert_w2), -1, bw_w2_scale.shape[-1]])
 
@@ -541,7 +544,7 @@ def bwd_dowm_input(self, expert_w2, unzipped_grad, o1):
                 (unzipped_grad_fp8, unzipped_grad_scale), (bw_w2_quant, bw_w2_scale), do2_s, m_indices=self.m_indices
             )
 
-        do1, probs_grad, o2_s = FQO.fused_swiglu_probs_bwd(o1, do2_s, self.unzipped_probs)
+        do1, probs_grad, o2_s = paddle.incubate.nn.functional.fused_swiglu_weighted_bwd(o1, do2_s, self.unzipped_probs)
 
         return do1, o2_s, probs_grad
 
@@ -555,7 +558,9 @@ def bwd_gate_up_input(self, do1, expert_w1):
         [m_sum, k] = [m_sum, n] * [num_groups, n, k]
         """
         # recompute concated_w1_t
-        bw_w1_quant, bw_w1_scale = FQO.fused_stack_quant(expert_w1)
+        bw_w1_quant, bw_w1_scale = paddle.incubate.nn.functional.fused_stack_transpose_quant(
+            expert_w1, transpose=False
+        )
         bw_w1_quant = bw_w1_quant.reshape([len(expert_w1), -1, bw_w1_quant.shape[-1]])
         bw_w1_scale = bw_w1_scale.reshape([len(expert_w1), -1, bw_w1_scale.shape[-1]])
 
@@ -573,11 +578,7 @@ def bwd_gate_up_input(self, do1, expert_w1):
         return dx
 
     def fused_transpose_split_quant(self, x, tokens_per_expert, pow_2_scales):
-        out, scale = [], []
-        for tokens in tokens_per_expert:
-            out.append(paddle.empty([x.shape[1], tokens], dtype="float8_e4m3fn"))
-            scale.append(paddle.empty([tokens // 128, x.shape[1]], dtype="float32"))
-        FQO.fused_transpose_split_quant(x, out, scale, pow_2_scales)
+        out, scale = paddle.incubate.nn.functional.fused_transpose_split_quant(x, tokens_per_expert, pow_2_scales)
         return out, scale
 
     def bwd_down_weight(self, do3, o2, expert_w2):
@@ -681,7 +682,7 @@ def backward(self, out_grad):
         expert_w2 = [x.w2 for x in self.custom_map.experts if x is not None]
 
         if self.mem_efficient:
-            input = FQO.fused_act_dequant(self.input_fp8, self.input_scale)
+            input = paddle.incubate.nn.functional.fused_act_dequant(self.input_fp8, self.input_scale)
         else:
             input = self.input
 

paddlenlp/transformers/moe_utils.py

@@ -18,11 +18,6 @@
 
 from .fp8_utils import dequantize_fp8_to_fp32
 
-try:
-    import TokenDispatcherUtils as TDU
-except:
-    pass
-
 
 def topk_to_permuted_indices(x, num_tokens_per_expert_list, topk):
     x = paddle.flatten(x)
@@ -120,12 +115,11 @@ def forward(
         num_experts,
         tokens_per_expert,
     ):
-        (unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, _,) = TDU.tokens_unzip_stable(
+        (unzipped_tokens, zipped_expertwise_rowmap, unzipped_probs, _,) = paddle.nn.functional.moe_permute(
             hs_2d_dispatched,
             None,
             dispatched_indices,
             dispatched_probs,
-            topk=topk,
             num_experts=num_experts,
             tokens_per_expert=tokens_per_expert,
             padding_multiplex=128,
@@ -140,7 +134,7 @@ def forward(
 
     @paddle.no_grad()
     def backward(self, dx, hidden_states_out_grad, probs_grad, dispatched_indices, num_experts):
-        weighted_zipped_tokens, probs_grad_zipped = TDU.tokens_zip(
+        weighted_zipped_tokens, probs_grad_zipped = paddle.nn.functional.moe_unpermute(
             dx,
             self.zipped_expertwise_rowmap,
             dispatched_indices,
@@ -161,7 +155,7 @@ def __init__(self, token_dispatcher, name="zip"):
     def forward(
         self, expert_out, zipped_expertwise_rowmap, routemap_topk, unzipped_probs, total_zipped_tokens, num_experts
     ):
-        expert_out_zipped, zipped_probs_topk = TDU.tokens_zip(
+        expert_out_zipped, zipped_probs_topk = paddle.nn.functional.moe_unpermute(
             expert_out, zipped_expertwise_rowmap, routemap_topk, unzipped_probs, total_zipped_tokens, num_experts
         )
         return expert_out_zipped
@@ -176,12 +170,11 @@ def backward(
         num_experts,
         tokens_per_expert,
     ):
-        (unzipped_grad, zipped_expertwise_rowmap_grad, unzipped_probs_grad, _,) = TDU.tokens_unzip_stable(
+        (unzipped_grad, zipped_expertwise_rowmap_grad, unzipped_probs_grad, _,) = paddle.nn.functional.moe_permute(
             grad_output,
             None,
             dispatched_indices,
             dispatched_probs,
-            top_k,
             num_experts,
             tokens_per_expert,
             padding_multiplex=128,