From 78c9f31c74859ee8aebdf514067ee3820b55877b Mon Sep 17 00:00:00 2001
From: Etherll <61019402+Etherll@users.noreply.github.com>
Date: Sat, 17 May 2025 02:06:42 +0300
Subject: [PATCH] Fix trust remote code (#2357)

* Update _utils.py

* Update loader.py

* Update loader.py

* Update vision.py

* Update unsloth/models/vision.py

* Update unsloth/models/vision.py

* Update unsloth/models/vision.py

* Update unsloth/models/vision.py

* Update unsloth/models/_utils.py

* Update unsloth/models/vision.py

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---
 unsloth/models/_utils.py |  5 +++--
 unsloth/models/loader.py |  2 ++
 unsloth/models/vision.py | 12 +++++++-----
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index e290ac08fe..63f48af659 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1174,6 +1174,7 @@ def unsloth_compile_transformers(
     import_from_cache       = False,
     disable                 = False,
     return_logits           = False,
+    unsloth_force_compile   = False,
 ):
     if Version(torch_version) < Version("2.4.0"):
         print(
@@ -1184,12 +1185,12 @@ def unsloth_compile_transformers(
         )
         return
     pass
-    if trust_remote_code:
+    if trust_remote_code and unsloth_force_compile == False:
         print(
             "Unsloth: We can't trace models if `trust_remote_code = True`, "\
             "so turning off some optimizations!"
         )
-        return
+        return model_types, False
     model_types = list(dict().fromkeys(model_types).keys())
     if disable: return model_types, False
 
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index a233b26a86..9c5a7b68be 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -485,6 +485,7 @@ class FastModel(FastBaseModel):
         auto_model                 = None,
         whisper_language           = None,
         whisper_task               = None,
+        unsloth_force_compile      = False,
         *args, **kwargs,
     ):
         if token is None: token = get_token()
@@ -715,6 +716,7 @@ class FastModel(FastBaseModel):
                 disable                 = False,
                 return_logits           = return_logits,
                 trust_remote_code       = trust_remote_code,
+                unsloth_force_compile   = unsloth_force_compile,
             )
         pass
 
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 64f4d71b6f..2bff87d8d9 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -479,10 +479,12 @@ class FastBaseModel:
                 unsloth_base_fast_generate.__doc__ = model._old_generate.__doc__
                 model.generate = types.MethodType(unsloth_base_fast_generate, model)
         pass
+        model._unsloth_trust_remote_code = trust_remote_code
         # Post patches
         model = FastBaseModel.post_patch_model(
             model,
             use_gradient_checkpointing = use_gradient_checkpointing,
+            trust_remote_code  = trust_remote_code,
         )
         # Clear deleted GPU items
         for _ in range(3):
@@ -516,7 +518,7 @@ class FastBaseModel:
         loftq_config               = {},
         task_type                  = TaskType.CAUSAL_LM,
         temporary_location         = "_unsloth_temporary_saved_buffers",
-        **kwargs,
+        **kwargs
     ):
         if os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1":
             print("Unsloth: Full finetuning is enabled, so .get_peft_model has no effect")
@@ -572,10 +574,9 @@ class FastBaseModel:
         model = _get_peft_model(model, lora_config)
         # Enable gradients on modules which are trainable
         requires_grad_for_gradient_checkpointing(model)
-
-        model = FastBaseModel.post_patch_model(model, use_gradient_checkpointing)
+        trust_remote_code = getattr(model, "_unsloth_trust_remote_code", False)
+        model = FastBaseModel.post_patch_model(model, use_gradient_checkpointing, trust_remote_code = trust_remote_code)
         model.max_seq_length = max_seq_length
-
         # Clear deleted GPU items
         for _ in range(3):
             gc.collect()
@@ -594,6 +595,7 @@ class FastBaseModel:
     def post_patch_model(
         model,
         use_gradient_checkpointing = True,
+        trust_remote_code = False,
     ):
         full_finetuning = os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1"
 
@@ -614,7 +616,7 @@ class FastBaseModel:
         )
 
         from transformers.trainer import Trainer 
-        if Trainer._inner_training_loop.__name__ != "_fast_inner_training_loop":
+        if Trainer._inner_training_loop.__name__ != "_fast_inner_training_loop" and trust_remote_code == False:
             raise RuntimeError('Unsloth: Unsuccessfully patched inner_training_loop')
         pass
         patch_saving_functions(model, vision = True)