From 3b3a591640c64fa46fd850be9087b018be1d1011 Mon Sep 17 00:00:00 2001 From: phlrain Date: Wed, 9 Apr 2025 15:33:11 +0800 Subject: [PATCH] optmize fp8 expert gemm --- .../transformers/deepseek_v2/fp8_linear.py | 10 +- .../transformers/deepseek_v2/modeling.py | 34 ++- paddlenlp/transformers/fp8_utils.py | 238 +++++++++--------- paddlenlp/transformers/moe_layer.py | 6 + 4 files changed, 154 insertions(+), 134 deletions(-) diff --git a/paddlenlp/transformers/deepseek_v2/fp8_linear.py b/paddlenlp/transformers/deepseek_v2/fp8_linear.py index d3ed61ddac34..767383e56005 100644 --- a/paddlenlp/transformers/deepseek_v2/fp8_linear.py +++ b/paddlenlp/transformers/deepseek_v2/fp8_linear.py @@ -348,6 +348,8 @@ def backward(ctx, dout): w_quant, w_scale = kitchen_quant( weight, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=False ) + + #print("shape", dout_quant.shape, w_quant.shape, dx.shape) deep_gemm.gemm_fp8_fp8_bf16_nt((dout_quant, dout_scale), (w_quant, w_scale), dx) dx = dx.reshape(dx_orig_shape) @@ -416,11 +418,9 @@ def forward(ctx, x, w1, w2): x_fp8, x_scale = kitchen_quant( x, backend=kitchen.ops.Backend.CUTLASS, is_1d_scaled=True, return_transpose=False ) - - w_t = w1.T.contiguous() - - w1_fp8, w1_sacle = kitchen_quant( - w_t, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=False + + _, _, w1_fp8, w1_sacle = kitchen_quant( + w1, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=True ) o1 = paddle.empty([x_fp8.shape[0], w1_fp8.shape[0]], dtype=x.dtype) deep_gemm.gemm_fp8_fp8_bf16_nt((x_fp8, x_scale), (w1_fp8, w1_sacle), o1) diff --git a/paddlenlp/transformers/deepseek_v2/modeling.py b/paddlenlp/transformers/deepseek_v2/modeling.py index 1d3e8e922bcb..e8749efed4f2 100644 --- a/paddlenlp/transformers/deepseek_v2/modeling.py +++ b/paddlenlp/transformers/deepseek_v2/modeling.py @@ -82,10 +82,11 @@ from ..utils import device_guard from . import fp8_linear as linear_utils from .configuration import DeepseekV2Config -from .fp8_linear import FP8DeepseekV2MLP, FP8KeepXLinear, FP8Linear, Linear +from .fp8_linear import FP8DeepseekV2MLP, FP8KeepXLinear, FP8Linear, Linear, FusedFP8DeepseekV2MLP DSV3_USE_FP8_GEMM = os.getenv("DSV3_USE_FP8_GEMM", "False").lower() == "true" DSV3_USE_ATTEN_RECOMPUTE = os.getenv("DSV3_USE_ATTEN_RECOMPUTE", "False").lower() == "true" +DSV3_USE_FUSED_Expert = os.getenv("DSV3_USE_FUSED_Expert", "False").lower() == "true" FA_VERSION = int(os.getenv("FA_VERSION", 2)) @@ -852,6 +853,9 @@ def __init__(self, config: DeepseekV2Config): drop_tokens=False, ) DeepseekV2MLPClass = FP8DeepseekV2MLP if DSV3_USE_FP8_GEMM else DeepseekV2MLP + + if DSV3_USE_FUSED_Expert: + DeepseekV2MLPClass = FusedFP8DeepseekV2MLP super().__init__( config=config, @@ -874,7 +878,12 @@ def __init__(self, config: DeepseekV2Config): self.alpha = config.aux_loss_alpha if config.n_shared_experts is not None: intermediate_size = config.moe_intermediate_size * config.n_shared_experts - self.shared_experts = DeepseekV2MLPClass(config=config, intermediate_size=intermediate_size, is_moe=False) + if DSV3_USE_FP8_GEMM: + self.shared_experts = FP8DeepseekV2MLP( + config=config, intermediate_size=intermediate_size, is_moe=False + ) + else: + self.shared_experts = DeepseekV2MLP(config=config, intermediate_size=intermediate_size, is_moe=False) def forward(self, hidden_states): final_hidden_states, l_aux, l_zloss = super().forward(hidden_states) @@ -1063,9 +1072,24 @@ def forward( ) elif FA_VERSION == 3: - attn_out, softmax_lse = flash_attn_v3( - query_states, key_states, value_states, softmax_scale=softmax_scale, causal=True - ) + attn_out, softmax_lse = _C_ops.flash_attn_v3( + query_states, + key_states, + value_states, + None, # q_v_ + None, # q_descale_ + None, # k_descale_ + None, # v_descale_ + softmax_scale, + True, + -1, # window_size_left + -1, # window_size_right + 0.0, # softcap + 1, # num_splits + False, # manual_set_pack_gqa + False, # pack_gqa_ + 0, # sm_margin + ) else: assert False, f"invalid {FA_VERSION=}" diff --git a/paddlenlp/transformers/fp8_utils.py b/paddlenlp/transformers/fp8_utils.py index cf0d8e18e1f6..eed83673d8e9 100644 --- a/paddlenlp/transformers/fp8_utils.py +++ b/paddlenlp/transformers/fp8_utils.py @@ -47,6 +47,7 @@ def swiglu(x, y=None): ] IF_USE_GROUP_GEMM_MASK = os.getenv("IF_USE_GROUP_GEMM_MASK", "False").lower() == "true" +DSV3_USE_FUSED_Expert = os.getenv("DSV3_USE_FUSED_Expert", "False").lower() == "true" def kitchen_quant(x, backend=None, is_1d_scaled=True, return_transpose=False): @@ -124,23 +125,19 @@ def reset_statue(self): self.unzipped_tokens = None self.unzipped_probs = None self.tokens_per_expert = None + self.custom_map = None def fwd_gate_up(self, x_fp8, x_scale, expert_w1, expert_w_count, tokens_per_expert): - # concat w1 - stacked_w1 = paddle.stack(expert_w1, axis=0) - stacked_w1_t = paddle.transpose(stacked_w1, [0, 2, 1]).contiguous() - concated_w1_t = stacked_w1_t.reshape([-1, stacked_w1_t.shape[-1]]) - # quant w1 w1_t_quant, w1_t_scale = kitchen_quant( - concated_w1_t, + expert_w1, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=False, ) - w1_t_quant = w1_t_quant.reshape([expert_w_count, -1, w1_t_quant.shape[-1]]) - w1_t_scale = w1_t_scale.reshape([expert_w_count, -1, w1_t_scale.shape[-1]]) + w1_t_quant = w1_t_quant.reshape([expert_w_count, -1, w1_t_quant.shape[-1]]).transpose([0, 2, 1]).contiguous() + w1_t_scale = w1_t_scale.reshape([expert_w_count, -1, w1_t_scale.shape[-1]]).transpose([0, 2, 1]).contiguous() # mask group gemm需要输入x是[group,m,n] x_fp8 = x_fp8.reshape([expert_w_count, -1, x_fp8.shape[-1]]) @@ -177,19 +174,12 @@ def fwd_swiglu(self, o1): return o2 def fwd_down(self, o2, expert_w2, expert_w_count, tokens_per_expert): - # concat and transpose w2 - expert_w2 = [x.w2 for x in self.custom_map.experts if x is not None] - - stacked_w2 = paddle.stack(expert_w2, axis=0) - stacked_w2_t = paddle.transpose(stacked_w2, [0, 2, 1]).contiguous() - concated_w2_t = stacked_w2_t.reshape([-1, stacked_w2_t.shape[-1]]) - # quant w2 w2_quant, w2_sacle = kitchen_quant( - concated_w2_t, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=False + expert_w2, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=False ) - w2_quant = w2_quant.reshape([expert_w_count, -1, w2_quant.shape[-1]]) - w2_sacle = w2_sacle.reshape([expert_w_count, -1, w2_sacle.shape[-1]]) + w2_quant = w2_quant.reshape([expert_w_count, -1, w2_quant.shape[-1]]).transpose([0, 2, 1]).contiguous() + w2_sacle = w2_sacle.reshape([expert_w_count, -1, w2_sacle.shape[-1]]).transpose([0, 2, 1]).contiguous() # quant o2 o2_reshape = o2.reshape([-1, o2.shape[-1]]) @@ -229,25 +219,22 @@ def fwd_down(self, o2, expert_w2, expert_w_count, tokens_per_expert): return o3 # ===== do2 = deep_gemm(do3_fp8, w2_fp8) - def bwd_dowm_input(self, expert_w2, unzipped_grad, unzipped_scale, tokens_per_expert, expected_m): - # recompute concated_w2_2d - stacked_w2 = paddle.stack(expert_w2, axis=0) - concated_w2 = stacked_w2.reshape([-1, stacked_w2.shape[-1]]) - + def bwd_dowm_input(self, expert_w2, expert_w_count, unzipped_grad, unzipped_scale, tokens_per_expert, expected_m): # quant w2 bw_w2_quant, bw_w2_scale = kitchen_quant( - concated_w2, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=False + expert_w2, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_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]]) + + bw_w2_quant = bw_w2_quant.reshape([expert_w_count, -1, bw_w2_quant.shape[-1]]) + bw_w2_scale = bw_w2_scale.reshape([expert_w_count, -1, bw_w2_scale.shape[-1]]) # do2 - unzipped_grad = unzipped_grad.reshape([len(expert_w2), -1, unzipped_grad.shape[-1]]) - unzipped_scale = unzipped_scale.reshape([len(expert_w2), -1, unzipped_scale.shape[-1]]) + unzipped_grad = unzipped_grad.reshape([expert_w_count, -1, unzipped_grad.shape[-1]]) + unzipped_scale = unzipped_scale.reshape([expert_w_count, -1, unzipped_scale.shape[-1]]) # do2 = paddle.empty([len(expert_w2), unzipped_grad.shape[1], bw_w2_quant.shape[1]], dtype="bfloat16") if IF_USE_GROUP_GEMM_MASK: - do2 = paddle.zeros([len(expert_w2), unzipped_grad.shape[1], bw_w2_quant.shape[1]], dtype="bfloat16") + do2 = paddle.zeros([expert_w_count, unzipped_grad.shape[1], bw_w2_quant.shape[1]], dtype="bfloat16") deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked( (unzipped_grad, unzipped_scale), (bw_w2_quant, bw_w2_scale), @@ -309,30 +296,26 @@ def bwd_swiglu(self, o1, do2): return do1 # ===== dx = deep_gemm(do1_fp8, w1_fp8) - - def bwd_gate_up_input(self, do1, expert_w1, tokens_per_expert, expected_m): - # recompute concated_w1_t - stacked_w1 = paddle.stack(expert_w1, axis=0) - concated_w1_t_2d = stacked_w1.reshape([-1, stacked_w1.shape[-1]]) - + def bwd_gate_up_input(self, do1, expert_w1, expert_w_count, tokens_per_expert, expected_m): # quant w1 bw_w1_quant, bw_w1_scale = kitchen_quant( - concated_w1_t_2d, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_transpose=False + expert_w1, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=False, return_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]]) + bw_w1_quant = bw_w1_quant.reshape([expert_w_count, -1, bw_w1_quant.shape[-1]]) + bw_w1_scale = bw_w1_scale.reshape([expert_w_count, -1, bw_w1_scale.shape[-1]]) + # print("bwd down input 22") # quant do1 do1_fp8_reshape = do1.reshape([-1, do1.shape[-1]]) do1_fp8, do1_scale = kitchen_quant( do1_fp8_reshape, backend=kitchen.ops.Backend.CUTLASS, is_1d_scaled=True, return_transpose=False ) - do1_fp8 = do1_fp8.reshape([len(expert_w1), -1, do1_fp8.shape[-1]]) - do1_scale = do1_scale.reshape([len(expert_w1), -1, do1_scale.shape[-1]]) + do1_fp8 = do1_fp8.reshape([expert_w_count, -1, do1_fp8.shape[-1]]) + do1_scale = do1_scale.reshape([expert_w_count, -1, do1_scale.shape[-1]]) # group gemm if IF_USE_GROUP_GEMM_MASK: - dx = paddle.zeros(shape=[len(expert_w1), do1_fp8.shape[1], bw_w1_quant.shape[1]], dtype=paddle.bfloat16) + dx = paddle.zeros(shape=[expert_w_count, do1_fp8.shape[1], bw_w1_quant.shape[1]], dtype=paddle.bfloat16) deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked( (do1_fp8, do1_scale), (bw_w1_quant, bw_w1_scale), dx, tokens_per_expert, expected_m ) @@ -358,120 +341,118 @@ def bwd_gate_up_input(self, do1, expert_w1, tokens_per_expert, expected_m): return dx # ===== dw2 = deep_gemm(o2_t_fp8, do3_t_fp8) - - def bwd_down_weight(self, out_grad, o2, expert_w2): + def bwd_down_weight(self, out_grad, o2, expert_w2_list, expert_w_count): # transpose o2 - group_num = len(expert_w2) H2 = o2.shape[-1] + o2 = o2.reshape([-1, H2]) - o2_t = o2.reshape([group_num, -1, H2]).transpose([0, 2, 1]).contiguous().reshape([group_num * H2, -1]) - o2_t_fp8, o2_t_scale = kitchen_quant( - o2_t, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=True, return_transpose=False + _, _, o2_t_fp8, o2_t_scale = kitchen_quant( + o2, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=True, return_transpose=True ) - o2_t_fp8 = o2_t_fp8.reshape([group_num, H2, -1]) - - o2_t_scale = o2_t_scale.reshape([group_num, H2, -1]) - + o2_t_fp8 = o2_t_fp8.reshape([H2, expert_w_count, -1]).transpose([1, 0, 2]).contiguous() + o2_t_scale = o2_t_scale.reshape([expert_w_count, -1, H2]) # quant out_grad - H1 = out_grad.shape[-1] - out_grad = ( - out_grad.reshape([group_num, -1, H1]).transpose([0, 2, 1]).contiguous().reshape([group_num * H1, -1]) - ) + H1 = out_grad.shape[-1] + out_grad = out_grad.reshape([-1, H1]) - out_grad_fp8, out_grad_scale = kitchen_quant( - out_grad, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=True, return_transpose=False + _, _, out_grad_fp8, out_grad_scale = kitchen_quant( + out_grad, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=True, return_transpose=True ) - out_grad_fp8 = out_grad_fp8.reshape([group_num, H1, -1]) # [4, 8448, 7196] - out_grad_scale = paddle.split(out_grad_scale, num_or_sections=group_num, axis=-1) - # out_grad_scale = out_grad_scale.T.contiguous().reshape([group_num, H1, -1]) - # out_grad_scale = out_grad_scale.reshape([group_num, H1, -1]) + out_grad_fp8 = out_grad_fp8.reshape([H1, expert_w_count, -1]).transpose([1, 0, 2]).contiguous() + out_grad_scale = out_grad_scale.reshape([expert_w_count, -1, H1]) - for i in range(len(expert_w2)): - if hasattr(expert_w2[i], "main_grad"): - expert_w2[i].main_grad = kitchen_fp8_gemm( - o2_t_fp8[i], - o2_t_scale[i], - out_grad_fp8[i], - out_grad_scale[i], - True, - True, - expert_w2[i].main_grad, - ) + for i in range(expert_w_count): + if DSV3_USE_FUSED_Expert: + if hasattr(expert_w2_list[0], "main_grad"): + weight_grad = expert_w2_list[0].main_grad.reshape([expert_w_count, -1, expert_w2_list[0].shape[-1]])[i] + else: + weight_grad = expert_w2_list[0].grad.reshape([expert_w_count, -1, expert_w2_list[0].shape[-1]])[i] else: - expert_w2[i].grad = kitchen_fp8_gemm( + if hasattr(expert_w2_list[i], "main_grad"): + weight_grad = expert_w2_list[i].main_grad + else: + weight_grad = expert_w2_list[i].grad + + kitchen_fp8_gemm( o2_t_fp8[i], o2_t_scale[i], out_grad_fp8[i], out_grad_scale[i], True, True, - expert_w2[i].grad, - ) + weight_grad + ) + # ===== dw1 = deep_gemm(input_x_t_fp8, do1_t_fp8) - def bwd_gate_up_weight(self, do1, input_x, expert_w1): + def bwd_gate_up_weight(self, do1, input_x, expert_w1_list, expert_w_count): # transpose input_x and quant input_x - group_num = len(expert_w1) H1 = input_x.shape[-1] - input_x = input_x.reshape([group_num, -1, H1]).transpose([0, 2, 1]).contiguous().reshape([group_num * H1, -1]) - # input_x = input_x.reshape([group_num, -1, H1]).transpose([0, 2, 1]).reshape([group_num * H1, -1]).contiguous() + input_x = input_x.reshape([-1, H1]) - input_x_fp8, input_x_scale = kitchen_quant( - input_x, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=True, return_transpose=False - ) - input_x_fp8 = input_x_fp8.reshape([group_num, H1, -1]) - input_x_scale = input_x_scale.reshape([group_num, H1, -1]) + _, _, input_x_fp8, input_x_scale = kitchen_quant( + input_x, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=True, return_transpose=True + ) + + input_x_fp8 = input_x_fp8.reshape([H1, expert_w_count, -1]).transpose([1, 0, 2]).contiguous() + input_x_scale = input_x_scale.reshape([expert_w_count, -1, H1]) # transpose do1 and quant do1 H2 = do1.shape[-1] - do1 = do1.reshape([group_num, -1, H2]).transpose([0, 2, 1]).contiguous().reshape([group_num * H2, -1]) - do1_fp8, do1_scale = kitchen_quant( - do1, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=True, return_transpose=False + do1 = do1.reshape([-1, H2]) + _, _, do1_fp8, do1_scale = kitchen_quant( + do1, backend=kitchen.ops.Backend.CUBLAS, is_1d_scaled=True, return_transpose=True ) - do1_fp8 = do1_fp8.reshape([group_num, H2, -1]) - # do1_scale = do1_scale.T.contiguous().reshape([group_num, H2, -1]) - do1_scale = paddle.split(do1_scale, num_or_sections=group_num, axis=-1) - # do1_scale = do1_scale.reshape([group_num, H2, -1]) + do1_fp8 = do1_fp8.reshape([H2, expert_w_count, -1]).transpose([1, 0, 2]).contiguous() + do1_scale = do1_scale.reshape([expert_w_count, -1, H2]) # dw1 - for i in range(len(expert_w1)): - if hasattr(expert_w1[i], "main_grad"): - expert_w1[i].main_grad = kitchen_fp8_gemm( - input_x_fp8[i], - input_x_scale[i], - do1_fp8[i], - do1_scale[i], - True, - True, - expert_w1[i].main_grad, - ) + for i in range(expert_w_count): + if DSV3_USE_FUSED_Expert: + if hasattr(expert_w1_list[0], "main_grad"): + weight_grad = expert_w1_list[0].main_grad.reshape([expert_w_count, -1, expert_w1_list[0].shape[-1]])[i] + else: + weight_grad = expert_w1_list[0].grad.reshape([expert_w_count, -1, expert_w1_list[0].shape[-1]])[i] else: - expert_w1[i].grad = kitchen_fp8_gemm( - input_x_fp8[i], - input_x_scale[i], - do1_fp8[i], - do1_scale[i], - True, - True, - expert_w1[i].grad, - ) - + if hasattr(expert_w1_list[i], "main_grad"): + weight_grad = expert_w1_list[i].main_grad + else: + weight_grad = expert_w1_list[i].grad + + kitchen_fp8_gemm( + input_x_fp8[i], + input_x_scale[i], + do1_fp8[i], + do1_scale[i], + True, + True, + weight_grad, + ) + @paddle.no_grad() def forward(self, hs_out, hs_scale_out, unzipped_probs, tokens_per_expert): # self.tokens_per_expert = tokens_per_expert - # get w1 + # get w1, w2 expert_w1 = [x.w1 for x in self.custom_map.experts if x is not None] - - expert_w_count = len(expert_w1) - - # get w2 expert_w2 = [x.w2 for x in self.custom_map.experts if x is not None] + if DSV3_USE_FUSED_Expert: + # only one weight + assert( len(expert_w1) == 1, "only one weight 1 when enable fused expoert") + assert( len(expert_w2) == 1, "only one weight 2 when enable fused expoert") + expert_w1 = expert_w1[0] + expert_w2 = expert_w2[0] + else: + expert_w1 = paddle.concat( expert_w1, axis=0) + expert_w2 = paddle.concat( expert_w2, axis=0) + + expert_w_count = self.custom_map.num_local_experts + # o1 o1 = self.fwd_gate_up(hs_out, hs_scale_out, expert_w1, expert_w_count, tokens_per_expert) self.o1 = o1 @@ -494,28 +475,37 @@ def forward(self, hs_out, hs_scale_out, unzipped_probs, tokens_per_expert): @paddle.no_grad() def backward(self, out_grad, out_grad_scale, tokens_per_expert, dispatched_indices, expected_m): # recompute expert_w2 and expert_w1 - expert_w2 = [x.w2 for x in self.custom_map.experts if x is not None] - expert_w1 = [x.w1 for x in self.custom_map.experts if x is not None] + expert_w2_list = [x.w2 for x in self.custom_map.experts if x is not None] + expert_w1_list = [x.w1 for x in self.custom_map.experts if x is not None] + if DSV3_USE_FUSED_Expert: + # only one weight + assert( len(expert_w1_list) == 1, "only one weight 1 when enable fused expoert") + assert( len(expert_w2_list) == 1, "only one weight 2 when enable fused expoert") + expert_w1 = expert_w1_list[0] + expert_w2 = expert_w2_list[0] + else: + expert_w1 = paddle.concat( expert_w1_list, axis=0) + expert_w2 = paddle.concat( expert_w2_list, axis=0) + expert_w_count = self.custom_map.num_local_experts # do2 - do2, probs_grad, o2 = self.bwd_dowm_input(expert_w2, out_grad, out_grad_scale, tokens_per_expert, expected_m) + do2, probs_grad, o2 = self.bwd_dowm_input(expert_w2, expert_w_count, out_grad, out_grad_scale, tokens_per_expert, expected_m) # do1 do1 = self.bwd_swiglu(self.o1, do2) # dx - dx = self.bwd_gate_up_input(do1, expert_w1, tokens_per_expert, expected_m) + dx = self.bwd_gate_up_input(do1, expert_w1, expert_w_count, tokens_per_expert, expected_m) dx = dx.reshape([-1, dx.shape[-1]]) # dequant dout out_grad_dequant_fp16 = FQO.fused_act_dequant(out_grad, out_grad_scale) # dw2 - self.bwd_down_weight(out_grad_dequant_fp16, o2, expert_w2) + self.bwd_down_weight(out_grad_dequant_fp16, o2, expert_w2_list, expert_w_count) input_x = FQO.fused_act_dequant(self.unzipped_tokens, self.unzipped_scale) # dw1 - self.bwd_gate_up_weight(do1, input_x, expert_w1) - + self.bwd_gate_up_weight(do1, input_x, expert_w1_list, expert_w_count) self.reset_statue() return dx, probs_grad diff --git a/paddlenlp/transformers/moe_layer.py b/paddlenlp/transformers/moe_layer.py index 125dfc9de9e9..4b6d381acd5b 100644 --- a/paddlenlp/transformers/moe_layer.py +++ b/paddlenlp/transformers/moe_layer.py @@ -41,6 +41,8 @@ DSV3_USE_FP8_GROUP_GEMM = os.getenv("DSV3_USE_FP8_GROUP_GEMM", "False").lower() == "true" +DSV3_USE_FUSED_Expert = os.getenv("DSV3_USE_FUSED_Expert", "False").lower() == "true" + def dispatching(x, dispatch_mask, scatter_index, num_experts, capacity): """ @@ -218,6 +220,9 @@ def __init__( for i in range(self.moe_num_experts): if i // self.moe_num_experts_per_device == self.moe_rank: self.experts.append(expert_class(**expert_kwargs)) + + if DSV3_USE_FUSED_Expert: + break else: self.experts.append(None) @@ -663,6 +668,7 @@ def forward(self, hs_fp8_dispatched, hs_scale_dispatched, dispatched_indices, di total_zipped_tokens=hs_fp8_dispatched.shape[0], num_experts=4, ) + self.dispatched_probs = dispatched_probs expert_out_zipped.stop_gradient = False return expert_out_zipped