update trainer

llm/run_embedding.py

@@ -155,6 +155,11 @@ def main():
     if tokenizer.chat_template is not None:
         data_args.eval_with_do_generation = False
 
+    if training_args.do_eval:
+        logger.warning("Warning: 'do_eval' is set to True, but will be set to False for Embedding training currently.")
+        training_args.do_eval = False
+        training_args.evaluation_strategy = "no"
+
     if data_args.dataset_name_or_path is None:
         raise ValueError(f"Please specific dataset name or path (got {data_args.dataset_name_or_path})")
     elif os.path.exists(os.path.join(data_args.dataset_name_or_path, "train.json")) or os.path.exists(
@@ -259,7 +264,7 @@ def main():
         padding = True
 
     if training_args.pipeline_parallel_degree > 1:
-        metrics = None
+        raise NotImplementedError("Cannot support pipeline parallel for Embedding training now.")
     else:
         metrics = compute_metrics
 

paddlenlp/transformers/qwen2/modeling.py

@@ -1696,13 +1696,13 @@ def forward(
         q_reps = nn.functional.normalize(q_reps, axis=-1)
         p_reps = nn.functional.normalize(p_reps, axis=-1)
 
+        if return_encode:
+            return q_reps, p_reps
+
         if self.embedding_negatives_cross_device:
             q_reps = self._dist_gather_tensor(q_reps)
             p_reps = self._dist_gather_tensor(p_reps)
 
-        if return_encode:
-            return q_reps, p_reps
-
         loss = self.in_batch_negative_loss(q_reps, p_reps)
         return loss
 

paddlenlp/trl/embedding_trainer.py

@@ -28,12 +28,11 @@
 
 
 class EmbeddingTrainer(Trainer):
-    def __init__(self, model_args, use_gradient_cache=False, **kwargs):
+    def __init__(self, model_args, **kwargs):
         super().__init__(**kwargs)
 
         self.model_args = model_args
         self.embedding_negatives_cross_device = model_args.embedding_negatives_cross_device
-        self.use_gradient_cache = use_gradient_cache
         self.accum_data = []
         self.accum_freq = 0
         self.accum_q_features = []
@@ -168,7 +167,7 @@ def training_step(
         if self.args.pipeline_parallel_degree > 1:
             raise NotImplementedError("Cannot support pipeline parallel for Embedding training now.")
 
-        if self.args.gradient_accumulation_steps == 1 or not self.use_gradient_cache:
+        if self.args.gradient_accumulation_steps == 1:
             return super().training_step(model, inputs)
         else:
             self.forward_no_grad(model, inputs)