diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index cdc2cd45fd..903093e60f 100644 --- a/unsloth/models/_utils.py +++ b/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(