diff --git a/paddlenlp/trainer/trainer.py b/paddlenlp/trainer/trainer.py
index b48679dbf26a..5e8adb68f051 100644
--- a/paddlenlp/trainer/trainer.py
+++ b/paddlenlp/trainer/trainer.py
@@ -2526,6 +2526,9 @@ def training_step(
         if self.args.pipeline_parallel_degree > 1:
             return self.training_pipeline_step(model, inputs)
 
+        if hasattr(model, "_prepare_unified_non_pp_data"):
+            model._prepare_unified_non_pp_data(inputs)
+
         model.train()
         inputs = self._prepare_inputs(inputs)
         with self.autocast_smart_context_manager():