From 2bcd4c860aab53aa49ba730e0199d84fb4acabe7 Mon Sep 17 00:00:00 2001 From: Yichen Zhang Date: Wed, 1 Jul 2026 11:38:15 +0800 Subject: [PATCH 1/6] add config in SonicMoEExpert and unit test for sonicmoe recompute_z, bf16_wgrad --- .../transformer/moe/fusion_layer_utils.py | 6 +- src/paddlefleet/transformer/moe/moe_expert.py | 21 ++++++- src/paddlefleet/transformer/moe/moe_layer.py | 16 ++++-- .../multi_card_tests/moe/test_sonic_moe_ep.py | 46 +++++++++++++++ .../model/test_gpt_model_sonic_moe.py | 57 ++++++++++++++++++- 5 files changed, 138 insertions(+), 8 deletions(-) diff --git a/src/paddlefleet/transformer/moe/fusion_layer_utils.py b/src/paddlefleet/transformer/moe/fusion_layer_utils.py index 7d7106293..a1a5fd438 100644 --- a/src/paddlefleet/transformer/moe/fusion_layer_utils.py +++ b/src/paddlefleet/transformer/moe/fusion_layer_utils.py @@ -42,7 +42,6 @@ ) from paddlefleet_ops.sonicmoe.functional import ( _DownProjection, - _refresh_fp8_config, _UpProjection, ) from paddlefleet_ops.sonicmoe.functional.utils import enable_fp8 @@ -3187,6 +3186,7 @@ def run_sonic_moe( tokens_per_expert=None, fp8_scale=None, fp8_combine_grad_handle=None, + fp8_config=None, ): T = hidden_states.shape[0] stream_id = paddle.device.current_stream() @@ -3252,7 +3252,7 @@ def run_sonic_moe( w2_sonic = w2.permute([1, 2, 0]) with enable_fp8(fp8): - _refresh_fp8_config() + # _refresh_fp8_config() y1, z = _UpProjection.apply( hidden_states, w1_sonic, @@ -3270,6 +3270,7 @@ def run_sonic_moe( is_inference_mode_enabled=False, use_low_precision_postact_buffer=False, prequant_activation_payload=fp8_hidden_states, + fp8_config=fp8_config, ) hidden_states = _DownProjection.apply( y1, @@ -3290,6 +3291,7 @@ def run_sonic_moe( activation_type, None, fp8_combine_grad_handle, + fp8_config=fp8_config, ) return hidden_states diff --git a/src/paddlefleet/transformer/moe/moe_expert.py b/src/paddlefleet/transformer/moe/moe_expert.py index d31d00451..58aaa0f90 100644 --- a/src/paddlefleet/transformer/moe/moe_expert.py +++ b/src/paddlefleet/transformer/moe/moe_expert.py @@ -40,7 +40,10 @@ try: from paddlefleet_ops import deep_gemm as paddlefleet_deep_gemm - from paddlefleet_ops.sonicmoe.functional import clear_all_fp8_weight_caches + from paddlefleet_ops.sonicmoe.functional import ( + _refresh_fp8_config, + clear_all_fp8_weight_caches, + ) from paddlefleet_ops.sonicmoe.quack_utils import quantize_native_fp8_weights except (ImportError, RuntimeError): pass @@ -511,6 +514,11 @@ def __init__( self.hidden_size = self.config.hidden_size self.K = topk self._weights_layout = self._GROUPED_LAYOUT + self.sonic_moe_config = _refresh_fp8_config() + self.sonic_moe_config.enabled = self.config.fp8 is not None + self.sonic_moe_config.fp8_wgrad = self.config.fp8_wgrad + self.sonic_moe_config.fuse_y1_quant = True + self.sonic_moe_config.fuse_y1_bf16_trunc = True def _convert_grad_layout(self, param, converter): main_grad = getattr(param, "main_grad", None) @@ -583,6 +591,14 @@ def clear_fp8_weights(self): weight.fp8 = None weight.transposed_fp8 = None + def need_quant_weight(self): + for w in [self.weight1, self.weight2]: + if not hasattr(w, "fp8") or w.fp8 is None: + return True + if not hasattr(w, "transposed_fp8") or w.transposed_fp8 is None: + return True + return False + def forward( self, hidden_states, @@ -594,6 +610,8 @@ def forward( fp8_combine_grad_handle=None, ): self.convert_weights_to_sonic_layout() + if self.sonic_moe_config.enabled is True and self.need_quant_weight(): + self.quant_weight() hidden_states = run_sonic_moe( hidden_states, topk_indices, @@ -606,6 +624,7 @@ def forward( tokens_per_expert=tokens_per_expert, fp8_scale=fp8_scale, fp8_combine_grad_handle=fp8_combine_grad_handle, + fp8_config=self.sonic_moe_config, ) return hidden_states diff --git a/src/paddlefleet/transformer/moe/moe_layer.py b/src/paddlefleet/transformer/moe/moe_layer.py index 2e8acd53c..93ffbc043 100644 --- a/src/paddlefleet/transformer/moe/moe_layer.py +++ b/src/paddlefleet/transformer/moe/moe_layer.py @@ -180,6 +180,9 @@ def __init__( self.using_sonic_moe = self.config.using_sonic_moe self.fp8_dispatch = bool(config.fp8) self.fp8_wgrad = config.fp8_wgrad + self.fp8_dispatch_bwd = ( + self.fp8_dispatch and self.using_sonic_moe and self.fp8_wgrad + ) self.moe_expert_fusion = config.moe_expert_fusion self.moe_subbatch_token_num_after_dispatch = ( config.moe_subbatch_token_num_after_dispatch @@ -917,9 +920,12 @@ def fusion_moe_forward( self.num_experts_per_tok, tokens_per_expert, ) - fp8_combine_grad_handle = ( - {} if self.fp8_dispatch and self.using_sonic_moe else None + dispatched_indices = ( + self.token_dispatcher._comm_manager.dispatched_indices ) + dispatched_probs = self.token_dispatcher._comm_manager.dispatched_probs + fp8_combine_grad_handle = {} if self.fp8_dispatch_bwd else None + # fp8_combine_grad_handle = None with profile("fusion_mlp"): if self._use_hybrid_ep_fusion(): @@ -972,8 +978,10 @@ def fusion_moe_forward( with profile("combine"): hidden_states = self.combine( hidden_states, - combine_overlap_handle=combine_overlap_handle, - fp8_combine_grad_handle=fp8_combine_grad_handle, + combine_overlap_handle, + use_rr_deepep_combine=self.use_rr_deepep_combine, + fp8_dispatch=self.fp8_dispatch_bwd, + combine_grad_handle=fp8_combine_grad_handle, ) # Latent MoE: project back from latent space to hidden_size diff --git a/tests/multi_card_tests/moe/test_sonic_moe_ep.py b/tests/multi_card_tests/moe/test_sonic_moe_ep.py index 86bf687e4..1747fc9e1 100644 --- a/tests/multi_card_tests/moe/test_sonic_moe_ep.py +++ b/tests/multi_card_tests/moe/test_sonic_moe_ep.py @@ -605,9 +605,55 @@ def run_test_ep_precision(self): print("Expert-parallel FP8 precision checks passed!") + def run_test_bf16_wgrad(self): + acc_steps = 1 + + paddle.seed(self.seed) + model_parallel_cuda_manual_seed(self.seed) + moe_layer_sonic_bf16 = self._build_moe_layer(using_sonic_moe=True) + paddle.seed(self.seed) + model_parallel_cuda_manual_seed(self.seed) + moe_layer_sonic_fp8 = self._build_moe_layer( + using_sonic_moe=True, + fp8="e4m3", + fp8_wgrad=False, + ) + + input_data_list = [] + for step_idx in range(acc_steps): + paddle.seed(self.seed + step_idx) + input_data_list.append( + paddle.randn( + [4, 256, self.hidden_size], + dtype=paddle.bfloat16, + ) + ) + + output_bf16, grads_bf16 = self._run_accumulated_forward_backward( + moe_layer_sonic_bf16, input_data_list + ) + output_fp8, grads_fp8 = self._run_accumulated_forward_backward( + moe_layer_sonic_fp8, input_data_list + ) + + fp8_tol = 5e-3 + self._assert_tensor_diff_less( + output_fp8, + output_bf16, + tol=fp8_tol, + title="Sonic-MoE FP8 vs BF16 final output", + ) + self._assert_grad_diff_less( + grads_fp8, + grads_bf16, + tol=fp8_tol, + title="Sonic-MoE FP8 vs BF16 accumulated grad", + ) + def test_sonic_moe_all(self): self.run_test_sonic_moe_ep_grad_accumulation() self.run_test_ep_precision() + self.run_test_bf16_wgrad() if __name__ == "__main__": diff --git a/tests/single_card_tests/model/test_gpt_model_sonic_moe.py b/tests/single_card_tests/model/test_gpt_model_sonic_moe.py index 71c29cf88..329d03801 100644 --- a/tests/single_card_tests/model/test_gpt_model_sonic_moe.py +++ b/tests/single_card_tests/model/test_gpt_model_sonic_moe.py @@ -16,6 +16,7 @@ import subprocess import unittest +import numpy as np import paddle import paddle.nn.functional as F import paddlefleet_ops @@ -241,7 +242,7 @@ def _run_accumulated_forward_backward(self, moe_layer, input_data_list): self._collect_grads(moe_layer), ) - def test_moe_layer_precision(self): + def run_test_moe_layer_precision(self): """Test MoELayer precision: BF16 sonic-moe vs baseline, FP8 vs BF16. Both baseline and sonic layers are built with the same seed. @@ -262,9 +263,14 @@ def test_moe_layer_precision(self): using_sonic_moe=False, moe_deep_gemm=False ) moe_layer_sonic_bf16 = self._build_moe_layer(using_sonic_moe=True) + moe_layer_sonic_bf16.grouped_gemm_experts.sonic_moe_config.enabled = ( + False + ) + moe_layer_sonic_fp8 = self._build_moe_layer( using_sonic_moe=True, fp8="e4m3" ) + moe_layer_sonic_fp8.grouped_gemm_experts.sonic_moe_config.enabled = True input_data_list = [] for step_idx in range(self.acc_steps): @@ -282,6 +288,7 @@ def test_moe_layer_precision(self): moe_layer_sonic_bf16, input_data_list ) ) + moe_layer_sonic_fp8.grouped_gemm_experts.quant_weight() loss_fp8, output_fp8, grads_fp8 = ( self._run_accumulated_forward_backward( moe_layer_sonic_fp8, input_data_list @@ -369,6 +376,54 @@ def test_moe_layer_precision(self): print("All precision checks passed!") + def run_test_z_bf16_recompute_z(self): + sonic_fp8 = self._build_moe_layer(using_sonic_moe=True, fp8="e4m3") + sonic_fp8_recompute_z = self._build_moe_layer( + using_sonic_moe=True, fp8="e4m3" + ) + + sonic_fp8.grouped_gemm_experts.sonic_moe_config.save_z_fp8 = False + sonic_fp8_recompute_z.grouped_gemm_experts.sonic_moe_config.save_z_fp8 = False + sonic_fp8_recompute_z.grouped_gemm_experts.sonic_moe_config.recompute_z = True + # sonic_fp8.grouped_gemm_experts.sonic_moe_config.enabled = True + input_data_list = [] + for step_idx in range(self.acc_steps): + paddle.seed(self.seed + step_idx) + data = paddle.randn( + [2, 64, self.hidden_size], dtype=paddle.bfloat16 + ) + input_data_list.append(data) + + loss, output, grads = self._run_accumulated_forward_backward( + sonic_fp8, input_data_list + ) + loss_recompute, output_recompute, grads_recompute = ( + self._run_accumulated_forward_backward( + sonic_fp8_recompute_z, input_data_list + ) + ) + self.assertEqual(loss, loss_recompute, "Loss not equal.") + self.assertTrue( + np.array_equal(output.numpy(), output_recompute.numpy()), + "Output not equal.", + ) + + common_grads = set(grads) & set(grads_recompute) + self.assertTrue( + common_grads, + "No common FP8 grad tensors found", + ) + for name in sorted(common_grads): + g1 = grads[name] + g2 = grads_recompute[name] + self.assertTrue( + np.array_equal(g1.numpy(), g2.numpy()), "grad not equal." + ) + + def test_all_cases(self): + self.run_test_moe_layer_precision() + self.run_test_z_bf16_recompute_z() + if __name__ == "__main__": unittest.main() From 76391c500b7335226d944c2a172c674d5cc0285d Mon Sep 17 00:00:00 2001 From: Yichen Zhang Date: Wed, 1 Jul 2026 11:47:05 +0800 Subject: [PATCH 2/6] update sonicmoe --- packages/paddlefleet_ops/third_party/sonic-moe | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/paddlefleet_ops/third_party/sonic-moe b/packages/paddlefleet_ops/third_party/sonic-moe index 226ebb059..e7bbf3fb7 160000 --- a/packages/paddlefleet_ops/third_party/sonic-moe +++ b/packages/paddlefleet_ops/third_party/sonic-moe @@ -1 +1 @@ -Subproject commit 226ebb059e373540731167b27d152ae329dcbfb2 +Subproject commit e7bbf3fb7be741ed91ee91cb02dc6f5abbd55991 From 0059136f0036f4278a0e45f48951d2b7c6c6e203 Mon Sep 17 00:00:00 2001 From: Yichen Zhang Date: Wed, 1 Jul 2026 15:48:56 +0800 Subject: [PATCH 3/6] resolve conflicts --- src/paddlefleet/transformer/moe/moe_layer.py | 8 +++----- .../ai_edited_test/transformer/test_ai_moe_layer_7.py | 1 + 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/paddlefleet/transformer/moe/moe_layer.py b/src/paddlefleet/transformer/moe/moe_layer.py index 93ffbc043..fc943f1ea 100644 --- a/src/paddlefleet/transformer/moe/moe_layer.py +++ b/src/paddlefleet/transformer/moe/moe_layer.py @@ -752,7 +752,7 @@ def combine( hidden_states, combine_overlap_handle, use_rr_deepep_combine=self.use_rr_deepep_combine, - fp8_dispatch=self.fp8_dispatch and self.using_sonic_moe, + fp8_dispatch=self.fp8_dispatch_bwd, combine_grad_handle=fp8_combine_grad_handle, ) @@ -978,10 +978,8 @@ def fusion_moe_forward( with profile("combine"): hidden_states = self.combine( hidden_states, - combine_overlap_handle, - use_rr_deepep_combine=self.use_rr_deepep_combine, - fp8_dispatch=self.fp8_dispatch_bwd, - combine_grad_handle=fp8_combine_grad_handle, + combine_overlap_handle=combine_overlap_handle, + fp8_combine_grad_handle=fp8_combine_grad_handle, ) # Latent MoE: project back from latent space to hidden_size diff --git a/tests/single_card_tests/ai_edited_test/transformer/test_ai_moe_layer_7.py b/tests/single_card_tests/ai_edited_test/transformer/test_ai_moe_layer_7.py index 3dc365919..d46d5409e 100644 --- a/tests/single_card_tests/ai_edited_test/transformer/test_ai_moe_layer_7.py +++ b/tests/single_card_tests/ai_edited_test/transformer/test_ai_moe_layer_7.py @@ -286,6 +286,7 @@ def test_combine_delegates(self): layer.use_rr_deepep_combine = False layer.fp8_dispatch = False layer.using_sonic_moe = False + layer.fp8_dispatch_bwd = False hidden = paddle.randn([4, 64]) layer.token_dispatcher._comm_manager.combine.return_value = hidden result = layer.combine(hidden) From 317cbadaa1710679b4db17a329aadc7a027f1445 Mon Sep 17 00:00:00 2001 From: Yichen Zhang Date: Sat, 4 Jul 2026 18:00:58 +0800 Subject: [PATCH 4/6] pass original weight to sonicmoe to support main grad accumulation --- src/paddlefleet/transformer/moe/fusion_layer_utils.py | 4 ++++ src/paddlefleet/transformer/moe/moe_expert.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/src/paddlefleet/transformer/moe/fusion_layer_utils.py b/src/paddlefleet/transformer/moe/fusion_layer_utils.py index a1a5fd438..92a1fe0aa 100644 --- a/src/paddlefleet/transformer/moe/fusion_layer_utils.py +++ b/src/paddlefleet/transformer/moe/fusion_layer_utils.py @@ -3187,6 +3187,8 @@ def run_sonic_moe( fp8_scale=None, fp8_combine_grad_handle=None, fp8_config=None, + w1_original=None, + w2_original=None, ): T = hidden_states.shape[0] stream_id = paddle.device.current_stream() @@ -3271,6 +3273,7 @@ def run_sonic_moe( use_low_precision_postact_buffer=False, prequant_activation_payload=fp8_hidden_states, fp8_config=fp8_config, + w1_original=w1_original, ) hidden_states = _DownProjection.apply( y1, @@ -3292,6 +3295,7 @@ def run_sonic_moe( None, fp8_combine_grad_handle, fp8_config=fp8_config, + w2_original=w2_original, ) return hidden_states diff --git a/src/paddlefleet/transformer/moe/moe_expert.py b/src/paddlefleet/transformer/moe/moe_expert.py index 58aaa0f90..bb1c24544 100644 --- a/src/paddlefleet/transformer/moe/moe_expert.py +++ b/src/paddlefleet/transformer/moe/moe_expert.py @@ -625,6 +625,8 @@ def forward( fp8_scale=fp8_scale, fp8_combine_grad_handle=fp8_combine_grad_handle, fp8_config=self.sonic_moe_config, + w1_original=self.weight1, + w2_original=self.weight2, ) return hidden_states From 48b01cfa40360bb53e211075636ac6803da7650e Mon Sep 17 00:00:00 2001 From: Yichen Zhang Date: Mon, 6 Jul 2026 12:38:53 +0800 Subject: [PATCH 5/6] remove fp8 carrier --- .../transformer/moe/fusion_layer_utils.py | 20 ++++++++----------- src/paddlefleet/transformer/moe/moe_expert.py | 2 -- 2 files changed, 8 insertions(+), 14 deletions(-) diff --git a/src/paddlefleet/transformer/moe/fusion_layer_utils.py b/src/paddlefleet/transformer/moe/fusion_layer_utils.py index 92a1fe0aa..d93d02c7e 100644 --- a/src/paddlefleet/transformer/moe/fusion_layer_utils.py +++ b/src/paddlefleet/transformer/moe/fusion_layer_utils.py @@ -3187,8 +3187,6 @@ def run_sonic_moe( fp8_scale=None, fp8_combine_grad_handle=None, fp8_config=None, - w1_original=None, - w2_original=None, ): T = hidden_states.shape[0] stream_id = paddle.device.current_stream() @@ -3246,18 +3244,18 @@ def run_sonic_moe( if fp8_scale is not None: fp8_hidden_states = (hidden_states, fp8_scale) - if fp8: - w1_sonic = _make_sonic_fp8_weight_carrier(w1) - w2_sonic = _make_sonic_fp8_weight_carrier(w2) - else: - w1_sonic = w1.permute([1, 2, 0]) - w2_sonic = w2.permute([1, 2, 0]) + # if fp8: + # w1_sonic = _make_sonic_fp8_weight_carrier(w1) + # w2_sonic = _make_sonic_fp8_weight_carrier(w2) + # else: + # w1_sonic = w1.permute([1, 2, 0]) + # w2_sonic = w2.permute([1, 2, 0]) with enable_fp8(fp8): # _refresh_fp8_config() y1, z = _UpProjection.apply( hidden_states, - w1_sonic, + w1, None, expert_frequency_offset, total_expert_freq, @@ -3273,12 +3271,11 @@ def run_sonic_moe( use_low_precision_postact_buffer=False, prequant_activation_payload=fp8_hidden_states, fp8_config=fp8_config, - w1_original=w1_original, ) hidden_states = _DownProjection.apply( y1, z, - w2_sonic, + w2, None, scores_for_down, s_scatter_idx, @@ -3295,7 +3292,6 @@ def run_sonic_moe( None, fp8_combine_grad_handle, fp8_config=fp8_config, - w2_original=w2_original, ) return hidden_states diff --git a/src/paddlefleet/transformer/moe/moe_expert.py b/src/paddlefleet/transformer/moe/moe_expert.py index bb1c24544..58aaa0f90 100644 --- a/src/paddlefleet/transformer/moe/moe_expert.py +++ b/src/paddlefleet/transformer/moe/moe_expert.py @@ -625,8 +625,6 @@ def forward( fp8_scale=fp8_scale, fp8_combine_grad_handle=fp8_combine_grad_handle, fp8_config=self.sonic_moe_config, - w1_original=self.weight1, - w2_original=self.weight2, ) return hidden_states From 4ababba786a8276d8b85d5fa36967c842d635272 Mon Sep 17 00:00:00 2001 From: Yichen Zhang Date: Fri, 10 Jul 2026 17:55:16 +0800 Subject: [PATCH 6/6] add fused convert layout function, fused router score, and remove redundant permute --- .../paddlefleet_ops/third_party/sonic-moe | 2 +- .../transformer/moe/fusion_layer_utils.py | 158 ++++++++++++++---- src/paddlefleet/transformer/moe/moe_expert.py | 63 +++++-- src/paddlefleet/transformer/moe/moe_layer.py | 4 - .../multi_card_tests/moe/test_sonic_moe_ep.py | 9 +- .../model/test_gpt_model_sonic_moe.py | 9 +- 6 files changed, 183 insertions(+), 62 deletions(-) diff --git a/packages/paddlefleet_ops/third_party/sonic-moe b/packages/paddlefleet_ops/third_party/sonic-moe index e7bbf3fb7..cd918c5f0 160000 --- a/packages/paddlefleet_ops/third_party/sonic-moe +++ b/packages/paddlefleet_ops/third_party/sonic-moe @@ -1 +1 @@ -Subproject commit e7bbf3fb7be741ed91ee91cb02dc6f5abbd55991 +Subproject commit cd918c5f0bf9ac8e251b3f30223f060cef8df600 diff --git a/src/paddlefleet/transformer/moe/fusion_layer_utils.py b/src/paddlefleet/transformer/moe/fusion_layer_utils.py index d93d02c7e..7ae1f0d9c 100644 --- a/src/paddlefleet/transformer/moe/fusion_layer_utils.py +++ b/src/paddlefleet/transformer/moe/fusion_layer_utils.py @@ -32,10 +32,12 @@ tokens_zip_unique_add_with_subbatch, ) +_scatter_router_scores_i32 = None if paddlefleet_ops.is_sonic_moe_available(): from paddlefleet_ops.sonicmoe.enums import ActivationType from paddlefleet_ops.sonicmoe.ernie_compat.deepep_metadata import ( deepep_topk_to_sonic_metadata, + deepep_topk_to_sonic_metadata_with_scales, ) from paddlefleet_ops.sonicmoe.ernie_compat.mlp_node_v2 import ( _differentiable_router_scores, @@ -46,6 +48,13 @@ ) from paddlefleet_ops.sonicmoe.functional.utils import enable_fp8 + try: + from paddlefleet_ops.sonicmoe.quack_utils.blockscaled_fp8_gemm import ( + _scatter_router_scores_i32, + ) + except (ImportError, RuntimeError): + pass + logger = logging.getLogger(__name__) @@ -81,6 +90,57 @@ def _make_sonic_fp8_weight_carrier(weight): return _SonicFP8WeightCarrier.apply(weight) +class _SonicRouterScoresFromMetadata(paddle.autograd.PyLayer): + @staticmethod + def forward(ctx, topk_scores, metadata_scores, score_src_idx): + if len(topk_scores.shape) != 2: + raise ValueError( + f"topk_scores: expected rank 2, got shape {topk_scores.shape}" + ) + if len(metadata_scores.shape) != 1: + raise ValueError( + "metadata_scores: expected rank 1, got shape " + f"{metadata_scores.shape}" + ) + if len(score_src_idx.shape) != 1: + raise ValueError( + f"score_src_idx: expected rank 1, got shape {score_src_idx.shape}" + ) + if metadata_scores.shape[0] < score_src_idx.shape[0]: + raise ValueError( + "metadata_scores must include every real score referenced by " + f"score_src_idx; got {metadata_scores.shape[0]} scores and " + f"{score_src_idx.shape[0]} indices" + ) + if "int32" not in str(score_src_idx.dtype): + raise ValueError( + f"score_src_idx: expected int32, got {score_src_idx.dtype}" + ) + metadata_scores.stop_gradient = True + score_src_idx.stop_gradient = True + ctx.save_for_backward(score_src_idx) + ctx.input_shape = list(topk_scores.shape) + ctx.n_total = int(topk_scores.shape[0]) * int(topk_scores.shape[1]) + scores = metadata_scores.clone() + scores.stop_gradient = topk_scores.stop_gradient + return scores + + @staticmethod + def backward(ctx, grad_out): + (score_src_idx,) = ctx.saved_tensor() + if _scatter_router_scores_i32 is None: + raise RuntimeError( + "SonicMoE metadata router score backward requires " + "paddlefleet_ops.sonicmoe.quack_utils.blockscaled_fp8_gemm." + "_scatter_router_scores_i32; update paddlefleet_ops or use " + "the differentiable router-score fallback." + ) + grad_flat = _scatter_router_scores_i32( + grad_out.contiguous(), score_src_idx, ctx.n_total + ) + return grad_flat.reshape(ctx.input_shape), None, None + + class UnZipNode: """ UnZipNode 类用于对输入的token 矩阵根据分发索引进行解压操作,得到专家需要处理的 token。 @@ -3198,51 +3258,75 @@ def run_sonic_moe( paddle.int32 ) - # paddle.cast is not a no-op for matching dtype (it allocates + copies), - # so cast once with a guard instead of twice below. Allgather feeds int32 - # (zero copies); deepep feeds int64 (one copy instead of two). - topk_indices_i32 = ( - topk_indices - if topk_indices.dtype == paddle.int32 - else topk_indices.cast(paddle.int32) - ) - - ( - expert_frequency_offset, - x_gather_idx, - s_scatter_idx, - s_reverse_scatter_idx, - num_activated_expert_per_token_offset, - _router_scores, - TK_padded, - total_pad_rows, - _N_recv, - _score_src_idx, - ) = deepep_topk_to_sonic_metadata( - topk_indices_i32, - topk_scores, - tokens_per_expert, - E, - block=128 if fp8 else 1, - ) + fp8_scale_packed = None + if fp8 and fp8_scale is not None: + ( + expert_frequency_offset, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + num_activated_expert_per_token_offset, + _router_scores, + TK_padded, + total_pad_rows, + _N_recv, + _score_src_idx, + fp8_scale_packed, + ) = deepep_topk_to_sonic_metadata_with_scales( + topk_indices.cast(paddle.int32), + topk_scores, + tokens_per_expert, + E, + fp8_scale, + int(hidden_states.shape[1]), + block=128, + ) + else: + ( + expert_frequency_offset, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + num_activated_expert_per_token_offset, + _router_scores, + TK_padded, + total_pad_rows, + _N_recv, + _score_src_idx, + ) = deepep_topk_to_sonic_metadata( + topk_indices.cast(paddle.int32), + topk_scores, + tokens_per_expert, + E, + block=128 if fp8 else 1, + ) s_scatter_idx.stop_gradient = True activation_type = ActivationType("swiglu") total_expert_freq = TK_padded - scores_for_down = _differentiable_router_scores( - topk_scores, - topk_indices_i32, - num_activated_expert_per_token_offset, - TK_padded - total_pad_rows, - TK_padded, - E, - score_src_idx=_score_src_idx, - ) + if _score_src_idx is not None and _scatter_router_scores_i32 is not None: + scores_for_down = _SonicRouterScoresFromMetadata.apply( + topk_scores, _router_scores, _score_src_idx + ) + else: + scores_for_down = _differentiable_router_scores( + topk_scores, + topk_indices.cast(paddle.int32), + num_activated_expert_per_token_offset, + TK_padded - total_pad_rows, + TK_padded, + E, + score_src_idx=_score_src_idx, + ) fp8_hidden_states = None if fp8_scale is not None: - fp8_hidden_states = (hidden_states, fp8_scale) + fp8_hidden_states = ( + (hidden_states, fp8_scale, fp8_scale_packed) + if fp8_scale_packed is not None + else (hidden_states, fp8_scale) + ) # if fp8: # w1_sonic = _make_sonic_fp8_weight_carrier(w1) diff --git a/src/paddlefleet/transformer/moe/moe_expert.py b/src/paddlefleet/transformer/moe/moe_expert.py index 58aaa0f90..4758c8138 100644 --- a/src/paddlefleet/transformer/moe/moe_expert.py +++ b/src/paddlefleet/transformer/moe/moe_expert.py @@ -49,6 +49,18 @@ pass +try: + from paddlefleet_ops.sonicmoe.ernie_compat.weight_layout_fusion import ( + fused_grouped_w1_to_sonic, + fused_sonic_w1_to_grouped, + fused_transpose_w2_layout, + ) +except (ImportError, RuntimeError): + fused_grouped_w1_to_sonic = None + fused_sonic_w1_to_grouped = None + fused_transpose_w2_layout = None + + class BMMFunction(paddle.autograd.PyLayer): @staticmethod def forward(ctx, x, y, batch_sizes, trans_y=False): @@ -462,25 +474,34 @@ def _is_tensor_initialized(tensor): @staticmethod def _grouped_w1_to_sonic(weight): - target_shape = [weight.shape[0], weight.shape[2], weight.shape[1]] - gate, up = paddle.chunk(weight, 2, axis=-1) - gate = gate.transpose([0, 2, 1]) - up = up.transpose([0, 2, 1]) - return paddle.stack([gate, up], axis=2).reshape(target_shape) + if fused_grouped_w1_to_sonic is not None: + return fused_grouped_w1_to_sonic(weight) + else: + target_shape = [weight.shape[0], weight.shape[2], weight.shape[1]] + gate, up = paddle.chunk(weight, 2, axis=-1) + gate = gate.transpose([0, 2, 1]) + up = up.transpose([0, 2, 1]) + return paddle.stack([gate, up], axis=2).reshape(target_shape) @staticmethod def _sonic_w1_to_grouped(weight): - target_shape = [weight.shape[0], weight.shape[2], weight.shape[1]] - weight = weight.reshape([weight.shape[0], -1, 2, weight.shape[2]]) - gate = weight[:, :, 0, :].transpose([0, 2, 1]) - up = weight[:, :, 1, :].transpose([0, 2, 1]) - return paddle.concat([gate, up], axis=-1) + if fused_sonic_w1_to_grouped is not None: + return fused_sonic_w1_to_grouped(weight) + else: + target_shape = [weight.shape[0], weight.shape[2], weight.shape[1]] + weight = weight.reshape([weight.shape[0], -1, 2, weight.shape[2]]) + gate = weight[:, :, 0, :].transpose([0, 2, 1]) + up = weight[:, :, 1, :].transpose([0, 2, 1]) + return paddle.concat([gate, up], axis=-1) @staticmethod def _transpose_w2_layout(weight): # if not SonicMoEExpert._is_tensor_initialized(weight): # return weight - return weight.transpose([0, 2, 1]) + if fused_transpose_w2_layout is not None: + return fused_transpose_w2_layout(weight) + else: + return weight.transpose([0, 2, 1]) @staticmethod def _assign_tensor(tensor, value): @@ -520,9 +541,9 @@ def __init__( self.sonic_moe_config.fuse_y1_quant = True self.sonic_moe_config.fuse_y1_bf16_trunc = True - def _convert_grad_layout(self, param, converter): + def _convert_grad_layout(self, param, converter, convert_main_grad=True): main_grad = getattr(param, "main_grad", None) - if main_grad is not None: + if convert_main_grad and main_grad is not None: self._assign_tensor(main_grad, converter(main_grad)) if param.grad is not None and ( main_grad is None or param.grad.data_ptr() != main_grad.data_ptr() @@ -544,7 +565,15 @@ def _convert_layout( param.get_tensor()._set_dims([shape[0], shape[2], shape[1]]) else: self._assign_tensor(param, converter(param)) - self._convert_grad_layout(param, converter) + # weight2's main_grad stays in the original grouped layout + # ([E, I, H]); its grouped->sonic->grouped transpose is elided and + # the down-proj bf16/fp8 wgrad accumulates into it via a permute + # view. weight1's main_grad must still be converted because the + # grouped->sonic conversion also interleaves gate/up (a perfect + # shuffle the wgrad kernel does not undo). + self._convert_grad_layout( + param, converter, convert_main_grad=(param is self.weight1) + ) self._weights_layout = target_layout def convert_weights_to_sonic_layout(self): @@ -572,15 +601,15 @@ def quant_weight(self): self.convert_weights_to_sonic_layout() payload = quantize_native_fp8_weights( - self.weight1.permute([1, 2, 0]), - self.weight2.permute([1, 2, 0]), + self.weight1, + self.weight2, ) assert payload["format"] == "1x32", ( f"quant strategy {payload.get('format')} is not supported." ) w1_fp8, w1_scale, w1t_fp8, w1t_scale = payload["w1"] w2_fp8, w2_scale, w2t_fp8, w2t_scale = payload["w2"] - self.weight1.fp8 = (w1_fp8.mT, w1_scale) + self.weight1.fp8 = (w1_fp8, w1_scale) self.weight1.transposed_fp8 = (w1t_fp8, w1t_scale) self.weight2.fp8 = (w2_fp8, w2_scale) self.weight2.transposed_fp8 = (w2t_fp8, w2t_scale) diff --git a/src/paddlefleet/transformer/moe/moe_layer.py b/src/paddlefleet/transformer/moe/moe_layer.py index fc943f1ea..68f1f5ec0 100644 --- a/src/paddlefleet/transformer/moe/moe_layer.py +++ b/src/paddlefleet/transformer/moe/moe_layer.py @@ -920,10 +920,6 @@ def fusion_moe_forward( self.num_experts_per_tok, tokens_per_expert, ) - dispatched_indices = ( - self.token_dispatcher._comm_manager.dispatched_indices - ) - dispatched_probs = self.token_dispatcher._comm_manager.dispatched_probs fp8_combine_grad_handle = {} if self.fp8_dispatch_bwd else None # fp8_combine_grad_handle = None diff --git a/tests/multi_card_tests/moe/test_sonic_moe_ep.py b/tests/multi_card_tests/moe/test_sonic_moe_ep.py index 1747fc9e1..25041ec05 100644 --- a/tests/multi_card_tests/moe/test_sonic_moe_ep.py +++ b/tests/multi_card_tests/moe/test_sonic_moe_ep.py @@ -264,11 +264,16 @@ def _build_moe_layer( transformer_config, num_experts=self.n_routed_experts, ) - return MoELayer( + moe_layer = MoELayer( transformer_config, transformer_layer_spec.sublayers_spec.mlp.extra_kwargs["sublayers"], pg_collection or self.pg_collection, ) + mix_precision_utils.MixPrecisionLayer(moe_layer, dtype="bfloat16") + for param in moe_layer.parameters(): + if hasattr(param, "main_grad") and param.main_grad is None: + param.main_grad = paddle.zeros_like(param, dtype=paddle.float32) + return moe_layer @staticmethod def _expert_slice_for_rank(tensor, ep_rank, ep_size): @@ -352,7 +357,7 @@ def _run_forward_backward(self, moe_layer, input_data): master_grad=True, master_weight=True, ) - mix_precision_utils.MixPrecisionLayer(moe_layer, dtype="bfloat16") + # mix_precision_utils.MixPrecisionLayer(moe_layer, dtype="bfloat16") hidden_states = input_data.detach().clone() hidden_states.stop_gradient = False with paddle.amp.auto_cast(level="O2", dtype="bfloat16"): diff --git a/tests/single_card_tests/model/test_gpt_model_sonic_moe.py b/tests/single_card_tests/model/test_gpt_model_sonic_moe.py index 329d03801..aa4b137f7 100644 --- a/tests/single_card_tests/model/test_gpt_model_sonic_moe.py +++ b/tests/single_card_tests/model/test_gpt_model_sonic_moe.py @@ -21,6 +21,7 @@ import paddle.nn.functional as F import paddlefleet_ops from paddle.distributed import fleet +from paddle.distributed.fleet.utils import mix_precision_utils from paddlefleet_ops.utils import get_cuda_version # from tests.unit_tests.test_utilities import Utils @@ -115,7 +116,7 @@ def setUp(self): self.pg_collection = ProcessGroupCollection.use_mpu_process_groups() self.seed = 46 - self.hidden_size = 2048 + self.hidden_size = 512 self.n_routed_experts = 8 self.acc_steps = 1 @@ -177,6 +178,12 @@ def _build_moe_layer( self.pg_collection, ) + mix_precision_utils.MixPrecisionLayer(moe_layer, dtype="bfloat16") + + for param in moe_layer.parameters(): + if hasattr(param, "main_grad") and param.main_grad is None: + param.main_grad = paddle.zeros_like(param, dtype=paddle.float32) + return moe_layer @staticmethod