add callback for fp8 sonicmoe

Author: pkuzyc

paddleformers/trainer/trainer.py

@@ -173,6 +173,7 @@
     InterleaveGateUpCallback,
     PrinterCallback,
     ProgressCallback,
+    SonicMoELayoutSwitchCallback,
     SPGradSyncCallback,
     TrainerCallback,
     TrainerControl,
@@ -1648,8 +1649,10 @@ def train(
             self.add_non_zcc_ema_callback(resume_from_checkpoint)
 
         if self.args.using_sonic_moe:
-            callback = InterleaveGateUpCallback(self.model, resume_from_checkpoint, self.args.output_dir)
-            self.add_callback(callback)
+            # callback = InterleaveGateUpCallback(self.model, resume_from_checkpoint, self.args.output_dir)
+            # self.add_callback(callback)
+            print("==== add sonicmoe callback ====")
+            self.add_callback(SonicMoELayoutSwitchCallback())
 
         self.log_trainable_numel(model)
 

paddleformers/trainer/trainer_callback.py

@@ -42,13 +42,17 @@
 # Conditionally import paddlefleet modules
 if is_paddlefleet_available():
     from paddlefleet.models.gpt import GPTModel
+    from paddlefleet.transformer.moe.moe_expert import SonicMoEExpert
     from paddlefleet.transformer.moe.moe_layer import MoELayer
     from paddlefleet.transformer.moe.moe_router import StandardMoERouter
 else:
 
     class GPTModel:
         pass
 
+    class SonicMoEExpert:
+        pass
+
     class MoELayer:
         pass
 
@@ -80,6 +84,7 @@ class StandardMoERouter:
     "MoEGateSpGradSyncCallBack",
     "SPGradSyncCallback",
     "EMAStateAssemblerCallback",
+    "SonicMoELayoutSwitchCallback",
 ]
 
 
@@ -715,6 +720,9 @@ def on_step_begin(self, args, state, control, **kwargs):
         """
         Quantize expert weights to FP8 before each training step
         """
+        if args.using_sonic_moe:
+            # sonicmoe cannot support offline quant now.
+            return
         model = kwargs["model"]
         optimizer = kwargs["optimizer"]
         global skip_count
@@ -754,6 +762,9 @@ def on_optimizer_begin(self, args, state, control, **kwargs):
         """
         Reload weights before optimizer step
         """
+        if args.using_sonic_moe:
+            # sonicmoe cannot support offline quant now.
+            return
         model = kwargs["model"]
         optimizer = kwargs["optimizer"]
         global skip_count
@@ -932,6 +943,32 @@ def on_step_end(self, args, state, control, **kwargs):
         logger.info(f"[EMAStateAssembler] Assembling EMA state took {duration:.3f} seconds.")
 
 
+class SonicMoELayoutSwitchCallback(TrainerCallback):
+    def _apply_to_sonic_moe_experts(self, model, fn_name):
+        def apply_layout_switch(layer):
+            if isinstance(layer, SonicMoEExpert):
+                getattr(layer, fn_name)()
+
+        model.apply(apply_layout_switch)
+
+    def _prepare_sonic_moe_fp8_weights(self, model):
+        def prepare_fp8_weights(layer):
+            if isinstance(layer, SonicMoEExpert):
+                layer.convert_weights_to_sonic_layout()
+                layer.quant_weight()
+
+        model.apply(prepare_fp8_weights)
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        if args.using_sonic_moe:
+            self._prepare_sonic_moe_fp8_weights(kwargs["model"])
+            # kwargs["optimizer"].clear_param_storage("moe_expert")
+
+    def on_optimizer_begin(self, args, state, control, **kwargs):
+        if args.using_sonic_moe:
+            self._apply_to_sonic_moe_experts(kwargs["model"], "convert_weights_to_grouped_layout")
+
+
 class InterleaveGateUpCallback(TrainerCallback):
     def __init__(self, model, resume_from_checkpoint=None, output_dir=None):
         self.model = model

paddleformers/transformers/configuration_utils.py

@@ -444,6 +444,12 @@ class LlmMetaConfig:
             True,
             "Whether to use FP8 for gradient storage during training (only effective if `fp8=True`). Further reduces memory footprint but may introduce minor numerical error. Defaults to False.",
         ),
+        (
+            "use_ue8m0",
+            bool,
+            False,
+            "Whether to use UE8M0 packed scaling factors for FP8 on Blackwell GPUs (SM100+). Enables deep_gemm backend for weight gradient computation. Defaults to False.",
+        ),
     ]
 
     model_conf = [

paddleformers/transformers/qwen3_moe/modeling.py

@@ -820,8 +820,8 @@ def _gen_aoa_config(cls, config: Qwen3MoeConfig):
                 group_gemm1 = ",".join(ep_weight1)
                 group_gemm2 = ",".join(ep_weight2)
                 aoa_config["aoa_statements"] += [
-                    f"{group_gemm1} -> {tgt_prefix}.mlp.grouped_gemm_experts.weight1, axis=0"
-                    f"{group_gemm2} -> {tgt_prefix}.mlp.grouped_gemm_experts.weight2, axis=0"
+                    f"{group_gemm1} -> {tgt_prefix}.mlp.grouped_gemm_experts.weight1, axis=0",
+                    f"{group_gemm2} -> {tgt_prefix}.mlp.grouped_gemm_experts.weight2, axis=0",
                 ]
         else:
             if config.get("fd_fallback", False):
@@ -903,8 +903,8 @@ def _gen_inv_aoa_config(cls, config: Qwen3MoeConfig):
                 group_gemm1 = ",".join(ep_weight1)
                 group_gemm2 = ",".join(ep_weight2)
                 aoa_statements += [
-                    f"{model_prefix}layers.{layer_id}.mlp.grouped_gemm_experts.weight1 -> {group_gemm1}, axis=0"
-                    f"{model_prefix}layers.{layer_id}.mlp.grouped_gemm_experts.weight2 -> {group_gemm2}, axis=0"
+                    f"{model_prefix}layers.{layer_id}.mlp.grouped_gemm_experts.weight1 -> {group_gemm1}, axis=0",
+                    f"{model_prefix}layers.{layer_id}.mlp.grouped_gemm_experts.weight2 -> {group_gemm2}, axis=0",
                 ]
         else:
             if config.get("fd_fallback", False):