Update _utils.py

unsloth/models/_utils.py

@@ -69,7 +69,6 @@ from unsloth_zoo.patching_utils import (
     patch_compiling_bitsandbytes,
     patch_layernorm,
     patch_torch_compile,
-    patch_regional_compilation,
     patch_model_and_tokenizer,
 )
 from unsloth_zoo.gradient_checkpointing import (
@@ -414,6 +413,26 @@ accelerate.utils.TorchDynamoPlugin.to_kwargs             = torch_compile_kwargs
 accelerate.accelerator.TorchDynamoPlugin.to_kwargs       = torch_compile_kwargs
 del accelerate
 
+def patch_regional_compilation():
+    # Regional torch 2.5 Recompilation - weirdly very slow??
+    if torch.nn.ModuleList.__name__ == "UnslothModuleList": return
+    # Only works for torch 2.5
+    if Version(torch.__version__) < Version("2.5.0"): return
+
+    old_module_list = torch.nn.ModuleList
+    os.environ["UNSLOTH_PATCHED"] = "1"
+
+    def UnslothModuleList(*args, **kwargs):
+        if len(args) == 1 and len(kwargs) == 0 and type(args[0]) is list:
+            args = [old_module_list([torch.compile(x, dynamic = True, options = torch_compile_options, fullgraph = False) for x in args[0]])]
+        return old_module_list(*args, **kwargs)
+    pass
+    UnslothModuleList.__doc__ = old_module_list.__doc__
+
+    torch.nn.ModuleList = UnslothModuleList
+    return
+pass
+
 # =============================================
 
 def prepare_model_for_kbit_training(