Skip to content

support dispatch both bf16 & fp8 #10817

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 4, 2025
+316 −181
Merged

support dispatch both bf16 & fp8 #10817

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 0 additions & 3 deletions paddlenlp/transformers/deepseek_v2/configuration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,6 @@ def __init__(
use_dualpipev=False,
send_mtp_embed=False,
recompute_fwd_gate_up=False,
dequant_input=False,
is_split_group_gemm=False,
**kwargs,
):
Expand Down Expand Up @@ -235,10 +234,8 @@ def __init__(
self.use_dualpipev = use_dualpipev
self.send_mtp_embed = send_mtp_embed
self.recompute_fwd_gate_up = recompute_fwd_gate_up
self.dequant_input = dequant_input
self.is_split_group_gemm = is_split_group_gemm


super().__init__(
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
Expand Down
70 changes: 34 additions & 36 deletions paddlenlp/transformers/deepseek_v2/modeling_pp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@


DSV3_USE_FP8_GEMM = os.getenv("DSV3_USE_FP8_GEMM", "False").lower() == "true"
DSV3_USE_FP8_DISPATCH = os.getenv("DSV3_USE_FP8_DISPATCH", "False").lower() == "true"


def parse_args(args):
Expand Down Expand Up @@ -510,7 +511,7 @@ def dispatch_forward(self, inputs, previous_event=None, async_finish=False, allo
token_probs,
) = inputs

(hs_fp16_dispatched, dispatched_indices, dispatched_probs,) = self.fp8_fusion_moe_node.dispatch_node.forward(
(hs_dispatched, dispatched_indices, dispatched_probs,) = self.fp8_fusion_moe_node.dispatch_node.forward(
hs_2d,
token_indices,
token_probs,
Expand All @@ -524,7 +525,7 @@ def dispatch_forward(self, inputs, previous_event=None, async_finish=False, allo
hidden_states,
residual,
l_aux,
hs_fp16_dispatched,
hs_dispatched,
dispatched_indices,
dispatched_probs,
)
Expand All @@ -533,7 +534,7 @@ def dispatch_forward(self, inputs, previous_event=None, async_finish=False, allo
hidden_states,
residual,
l_aux,
hs_fp16_dispatched,
hs_dispatched,
dispatched_indices,
dispatched_probs,
)
Expand All @@ -545,7 +546,7 @@ def mlp_forward(self, inputs):
hidden_states,
residual,
l_aux,
hs_fp16_dispatched,
hs_dispatched,
dispatched_indices,
dispatched_probs,
) = inputs
Expand All @@ -554,12 +555,12 @@ def mlp_forward(self, inputs):
hidden_states,
residual,
l_aux,
hs_fp16_dispatched,
hs_dispatched,
dispatched_indices,
dispatched_probs,
) = inputs
hidden_states_out = self.fp8_fusion_moe_node.mlp_node.forward(
hs_fp16_dispatched, dispatched_indices, dispatched_probs
hs_dispatched, dispatched_indices, dispatched_probs
)

if self.send_mtp_embed:
Expand All @@ -573,18 +574,18 @@ def combine_forward(self, inputs, async_finish=False):
else:
(hidden_states, residual, l_aux, hidden_states_out) = inputs

output_combie = self.fp8_fusion_moe_node.combine_node.forward(hidden_states_out, async_finish=async_finish)
output_combine = self.fp8_fusion_moe_node.combine_node.forward(hidden_states_out, async_finish=async_finish)
if self.send_mtp_embed:
return (inputs_embeds_mtp, hidden_states, residual, l_aux, output_combie)
return (inputs_embeds_mtp, hidden_states, residual, l_aux, output_combine)
else:
return (hidden_states, residual, l_aux, output_combie)
return (hidden_states, residual, l_aux, output_combine)

def post_process_forward(self, inputs):
if self.send_mtp_embed:
(inputs_embeds_mtp, hidden_states, residual, l_aux, output_combie) = inputs
(inputs_embeds_mtp, hidden_states, residual, l_aux, output_combine) = inputs
else:
(hidden_states, residual, l_aux, output_combie) = inputs
final_hidden_states = self.fp8_fusion_moe_node.combine_quant_node.forward(output_combie)
(hidden_states, residual, l_aux, output_combine) = inputs
final_hidden_states = self.fp8_fusion_moe_node.combine_quant_node.forward(output_combine)
if self.send_mtp_embed:
inputs = (inputs_embeds_mtp, hidden_states, residual, l_aux, final_hidden_states)
else:
Expand All @@ -609,21 +610,21 @@ def post_process_backward(self, output_grad):
l_aux_grad,
final_hidden_states_grad,
) = self.post_process_node.backward(output_grad)
output_combie_grad = self.fp8_fusion_moe_node.combine_quant_node.backward(final_hidden_states_grad)
output_combine_grad = self.fp8_fusion_moe_node.combine_quant_node.backward(final_hidden_states_grad)
if self.send_mtp_embed:
return (
inputs_embeds_mtp_grad,
hidden_states_grad,
residual_grad,
l_aux_grad,
output_combie_grad,
output_combine_grad,
)
else:
return (
hidden_states_grad,
residual_grad,
l_aux_grad,
output_combie_grad,
output_combine_grad,
)

def combine_backward(self, output_grad, async_finish=False):
Expand All @@ -633,18 +634,18 @@ def combine_backward(self, output_grad, async_finish=False):
hidden_states_grad,
residual_grad,
l_aux_grad,
output_combie_grad_bf16,
output_combine_grad,
) = output_grad
else:
(
hidden_states_grad,
residual_grad,
l_aux_grad,
output_combie_grad_bf16,
output_combine_grad,
) = output_grad

