PatchDPOTrainer

unsloth/models/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .loader import FastLanguageModel, FastDPOTrainer
+from .loader import FastLanguageModel
 from .llama import FastLlamaModel
 from .mistral import FastMistralModel
+from .dpo import PatchDPOTrainer

unsloth/models/dpo.py

@@ -17,9 +17,6 @@ from transformers.utils.notebook import (
     NotebookTrainingTracker,
     NotebookProgressCallback,
 )
-from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
-from trl import DPOTrainer
-import types
 
 DPOTrainer_metrics = [
     "rewards/chosen",
@@ -32,6 +29,7 @@ DPOTrainer_metrics = [
     "logits/chosen",
 ]
 
+
 def NotebookProgressCallback_on_train_begin(self, args, state, control, **kwargs):
     self.first_column = "Epoch" if args.evaluation_strategy == IntervalStrategy.EPOCH else "Step"
     self.training_loss = 0
@@ -67,17 +65,10 @@ def NotebookProgressCallback_on_log(self, args, state, control, logs=None, **kwa
 pass
 
 
-class FastDPOTrainer(DPOTrainer):
+def PatchDPOTrainer():
     # Patch DPO notebook printing
-    if (DEFAULT_PROGRESS_CALLBACK is NotebookProgressCallback):
-
-        DEFAULT_PROGRESS_CALLBACK.on_train_begin = types.MethodType(
-            NotebookProgressCallback_on_train_begin,
-            DEFAULT_PROGRESS_CALLBACK,
-        )
-        DEFAULT_PROGRESS_CALLBACK.on_log = types.MethodType(
-            NotebookProgressCallback_on_log,
-            DEFAULT_PROGRESS_CALLBACK,
-        )
-    pass
+    from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
+    DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin
+    DEFAULT_PROGRESS_CALLBACK.on_log         = NotebookProgressCallback_on_log
 pass
+

unsloth/models/loader.py

@@ -14,7 +14,6 @@
 
 from .llama import FastLlamaModel, logger
 from .mistral import FastMistralModel
-from .dpo import FastDPOTrainer
 from transformers import AutoConfig
 from transformers import __version__ as transformers_version