magic instance support (#4756)

Author: FeixLiu

paddleformers/trainer/trainer.py

@@ -4197,6 +4197,17 @@ def _dataset_process_function():
         with self.autocast_smart_context_manager():
             loss = model.forward_backward_pipeline(inputs, self.scaler if self.do_grad_scaling else None)
 
+        # MTP magic send: reset per-depth counters after each optimizer step
+        if getattr(self.args, "enable_mtp_magic_send", False):
+            try:
+                from paddlefleet.models.gpt.mtp_embedding_layer import (
+                    mtp_magic_instance,
+                )
+
+                mtp_magic_instance.clear_count_dict()
+            except (ImportError, ModuleNotFoundError):
+                pass
+
         return loss.detach()
 
     def save_model(
@@ -5399,6 +5410,16 @@ def _prepare_inputs_for_fleet(inputs):
         if need_clear:
             if hasattr(model, "_p2p_helper"):
                 model._p2p_helper.clear_meta_cache()
+            # MTP magic send: reset counters after eval step (shared singleton with train)
+            if getattr(self.args, "enable_mtp_magic_send", False):
+                try:
+                    from paddlefleet.models.gpt.mtp_embedding_layer import (
+                        mtp_magic_instance,
+                    )
+
+                    mtp_magic_instance.clear_count_dict()
+                except (ImportError, ModuleNotFoundError):
+                    pass
 
         return (loss, None, labels)