fix #6448

Former-commit-id: 04f78e85af5af14b4c195936623e426a6a128af2

src/llamafactory/train/pt/trainer.py

@@ -19,7 +19,7 @@ import torch
 from transformers import Trainer
 from typing_extensions import override
 
-from ...extras.packages import is_transformers_version_equal_to_4_46, is_transformers_version_greater_than
+from ...extras.packages import is_transformers_version_greater_than
 from ..callbacks import SaveProcessorCallback
 from ..trainer_utils import create_custom_optimizer, create_custom_scheduler
 
@@ -78,15 +78,13 @@ class CustomTrainer(Trainer):
         self, model: "PreTrainedModel", inputs: Dict[str, "torch.Tensor"], return_outputs: bool = False, **kwargs
     ) -> Union["torch.Tensor", Tuple["torch.Tensor", List["torch.Tensor"]]]:
         r"""
-        Fixes the loss value for transformers 4.46.0.
-        https://github.com/huggingface/transformers/blob/v4.46.0/src/transformers/trainer.py#L3605
+        Fixes the loss value. See https://github.com/huggingface/transformers/pull/35438 for details.
         """
         loss = super().compute_loss(model, inputs, return_outputs, **kwargs)
-        if is_transformers_version_equal_to_4_46() and not getattr(self, "model_accepts_loss_kwargs", False):
-            # other model should not scale the loss
+        if kwargs.get("num_items_in_batch") and not getattr(self, "model_accepts_loss_kwargs", False):
             if return_outputs:
-                return (loss[0] / self.args.gradient_accumulation_steps, *loss[1:])
+                loss = (loss[0] / self.args.gradient_accumulation_steps, *loss[1:])
             else:
-                return loss / self.args.gradient_accumulation_steps
+                loss = loss / self.args.gradient_accumulation_steps
 
         return loss