diff --git a/paddleformers/trainer/trainer_callback.py b/paddleformers/trainer/trainer_callback.py
index 2c1300a11c3..53bd3f2b326 100644
--- a/paddleformers/trainer/trainer_callback.py
+++ b/paddleformers/trainer/trainer_callback.py
@@ -1017,6 +1017,24 @@ def on_step_begin(self, args, state, control, **kwargs):
             if collect is not None:
                 collect()
 
+    def on_optimizer_begin(self, args, state, control, scaler=None, **kwargs):
+        if not self._setup_done:
+            return
+
+        for monitor in self._monitor_dict.values():
+            finalize = getattr(monitor, "finalize_scaled_grad_metrics", None)
+            if finalize is not None:
+                finalize(scaler)
+
+    def on_substep_end(self, args, state, control, **kwargs):
+        if not self._setup_done:
+            return
+
+        for monitor in self._monitor_dict.values():
+            finalize = getattr(monitor, "finalize_composite_microbatch", None)
+            if finalize is not None:
+                finalize()
+
     def on_step_end(self, args, state, control, **kwargs):
         if not self._setup_done:
             return