hidden_states_out_grad_bf16 = self.fp8_fusion_moe_node.combine_node.backward(
output_combie_grad_bf16,
hidden_states_out_grad = self.fp8_fusion_moe_node.combine_node.backward(
output_combine_grad,
async_finish=async_finish,
)

Expand All @@ -654,14 +655,14 @@ def combine_backward(self, output_grad, async_finish=False):
hidden_states_grad,
residual_grad,
l_aux_grad,
hidden_states_out_grad_bf16,
hidden_states_out_grad,
)
else:
return (
hidden_states_grad,
residual_grad,
l_aux_grad,
hidden_states_out_grad_bf16,
hidden_states_out_grad,
)

def mlp_backward(self, output_grad):
Expand All @@ -680,25 +681,23 @@ def mlp_backward(self, output_grad):
l_aux_grad,
hidden_states_out_grad,
) = output_grad
hs_fp16_dispatched_grad, dispatched_probs_grad = self.fp8_fusion_moe_node.mlp_node.backward(
hidden_states_out_grad
)
hs_dispatched_grad, dispatched_probs_grad = self.fp8_fusion_moe_node.mlp_node.backward(hidden_states_out_grad)

if self.send_mtp_embed:
return (
inputs_embeds_mtp_grad,
hidden_states_grad,
residual_grad,
l_aux_grad,
hs_fp16_dispatched_grad,
hs_dispatched_grad,
dispatched_probs_grad,
)
else:
return (
hidden_states_grad,
residual_grad,
l_aux_grad,
hs_fp16_dispatched_grad,
hs_dispatched_grad,
dispatched_probs_grad,
)

Expand All @@ -709,19 +708,19 @@ def dispatch_backward(self, output_grad, async_finish=False):
hidden_states_grad,
residual_grad,
l_aux_grad,
hs_bf16_dispatched_grad,
hs_dispatched_grad,
dispatched_probs_grad,
) = output_grad
else:
(
hidden_states_grad,
residual_grad,
l_aux_grad,
hs_bf16_dispatched_grad,
hs_dispatched_grad,
dispatched_probs_grad,
) = output_grad
hs_bf16_grad, token_probs_grad = self.fp8_fusion_moe_node.dispatch_node.backward(
hs_bf16_dispatched_grad, dispatched_probs_grad, async_finish=async_finish
hs_grad, token_probs_grad = self.fp8_fusion_moe_node.dispatch_node.backward(
hs_dispatched_grad, dispatched_probs_grad, async_finish=async_finish
)

if self.send_mtp_embed:
Expand All @@ -730,11 +729,11 @@ def dispatch_backward(self, output_grad, async_finish=False):
hidden_states_grad,
residual_grad,
l_aux_grad,
hs_bf16_grad,
hs_grad,
token_probs_grad,
)
else:
return (hidden_states_grad, residual_grad, l_aux_grad, hs_bf16_grad, token_probs_grad)
return (hidden_states_grad, residual_grad, l_aux_grad, hs_grad, token_probs_grad)

def attn_backward(self, output_grad):
if self.send_mtp_embed:
Expand All @@ -743,19 +742,19 @@ def attn_backward(self, output_grad):
hidden_states_grad,
residual_grad,
l_aux_grad,
hs_bf16_grad,
hs_grad,
token_probs_grad,
) = output_grad
else:
(
hidden_states_grad,
residual_grad,
l_aux_grad,
hs_bf16_grad,
hs_grad,
token_probs_grad,
) = output_grad
hidden_states_grad_, probs_grad, routing_map_grad = self.fp8_fusion_moe_node.dispatch_quant_node.backward(
hs_bf16_grad, token_probs_grad
hs_grad, token_probs_grad
)

if self.send_mtp_embed:
Expand Down Expand Up @@ -1234,7 +1233,6 @@ def build_schedule_node(self):
fp8_fusion_moe_node = FusionMoeNode(
self.mlp,
recompute_fwd_gate_up=self.config.recompute_fwd_gate_up,
dequant_input=self.config.dequant_input,
is_split_group_gemm=self.config.is_split_group_gemm,
name="fp8_fusion_moe_node",
)
Expand Down
Loading
Loading