diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py index 3e55332c80..6357ddaf87 100644 --- a/unsloth/kernels/__init__.py +++ b/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, diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index f2377d55cc..7ec1c258a9 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/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 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b2f4b5c66e..bd40fbd2eb 100644 --- a/unsloth/models/_utils.py +++ b/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