fix ema state assembler

paddleformers/trainer/trainer_utils.py

@@ -58,6 +58,13 @@
 
 # from ..ops import Topology
 from ..trainer.argparser import strtobool
+from ..utils.import_utils import is_paddlefleet_available
+
+if is_paddlefleet_available():
+    from ..transformers.gpt_provider import GPTModel
+else:
+    GPTModel = None
+
 from ..transformers.model_utils import (
     EMAStateHFFormatFullParamSaver,
     _add_variant,
@@ -2083,6 +2090,8 @@ def _build_ema_sharded_state_dict(self, ema_state_dict):
 
         ema_sharded_state_dict = {}
 
+        is_gpt_model = GPTModel is not None and isinstance(self.model, GPTModel)
+
         def _remove_layer_suffix(s):
             return re.sub(r"_layer_\d+$", "", s)
 
@@ -2102,7 +2111,8 @@ def _rename(key, add_mode=True):
                 assert (
                     self.expert_id_offset != -1
                 ), f"Your n_routed_experts is {self.model.config.n_routed_experts}, but you have param name:{key}, please check!"
-                key = _update_expert_number(key, self.expert_id_offset, add_mode)
+                if not is_gpt_model:
+                    key = _update_expert_number(key, self.expert_id_offset, add_mode)
             elif "_layer_" in key:
                 key = _remove_layer_suffix(key)
             return key
@@ -2146,12 +2156,6 @@ def _rename(key, add_mode=True):
                 v = paddle.reshape(v, expected_shape)
             ema_sharded_state_dict[k] = create_sharded_weight_with_new_local(k, v, ref_tensor)
 
-        # Fill missing params from model (e.g., e_score_correction_bias not tracked by EMA)
-        for k, v in self.model_sharded_state_dict.items():
-            if k not in ema_sharded_state_dict:
-                logger.debug(f"[EMAStateAssembler] Filling missing param {k} from model")
-                ema_sharded_state_dict[k] = v
-
         return ema_sharded_state_dict
 
     def _save_full_ema_states(self, step, ema_sharded_state_dict):