From 4ff247ab18600fb3bc474ddd89dfdce76b50c287 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 24 Oct 2024 00:17:26 -0700
Subject: [PATCH] Fix DPO, ORPO

---
 unsloth/kernels/cross_entropy_loss.py |  1 -
 unsloth/models/_utils.py              | 33 +++++++++++++++++++++++++--
 unsloth/save.py                       |  2 +-
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py
index 4895d27310..1c8f8c8d99 100644
--- a/unsloth/kernels/cross_entropy_loss.py
+++ b/unsloth/kernels/cross_entropy_loss.py
@@ -375,7 +375,6 @@ def fast_cross_entropy_loss(
     )
     if n_items is None:
         n_items = torch.count_nonzero(labels != -100)
-        print(n_items)
     return loss.sum() / n_items
 pass
 
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 25024be2ec..6611b4f638 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1172,10 +1172,10 @@ pass
 
 def patch_gradient_accumulation_fix(Trainer):
     # Fixes gradient accumulation 
+    import inspect
     if hasattr(Trainer, "get_batch_samples"):
-        from inspect import getsource
         if \
-            not getsource(Trainer.get_batch_samples).strip()\
+            not inspect.getsource(Trainer.get_batch_samples).strip()\
             .endswith("return batch_samples, num_items_in_batch"):
 
             raise NotImplementedError("Unsloth: Please make a Github issue immediately!!")
@@ -1198,4 +1198,33 @@ def patch_gradient_accumulation_fix(Trainer):
             '`pip install --upgrade --no-cache-dir unsloth git+https://github.com/huggingface/transformers.git git+https://github.com/huggingface/trl.git`'
         )
     pass
+
+    # Also fix up loss scaling ie negate loss *= self.args.gradient_accumulation_steps
+    if "num_items_in_batch" not in inspect.signature(Trainer.training_step).parameters: return
+
+    function = inspect.getsource(Trainer.training_step)
+    where = function.find("def")
+    function = function.split("\n")
+    function = "\n".join(x[where:] for x in function)
+
+    # Import all variables that need importing
+    import transformers.trainer
+    items_in_trainer = dir(transformers.trainer)
+    good_items = []
+    for item in items_in_trainer:
+        # TODO: Support Deepspeed
+        if item.startswith(("deepspeed", "xm", "met", "smp")): continue
+        if item in function: good_items.append(item)
+    pass
+    exec("from transformers.trainer import (" + ", ".join(x for x in good_items) + ")", globals())
+
+    # Accelerate does / self.args.gradient_accumulation_steps internally, so if we already
+    # summed it up and did the division before hand, we have to negate it.
+    function = function.replace(
+        "loss *= self.args.gradient_accumulation_steps",
+        "if num_items_in_batch is not None: loss *= self.args.gradient_accumulation_steps",
+    )
+    function = function.replace("def training_step", "def _unsloth_training_step", 1)
+    exec(function, globals())
+    Trainer.training_step = _unsloth_training_step
 pass
diff --git a/unsloth/save.py b/unsloth/save.py
index ab30e0fea5..ccda79aeee 100644
--- a/unsloth/save.py
+++ b/unsloth/save.py
@@ -145,7 +145,7 @@ pass
 
 def _merge_lora(layer, name):
 
-    bias = None
+    bias = getattr(layer, "bias", None)
     if isinstance(layer, (Bnb_Linear4bit, Peft_Linear4bit, Peft_Linear)):
         # Is LoRA so we need to merge!
         W, quant_state, A, B, s, bias = get_lora_parameters_bias(layer)