More patching

unsloth/kernels/__init__.py

@@ -16,6 +16,8 @@ from .cross_entropy_loss import (
     fast_cross_entropy_loss,
     patch_llama_for_causal_lm,
     unpatch_llama_for_causal_lm,
+    patch_transformers_losses,
+    patch_loss_function,
 )
 from .rms_layernorm import (
     fast_rms_layernorm,

unsloth/kernels/cross_entropy_loss.py

@@ -470,3 +470,52 @@ def unpatch_llama_for_causal_lm():
     transformers.models.llama.modeling_llama.LlamaForCausalLM = LlamaForCausalLM
     return
 pass
+
+
+@torch._disable_dynamo
+def UnslothForCausalLMLoss(
+    logits, labels, vocab_size: int, num_items_in_batch: int = None, ignore_index: int = -100, **kwargs
+):
+    shift_logits = logits
+    shift_labels = torch.empty_like(labels)
+    shift_labels[..., :-1] = labels[..., 1:]
+    shift_labels[..., -1] = -100
+    loss = fast_cross_entropy_loss(
+        logits  = shift_logits,
+        labels  = shift_labels,
+        n_items = num_items_in_batch,
+    )
+    return loss
+pass
+
+
+def patch_transformers_losses():
+    import re
+    try:
+        import transformers.loss.loss_utils
+    except:
+        logger.warning_once("Unsloth: Cannot patch loss functions - update transformers for faster modules!")
+
+    import transformers.modeling_utils
+    LOSS_MAPPING = transformers.loss.loss_utils.LOSS_MAPPING
+    LOSS_MAPPING["ForCausalLM"] = UnslothForCausalLMLoss
+
+    # Remove @property and @lru_cache
+    if hasattr(transformers.modeling_utils.PreTrainedModel.loss_function, "fget"):
+        transformers.modeling_utils.PreTrainedModel.loss_function = \
+            transformers.modeling_utils.PreTrainedModel.loss_function.fget.__wrapped__
+    pass
+pass
+
+
+def patch_loss_function(model):
+    try:
+        # model.loss_function starts as a dict to a loss fx
+        # We invoke it to save it
+        model.loss_function = model.loss_function()
+    except:
+        # Failed means we already invoked it, and we need args to the loss fx
+        pass
+    pass
+    return model
+pass

unsloth/models/_utils.py

@@ -44,6 +44,8 @@ __all__ = [
     "patch_gradient_checkpointing",
     "unpatch_gradient_checkpointing",
     "patch_gradient_accumulation_fix",
+    "patch_compiling_bitsandbytes",
+    "patch_regional_compilation",
 ]
 
 import torch
@@ -683,8 +685,19 @@ if Version(peft_version) < Version("0.12.0"):
         )
     pass
 pass
+
+# Also disable compiling on bitsandbytes
+def patch_compiling_bitsandbytes():
+    import peft.tuners.lora.bnb
+    peft.tuners.lora.bnb.Linear4bit.forward = \
+        torch._disable_dynamo(peft.tuners.lora.bnb.Linear4bit.forward)
+    peft.tuners.lora.bnb.Linear8bit.forward = \
+        torch._disable_dynamo(peft.tuners.lora.bnb.Linear8bit.forward)
+    return
+pass
 # =============================================
 
+
 import psutil
 def _get_statistics(statistics = None, force_download = True):
     # We log some basic stats about which environment is being used.
@@ -896,15 +909,39 @@ def unsloth_offloaded_gradient_checkpoint(function, *args, use_reentrant = None,
     return Unsloth_Offloaded_Gradient_Checkpointer.apply(function, *args)
 pass
 
-
 import torch.utils
-old_checkpoint = torch.utils.checkpoint
 def patch_gradient_checkpointing():
-    torch.utils.checkpoint = unsloth_offloaded_gradient_checkpoint
+    if torch.utils.checkpoint.checkpoint.__name__ == "unsloth_offloaded_gradient_checkpoint": return
+    torch.utils.checkpoint._old_checkpoint = torch.utils.checkpoint.checkpoint
+    torch.utils.checkpoint.checkpoint = unsloth_offloaded_gradient_checkpoint
 pass
 
 def unpatch_gradient_checkpointing():
-    torch.utils.checkpoint = old_checkpoint
+    if hasattr(torch.utils.checkpoint, "_old_checkpoint"):
+        torch.utils.checkpoint.checkpoint = torch.utils.checkpoint._old_checkpoint
+        del torch.utils.checkpoint._old_checkpoint
+    pass
+pass
+
+
+# =============================================
+# Regional torch 2.5 Recompilation - weirdly very slow??
+def patch_regional_compilation():
+    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
+
+    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