diff --git a/pyproject.toml b/pyproject.toml index fc9c8256ad..e0c5d93562 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,10 +33,10 @@ exclude = ["images*"] [project.optional-dependencies] huggingface = [ - "unsloth_zoo", + "unsloth_zoo>=2024.11.1", "packaging", "tyro", - "transformers>=4.44.2", + "transformers>=4.46.1", "datasets>=2.16.0", "sentencepiece>=0.2.0", "tqdm", @@ -244,10 +244,10 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3", ] colab-new = [ - "unsloth_zoo", + "unsloth_zoo>=2024.11.1", "packaging", "tyro", - "transformers>=4.44.2", + "transformers>=4.46.1", "datasets>=2.16.0", "sentencepiece>=0.2.0", "tqdm", diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 458c2696bc..5102d8f466 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -27,13 +27,6 @@ import numpy as np # pass # pass -# Check for unsloth_zoo -try: - import unsloth_zoo -except: - raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth-zoo`") -pass - # Unsloth currently does not work on multi GPU setups - sadly we are a 2 brother team so # enabling it will require much more work, so we have to prioritize. Please understand! # We do have a beta version, which you can contact us about! @@ -60,6 +53,14 @@ pass # Reduce VRAM usage by reducing fragmentation os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" +# Hugging Face Hub faster downloads +if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ: + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" +pass + +# Log Unsloth is being used +os.environ["UNSLOTH_IS_PRESENT"] = "1" + try: import torch except ModuleNotFoundError: @@ -71,12 +72,6 @@ except Exception as exception: raise exception pass -# Hugging Face Hub faster downloads (only enable during Colab and Kaggle sessions) -keynames = "\n" + "\n".join(os.environ.keys()) -if "\nCOLAB_" in keynames or "\nKAGGLE_" in keynames: - os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" -pass - # We support Pytorch 2 # Fixes https://github.com/unslothai/unsloth/issues/38 torch_version = torch.__version__.split(".") @@ -165,6 +160,13 @@ if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ pass pass +# Check for unsloth_zoo +try: + import unsloth_zoo +except: + raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth-zoo`") +pass + from .models import * from .save import * from .chat_templates import * diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py index 3e55332c80..82e7641693 100644 --- a/unsloth/kernels/__init__.py +++ b/unsloth/kernels/__init__.py @@ -14,8 +14,8 @@ from .cross_entropy_loss import ( fast_cross_entropy_loss, - patch_llama_for_causal_lm, - unpatch_llama_for_causal_lm, + post_patch_loss_function, + patch_loss_functions, ) from .rms_layernorm import ( fast_rms_layernorm, @@ -25,7 +25,6 @@ from .rms_layernorm import ( from .layernorm import ( fast_layernorm, patch_layernorm, - unpatch_layernorm, ) from .rope_embedding import fast_rope_embedding, inplace_rope_embedding from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel @@ -54,8 +53,12 @@ from .flex_attention import ( create_flex_attention_sliding_window_mask, ) -try: - print("🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.") -except: - print("Unsloth: Will patch your computer to enable 2x faster free finetuning.") +import os +if "UNSLOTH_ZOO_IS_PRESENT" not in os.environ: + try: + print("🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.") + except: + print("Unsloth: Will patch your computer to enable 2x faster free finetuning.") + pass pass +del os diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index f2377d55cc..f0193c74d8 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -17,24 +17,31 @@ import triton.language as tl import torch from .utils import calculate_settings, MAX_FUSED_SIZE, triton_tanh from transformers.models.llama.modeling_llama import logger +from packaging.version import Version + +from unsloth_zoo.loss_utils import ( + patch_loss_functions as _patch_loss_functions, + post_patch_loss_function, +) @triton.heuristics({ - "DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING" ], - "DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"], + "DO_SOFTCAPPING": lambda args: bool(args["DO_SOFTCAPPING" ]), + "DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]), }) @triton.jit def _cross_entropy_forward( - logits_ptr, logits_row_stride, - loss_ptr, - logsumexp_ptr, - labels_ptr, - VOCAB_SIZE : tl.constexpr, - BLOCK_SIZE : tl.constexpr, - DO_SOFTCAPPING : tl.constexpr, - SOFTCAP : tl.constexpr, - DO_LOGIT_SCALING: tl.constexpr, - LOGIT_SCALE : tl.constexpr, + logits_ptr , + logits_row_stride , + loss_ptr , + logsumexp_ptr , + labels_ptr , + VOCAB_SIZE , + BLOCK_SIZE : tl.constexpr, + DO_SOFTCAPPING , + SOFTCAP , + DO_LOGIT_SCALING , + LOGIT_SCALE , ): """ Cross Entropy Loss = 1/n sum [ -yi log(Pi) ] @@ -57,7 +64,7 @@ def _cross_entropy_forward( This ensures exp(x - max(x))'s maximum is 1 as exp(0) = 1. """ row_idx = tl.program_id(0) - logits_ptr += row_idx * logits_row_stride.to(tl.int64) + logits_ptr += row_idx * tl.cast(logits_row_stride, tl.int64) loss_ptr += row_idx logsumexp_ptr += row_idx labels_ptr += row_idx @@ -71,7 +78,7 @@ def _cross_entropy_forward( # Go logit scaling for Cohere: t * x if DO_LOGIT_SCALING: logits = LOGIT_SCALE * logits # Do logit softcapping for Gemma 2: t * tanh(1/t * x) - if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits / SOFTCAP) + if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits.to(tl.float32) / SOFTCAP).to(logits.dtype) logits = logits.to(tl.float32) c = tl.max(logits, 0) @@ -92,22 +99,23 @@ pass @triton.heuristics({ - "DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING" ], - "DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"], + "DO_SOFTCAPPING": lambda args: bool(args["DO_SOFTCAPPING" ]), + "DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]), }) @triton.jit def _chunked_cross_entropy_forward( - logits_ptr, logits_row_stride, - loss_ptr, - logsumexp_ptr, - labels_ptr, - VOCAB_SIZE : tl.constexpr, - N_CHUNKS : tl.constexpr, - BLOCK_SIZE : tl.constexpr, - DO_SOFTCAPPING : tl.constexpr, - SOFTCAP : tl.constexpr, - DO_LOGIT_SCALING: tl.constexpr, - LOGIT_SCALE : tl.constexpr, + logits_ptr , + logits_row_stride , + loss_ptr , + logsumexp_ptr , + labels_ptr , + VOCAB_SIZE , + N_CHUNKS , + BLOCK_SIZE : tl.constexpr, + DO_SOFTCAPPING , + SOFTCAP , + DO_LOGIT_SCALING , + LOGIT_SCALE , ): """ 256K vocab divided in 4 chunks @@ -135,7 +143,7 @@ def _chunked_cross_entropy_forward( """ row_idx = tl.program_id(0) chunk_idx = tl.program_id(1) - logits_ptr += row_idx * logits_row_stride.to(tl.int64) + logits_ptr += row_idx * tl.cast(logits_row_stride, tl.int64) loss_ptr += row_idx logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx labels_ptr += row_idx @@ -149,7 +157,7 @@ def _chunked_cross_entropy_forward( # Go logit scaling for Cohere: t * x if DO_LOGIT_SCALING: logits = LOGIT_SCALE * logits # Do logit softcapping for Gemma 2: t * tanh(1/t * x) - if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits / SOFTCAP) + if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits.to(tl.float32) / SOFTCAP).to(logits.dtype) logits = logits.to(tl.float32) c = tl.max(logits, 0) @@ -174,21 +182,23 @@ pass @triton.heuristics({ - "DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING" ], - "DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"], + "DO_SOFTCAPPING": lambda args: bool(args["DO_SOFTCAPPING" ]), + "DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]), }) @triton.jit def _cross_entropy_backward( - logits_ptr, logits_row_stride, - dloss_ptr, dloss_row_stride, - logsumexp_ptr, - labels_ptr, - VOCAB_SIZE : tl.constexpr, - BLOCK_SIZE : tl.constexpr, - DO_SOFTCAPPING : tl.constexpr, - SOFTCAP : tl.constexpr, - DO_LOGIT_SCALING: tl.constexpr, - LOGIT_SCALE : tl.constexpr, + logits_ptr , + logits_row_stride , + dloss_ptr , + dloss_row_stride , + logsumexp_ptr , + labels_ptr , + VOCAB_SIZE , + BLOCK_SIZE : tl.constexpr, + DO_SOFTCAPPING , + SOFTCAP , + DO_LOGIT_SCALING , + LOGIT_SCALE , ): """ CE_i = -y log(P) = y * (log[sum(exp(x))] - x) @@ -208,7 +218,7 @@ def _cross_entropy_backward( row_idx = tl.program_id(0) block_idx = tl.program_id(1) - logits_ptr += row_idx * logits_row_stride.to(tl.int64) + logits_ptr += row_idx * tl.cast(logits_row_stride, tl.int64) dloss_ptr += row_idx * dloss_row_stride col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = col_offsets < VOCAB_SIZE @@ -228,9 +238,10 @@ def _cross_entropy_backward( pass # Do logit softcapping for Gemma 2: t * tanh(1/t * x) + partial = x if DO_SOFTCAPPING: # d/dx [t * tanh(1/t * x)] = 1 - tanh^2(1/t * x) - partial = triton_tanh(x / SOFTCAP) + partial = triton_tanh(x.to(tl.float32) / SOFTCAP).to(x.dtype) x = SOFTCAP * partial pass @@ -261,16 +272,20 @@ MAX_FUSED_SIZE = 65536 # 2**16 class Fast_CrossEntropyLoss(torch.autograd.Function): @staticmethod - def forward(ctx, logits, labels, logit_softcapping = 0, logit_scaling = 0): + def forward(ctx, logits, labels, logit_softcapping : float = 0, logit_scaling : float = 0): + n_rows : int + vocab_size : int n_rows, vocab_size = logits.shape div, mod = divmod(vocab_size, MAX_FUSED_SIZE) - n_chunks = div + (mod != 0) + n_chunks : int = div + (mod != 0) losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") - DO_SOFTCAPPING = (logit_softcapping != 0) - DO_LOGIT_SCALING = (logit_scaling != 0) + DO_SOFTCAPPING : bool = bool(logit_softcapping != 0) + DO_LOGIT_SCALING : bool = bool(logit_scaling != 0) + BLOCK_SIZE : int + num_warps : int if n_chunks == 1: # For small vocabs <= 65336 like Llama, Mistral BLOCK_SIZE, num_warps = calculate_settings(vocab_size) @@ -325,11 +340,13 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): @staticmethod def backward(ctx, dlosses): logits, logsumexp, labels = ctx.saved_tensors + n_rows : int + vocab_size : int n_rows, vocab_size = logits.shape - BLOCK_SIZE = 4096 + BLOCK_SIZE : int = 4096 div, mod = divmod(vocab_size, BLOCK_SIZE) - n_blocks = div + (mod != 0) + n_blocks : int = div + (mod != 0) _cross_entropy_backward[(n_rows, n_blocks,)]( logits, logits.stride(0), @@ -342,14 +359,13 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): SOFTCAP = ctx.logit_softcapping, DO_LOGIT_SCALING = ctx.DO_LOGIT_SCALING, LOGIT_SCALE = ctx.logit_scaling, - num_warps = 8, + num_warps = 8, ) return logits, None, None, None, pass pass -@torch._disable_dynamo def fast_cross_entropy_loss( logits, labels, @@ -377,96 +393,12 @@ def fast_cross_entropy_loss( n_items = torch.count_nonzero(labels != -100) return loss.sum() / n_items pass - - -from transformers.models.llama.modeling_llama import ( - LlamaForCausalLM, - CausalLMOutputWithPast, - Optional, - Union, - Cache, - List, - Tuple, -) - -# Transformers 4.47 need Unpack, KwargsForCausalLM -try: - from transformers.models.llama.modeling_llama import Unpack, KwargsForCausalLM -except: - pass +if (Version(torch.__version__) < Version("2.4.0")) and \ + not hasattr(fast_cross_entropy_loss, "__wrapped__"): + fast_cross_entropy_loss = torch._disable_dynamo(fast_cross_entropy_loss) pass -import inspect, re -function = inspect.getsource(LlamaForCausalLM.forward) -function = function.split("\n") -i = re.match(r"[ ]{1,}", function[0]).span(0)[1] -function = [x[i:] for x in function] -function = "\n".join(function) -function = function[function.find("def forward"):] -replacement = """ loss = None - logit_softcapping = getattr(self.config, "final_logit_softcapping", 0) - logit_scaling = getattr(self.config, "logit_scale", 0) - if labels is not None: - shift_logits = logits - if not hasattr(self, "extra_ignored_labels"): - # Fixes https://github.com/unslothai/unsloth/issues/10 - self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0") - pass - - shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]])) - loss = fast_cross_entropy_loss( - logits = shift_logits, - labels = shift_labels, - logit_softcapping = logit_softcapping, - logit_scaling = logit_scaling, - n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None), - ) - else: - if logit_scaling != 0: - if logits.requires_grad: - logits = logit_scaling * logits - else: - logits *= logit_scaling - pass - pass - if logit_softcapping != 0: - if logits.requires_grad: - logits = (1.0 / logit_softcapping) * logits - logits = torch.tanh(logits) - logits = logit_softcapping * logits - else: - logits *= (1.0 / logit_softcapping) - torch.tanh(logits, out = logits) - logits *= logit_softcapping - pass - pass - pass -""" -function = \ - function[:function.find(" loss = None")] + \ - replacement + \ - function[ function.find(" if not return_dict"):] -function = function.replace("logits = logits.float()", "\n") -# Missed spaces -function = function.split("\n") -# Not the first one though! -function = [function[0]] + [" "*4 + x for x in function[1:]] -function = "\n".join(function) -function = f"class Unsloth_LlamaForCausalLM(LlamaForCausalLM):\n"\ -f" {function}\n" -exec(function, globals()) -del function, replacement, inspect, re - - -def patch_llama_for_causal_lm(): - import transformers.models.llama.modeling_llama - transformers.models.llama.modeling_llama.LlamaForCausalLM = Unsloth_LlamaForCausalLM - return -pass - - -def unpatch_llama_for_causal_lm(): - import transformers.models.llama.modeling_llama - transformers.models.llama.modeling_llama.LlamaForCausalLM = LlamaForCausalLM - return +# Patch CE Losses in transformers +def patch_loss_functions(): + _patch_loss_functions(fast_cross_entropy_loss) pass diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py index 48ade6d5ec..a5f7926e2e 100644 --- a/unsloth/kernels/layernorm.py +++ b/unsloth/kernels/layernorm.py @@ -17,6 +17,9 @@ import triton import triton.language as tl import torch from .utils import calculate_settings +from unsloth_zoo.patching_utils import ( + patch_layernorm, +) @triton.jit @@ -162,27 +165,6 @@ def fast_layernorm(layernorm, X): pass -from torch.nn import LayerNorm -class Unsloth_LayerNorm(LayerNorm): - def forward(self, X): - return fast_layernorm(self, X) - pass -pass - - -def patch_layernorm(): - import torch.nn - torch.nn.LayerNorm = Unsloth_LayerNorm - return -pass - - -def unpatch_layernorm(): - import torch.nn - torch.nn.LayerNorm = LayerNorm - return -pass - def test_layernorm( dim = 1024, eps = 1e-5, dtype = torch.float16, diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index 13faf08d6a..4b22f8c3e5 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -53,14 +53,14 @@ def _rms_layernorm_forward( pass -@triton.heuristics({"GEMMA": lambda args: args["GEMMA"],}) +@triton.heuristics({"GEMMA": lambda args: bool(args["GEMMA"]),}) @triton.jit def _rms_layernorm_backward( dY, dY_row_stride, X, X_row_stride, W, W_row_stride, r, r_row_stride, - dW, dW_row_stride, + # dW, dW_row_stride, n_cols, eps, GEMMA : tl.constexpr, BLOCK_SIZE : tl.constexpr, @@ -130,11 +130,15 @@ pass class Fast_RMS_Layernorm(torch.autograd.Function): @staticmethod - def forward(ctx, X, W, eps, gemma = False): + def forward(ctx, X : torch.Tensor, W : torch.Tensor, eps : float, gemma : bool = False): shape = X.shape - dim = shape[-1] + dim : int = shape[-1] X = X.view(-1, dim) + n_rows : int + n_cols : int n_rows, n_cols = X.shape + BLOCK_SIZE : int + num_warps : int BLOCK_SIZE, num_warps = calculate_settings(n_cols) Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0") @@ -159,20 +163,22 @@ class Fast_RMS_Layernorm(torch.autograd.Function): pass @staticmethod - def backward(ctx, dY): + def backward(ctx, dY : torch.Tensor): shape = dY.shape - dim = shape[-1] + dim : int = shape[-1] dY = dY.view(-1, dim) X, W, r = ctx.saved_tensors + n_rows : int + n_cols : int n_rows, n_cols = dY.shape - dW = X + # dW = X _rms_layernorm_backward[(n_rows,)]( dY, dY.stride(0), X, X .stride(0), W, W .stride(0), r, r .stride(0), - dW, dW.stride(0), + # dW, dW.stride(0), n_cols, ctx.eps, GEMMA = ctx.GEMMA, BLOCK_SIZE = ctx.BLOCK_SIZE, @@ -184,9 +190,11 @@ class Fast_RMS_Layernorm(torch.autograd.Function): pass -def fast_rms_layernorm(layernorm, X, gemma = False): - W = layernorm.weight - eps = layernorm.variance_epsilon if \ +# [TODO] Unsure why RMS Layernorm is not torch.compiling properly +@torch.compiler.disable +def fast_rms_layernorm(layernorm, X : torch.Tensor, gemma : bool = False): + W : torch.Tensor = layernorm.weight + eps : float = layernorm.variance_epsilon if \ hasattr(layernorm, "variance_epsilon") \ else layernorm.eps out = Fast_RMS_Layernorm.apply(X, W, eps, gemma) diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index 2934ac41c9..7fe15d0e3b 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/unsloth/kernels/rope_embedding.py @@ -16,9 +16,9 @@ import triton import triton.language as tl import torch from .utils import calculate_settings -ROPE_GROUP_SIZE = 4 +ROPE_GROUP_SIZE : int = 4 -@triton.heuristics({"BACKWARD_PASS": lambda args: args["BACKWARD_PASS"],}) +@triton.heuristics({"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),}) @triton.jit def _rope_embedding( Q, Q_row_stride, @@ -75,8 +75,14 @@ class Fast_RoPE_Embedding(torch.autograd.Function): @staticmethod def forward(ctx, Q, cos, sin): cos, sin = cos.squeeze(), sin.squeeze() + batch : int + seq_len : int + n_heads : int + head_dim : int batch, seq_len, n_heads, head_dim = Q.shape Q = Q.view(batch*seq_len, n_heads*head_dim) + n_rows : int + n_cols : int n_rows, n_cols = Q.shape assert(seq_len <= cos.shape[0]) @@ -85,8 +91,10 @@ class Fast_RoPE_Embedding(torch.autograd.Function): BLOCK_SIZE, num_warps = calculate_settings(head_dim//2) # (head_dim//2) # group_size = 4 # 4 or 8, too large group_size can hurt performance. + div : int + mod : int div, mod = divmod(n_heads, ROPE_GROUP_SIZE) - n_groups = div + (mod != 0) + n_groups : int = div + (mod != 0) _rope_embedding[(n_rows, n_groups, )]( Q, Q.stride(0), @@ -108,9 +116,15 @@ class Fast_RoPE_Embedding(torch.autograd.Function): @staticmethod def backward(ctx, dY): + batch : int + seq_len : int + n_heads : int + head_dim : int batch, seq_len, n_heads, head_dim = dY.shape dY = dY.reshape(batch*seq_len, n_heads*head_dim) # Must be reshape not view + n_rows : int + n_cols : int n_rows, n_cols = dY.shape cos = ctx.cos @@ -130,7 +144,8 @@ class Fast_RoPE_Embedding(torch.autograd.Function): pass pass - +# [TODO] Unsure why RoPE Embedding is not torch.compiling properly +@torch.compiler.disable def fast_rope_embedding(Q, K, cos, sin): Q = Fast_RoPE_Embedding.apply(Q.transpose(1, 2), cos, sin).transpose(1, 2) K = Fast_RoPE_Embedding.apply(K.transpose(1, 2), cos, sin).transpose(1, 2) diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index a8c20c75a4..b394d122fd 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -13,7 +13,7 @@ # limitations under the License. import triton -MAX_FUSED_SIZE = 65536 +MAX_FUSED_SIZE : int = 65536 next_power_of_2 = triton.next_power_of_2 # torch.cuda.amp.custom_fwd is deprecated >= 2.4 @@ -40,12 +40,12 @@ else: pass -def calculate_settings(n): - BLOCK_SIZE = next_power_of_2(n) +def calculate_settings(n : int) -> (int, int,): + BLOCK_SIZE : int = next_power_of_2(n) if BLOCK_SIZE > MAX_FUSED_SIZE: raise RuntimeError(f"Cannot launch Triton kernel since n = {n} exceeds "\ f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.") - num_warps = 4 + num_warps : int = 4 if BLOCK_SIZE >= 32768: num_warps = 32 elif BLOCK_SIZE >= 8192: num_warps = 16 elif BLOCK_SIZE >= 2048: num_warps = 8 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 51c63fc7dd..94cf1b74e0 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2024.10.7" +__version__ = "2024.11.1" __all__ = [ "prepare_model_for_kbit_training", @@ -41,9 +41,17 @@ __all__ = [ "torch_amp_custom_bwd", "accelerate_old_send_to_device", "accelerate_new_send_to_device", + "patch_gradient_accumulation_fix", + "patch_compiling_bitsandbytes", + "patch_regional_compilation", + "patch_layernorm", + "patch_torch_compile", + "patch_model_and_tokenizer", + + "patch_unsloth_gradient_checkpointing", + "unpatch_unsloth_gradient_checkpointing", "patch_gradient_checkpointing", "unpatch_gradient_checkpointing", - "patch_gradient_accumulation_fix", ] import torch @@ -54,6 +62,28 @@ import numpy as np import warnings, subprocess, re, inspect, psutil, os, math from packaging.version import Version +from unsloth_zoo.tokenizer_utils import ( + patch_tokenizer as _patch_tokenizer, +) +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 ( + Unsloth_Offloaded_Gradient_Checkpointer, + unsloth_offloaded_gradient_checkpoint, + patch_unsloth_gradient_checkpointing, + unpatch_unsloth_gradient_checkpointing, + + Unsloth_Gradient_Checkpointer, + unsloth_gradient_checkpoint, + patch_gradient_checkpointing, + unpatch_gradient_checkpointing, +) + # ============================================= # Disable some warnings which can get annoying warnings.filterwarnings(action = "ignore", category = UserWarning, module = "torch") @@ -70,6 +100,18 @@ warnings.filterwarnings(action = "ignore", category = RuntimeWarning, module = " # Stop "Special tokens have been added in the vocabulary, ..." import logging logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITICAL+1) + +# Ignore logging messages +class HideLoggingMessage(logging.Filter): + def __init__(self, text): self.text = text + def filter(self, x): return not x.getMessage().startswith(self.text) +pass + +# The speedups for torchdynamo mostly come wih GPU Ampere or higher and which is not detected here. +from transformers.training_args import logger as transformers_training_args_logger +transformers_training_args_logger.addFilter(HideLoggingMessage("The speedups")) +del transformers_training_args_logger + # ============================================= # ============================================= @@ -129,7 +171,6 @@ pass # ============================================= # torch.cuda.amp.custom_fwd is deprecated >= 2.4 -import torch torch_version = torch.__version__ if Version(torch_version) < Version("2.4.0"): torch_amp_custom_fwd = torch.cuda.amp.custom_fwd @@ -333,7 +374,8 @@ pass # ============================================= # Torch compile settings - +UNSLOTH_COMPILE_DEBUG = "UNSLOTH_COMPILE_DEBUG" in os.environ +UNSLOTH_COMPILE_MAXIMUM = "UNSLOTH_COMPILE_MAXIMUM" in os.environ # Just remove max_autotune_gemm warning import functools @functools.lru_cache(None) @@ -345,47 +387,27 @@ def is_big_gpu(index): return True import torch._inductor.utils torch._inductor.utils.is_big_gpu = is_big_gpu +patch_torch_compile(debug = UNSLOTH_COMPILE_DEBUG, O3 = UNSLOTH_COMPILE_MAXIMUM) - -# Torch compile arguments -torch_compile_arguments = [ - "config.dce = True", - "config.memory_planning = True", - "config.memory_pool = 'combined'", - "config.coordinate_descent_tuning = True", - "config.max_autotune_gemm = False", # GEMM is unnecessary - "config.autotune_multi_device = False", - "config.max_autotune_gemm_backends = 'TRITON,ATEN,CPP'", # Not much faster - "config.aggressive_fusion = False", # Careful changes results! - "config.cuda.enable_cuda_lto = True", - "config.cuda.use_fast_math = True", - "config.cuda.compile_opt_level = '-O2'", -] -# Torch dynamo arguments -torch_dynamo_arguments = [ - "config.accumulated_cache_size_limit = 1024", # Bump up a bit from 256 - "config.suppress_errors = True", # Supress errors for now - "config.do_not_emit_runtime_asserts = True", - "config.cache_size_limit = 1024", # Flex Attention - "config.inline_inbuilt_nn_modules = True", # Torch 2.5 Regional recompilation -] -import torch._inductor.config as config -for _try_compile_argument in torch_compile_arguments: - try: exec(_try_compile_argument) - except: pass -pass -import torch._dynamo.config as config -for _try_dynamo_argument in torch_dynamo_arguments: - try: exec(_try_dynamo_argument) - except: pass -pass torch_compile_options = { "epilogue_fusion" : True, "max_autotune" : True, "shape_padding" : True, - "trace.enabled" : False, # Output Triton kernel outputs! + "trace.enabled" : UNSLOTH_COMPILE_DEBUG, "triton.cudagraphs" : False, } + +import accelerate +def torch_compile_kwargs(*args, **kwargs): + print("Unsloth: Enabled auto compiling") + return {"dynamic" : True, "fullgraph" : False, "options" : torch_compile_options,} +pass + +accelerate.utils.dataclasses.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs +accelerate.utils.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs +accelerate.accelerator.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs +del accelerate + # ============================================= def prepare_model_for_kbit_training( @@ -455,137 +477,6 @@ def prepare_model_for_kbit_training( return model pass - -def patch_tokenizer(model, tokenizer): - """ - Phi3's pad_token isn't set. We set it to <|placeholder... - Llama-3 is <|reserved... - Llama-2 is - Check if pad_token is not the same as eos_token otherwise the loss will ignore it!! - Fixes https://github.com/unslothai/unsloth/issues/5 - """ - possible_reserved_tokens = ( - "<|finetune_right_pad_id|>", # Llama-3.1 - "", # Mistral Nemo - "<|reserved", # Llama-3 - "<|placeholder", # Phi-3 - "[control", # Mistral type models - ) - joiner = "\1\0=+=\0\1" - number_repetitions = 3 - 1 # Number of reserved tokens needed - - if model is not None: - model.config.update({"unsloth_version" : __version__}) - - bad_pad_token = False - if hasattr(tokenizer, "pad_token") and tokenizer.pad_token is not None: - # Check if pad_token is not the same as eos_token otherwise the loss will ignore it!! - bad_pad_token = tokenizer.eos_token == tokenizer.pad_token - elif hasattr(tokenizer, "pad_token") and tokenizer.pad_token is None: - bad_pad_token = True - else: - bad_pad_token = False - pass - - if bad_pad_token: - # Find a better pad token - added_tokens = [str(x) for x in tokenizer.added_tokens_decoder.values()] - all_added_tokens = joiner.join(added_tokens[::-1]) - all_added_tokens += joiner - - final_pad_token = None - final_good_match = False - - for possible_reserved_token in possible_reserved_tokens: - possible_reserved_token = re.escape(possible_reserved_token) - found = re.finditer(f"{possible_reserved_token}", all_added_tokens) - first_match = None - good_match = False - for j, x in enumerate(found): - if j == 0: first_match = x - if j >= number_repetitions: - good_match = True - break - pass - pass - - if first_match is None: continue - - # If it ends with |> or > etc, then set it as a good pad token! - start = first_match.span(0)[0] - possible_pad_token = first_match.group(0) - end = all_added_tokens.find(joiner, start) - first_match = all_added_tokens[start:end] - - if first_match is not None: - good_match = possible_pad_token.endswith((">", "|>", "]", ")")) - pass - possible_pad_token = first_match - - # Replace current pad token if another exact match is found - if not final_good_match and good_match: - final_good_match = True - final_pad_token = possible_pad_token - break - else: - final_good_match = False - final_pad_token = possible_pad_token - pass - pass - possible_pad_token = final_pad_token - - # Try unk_token - if possible_pad_token is None and hasattr(tokenizer, "unk_token"): - possible_pad_token = tokenizer.unk_token - pass - - # Check pad token's id must be less than vocab size - if possible_pad_token is not None: - check_pad_token = tokenizer(possible_pad_token, add_special_tokens = False).input_ids - if len(check_pad_token) != 1: - possible_pad_token = None - if model is not None and check_pad_token[0] >= model.config.vocab_size: - possible_pad_token = None - pass - - if possible_pad_token is None: - # Failure to find a good replacement!! We shall manually add one! - new_pad_token = "<|PAD_TOKEN|>" - while new_pad_token in tokenizer.get_vocab(): - new_pad_token = f"<{new_pad_token}>" - pass - possible_pad_token = new_pad_token - pass - - name = model.config._name_or_path if model is not None else "Model" - logger.warning_once( - f"{name} does not have a padding token! Will use pad_token = {possible_pad_token}." - ) - - # Edit pad_token - tokenizer.add_special_tokens({"pad_token" : possible_pad_token}) - tokenizer.pad_token = possible_pad_token - if model is not None: - model.config.update({"pad_token_id" : tokenizer.pad_token_id}) - if getattr(model, "generation_config") is not None: - model.generation_config.update(pad_token_id = tokenizer.pad_token_id) - else: - if model is not None: - if model.config.pad_token_id is None: - model.config.update({"pad_token_id" : tokenizer.pad_token_id}) - if getattr(model, "generation_config") is not None: - model.generation_config.update(pad_token_id = tokenizer.pad_token_id) - pass - pass - - if model is not None: - if getattr(model, "generation_config") is not None: - model.generation_config.update(max_length = model.config.max_position_embeddings) - - return model, tokenizer -pass - - # ============================================= # Weirdly LoraLayer.update_layer downcasts PEFT layers to float16?? # For mixed precision, we need it to be in float32 not float16. @@ -618,6 +509,7 @@ if Version(peft_version) < Version("0.12.0"): ) pass pass + # ============================================= import psutil @@ -678,7 +570,9 @@ def get_statistics(): # We log some basic stats about which environment is being used. # We simply download a README.md file from HF - all data is made public. # This is simply so we can check if some envs are broken or not. - # You can disable this by commenting the below out + # You can disable this by setting UNSLOTH_DISABLE_STATISTICS + import os + if "UNSLOTH_DISABLE_STATISTICS" in os.environ: return from huggingface_hub.utils import disable_progress_bars, enable_progress_bars, are_progress_bars_disabled disabled = False if not are_progress_bars_disabled(): @@ -710,139 +604,6 @@ def get_statistics(): pass -def _calculate_n_gradient_checkpoints( - n_layers : int, - method : Optional[Union[str, int]] = "sqrt", -) -> List[int]: - assert(type(n_layers) is int and n_layers > 0) - - if method is None: method = "sqrt" - - if method == "sqrt": - n_checkpoints = int(n_layers**0.5) - elif type(method) is int and method > 0: - n_checkpoints = int(np.ceil(n_layers / method)) - else: - raise ValueError("method must be 'sqrt' or an int >0 and <= n_layers.") - - size = n_layers // n_checkpoints - sizes = np.full(n_checkpoints, size, dtype = int) - leftovers = n_layers % n_checkpoints - # We append leftovers from the right - for k in range(leftovers): - sizes[n_checkpoints-1-k] += 1 - boundaries = np.hstack((0, np.cumsum(sizes))) - boundaries = boundaries.tolist() - return boundaries -pass - - -def calculate_n_gradient_checkpoints( - n_layers : int, - layers_per_checkpoint : Optional[Union[str, int]] = "sqrt", -) -> List[int]: - assert(type(n_layers) is int and n_layers > 0) - - if layers_per_checkpoint is None or layers_per_checkpoint == 1: - return None - - boundaries = _calculate_n_gradient_checkpoints(n_layers, layers_per_checkpoint) - - assert(boundaries[0] == 0 and boundaries[-1] == n_layers) - assert(min(boundaries) == 0 and max(boundaries) == n_layers) - assert(np.diff(boundaries).min() >= 0) - return boundaries -pass - - -def prepare_n_gradient_checkpoints( - model : Any, - layers_per_checkpoint : Optional[Union[str, int]] = "sqrt", - use_reentrant : Optional[bool] = True, -) -> None: - """ - Calculates where to place the gradient checkpoints given n_layers. - - Args: - model: Any LlamaModel with layers. - layers_per_checkpoint (`Union[str, int]`, *optional*): - Can either be `sqrt` or an integer for how many layers per checkpoint you want. - The more, the less memory usage, but can be slower. Default is `sqrt`. - Choose 1 for Pytorch gradient checkpointing. 2 to wrap 2 layers in 1 module etc. - use_reentrant (`bool`, *optional*): - https://github.com/pytorch/pytorch/blob/main/torch/utils/checkpoint.py#L354 - Optimal gradient checkpointing algorithm `use_reentrant=False` which will - be the default in future Pytorch versions doesn't seem to work?? - """ - _model = None - if hasattr(model, "layers"): - _model = model - elif hasattr(model, "model"): - if hasattr(model.model, "layers"): - _model = model.model - if _model is None: - raise TypeError("`model` or `model.model` does not have attribute `layers`. Are you sure this is a model?") - pass - - if use_reentrant is False: - use_reentrant = True - pass - - n_layers = len(_model.layers) - boundaries = calculate_n_gradient_checkpoints(n_layers, layers_per_checkpoint) - _model._gradient_checkpointing_boundaries = boundaries - _model._gradient_checkpointing_use_reentrant = use_reentrant -pass - - -class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function): - """ - Saves VRAM by smartly offloading to RAM. - Tiny hit to performance, since we mask the movement via non blocking calls. - """ - @staticmethod - @torch_amp_custom_fwd - def forward(ctx, forward_function, hidden_states, *args): - saved_hidden_states = hidden_states.to("cpu", non_blocking = True) - with torch.no_grad(): - output = forward_function(hidden_states, *args) - ctx.save_for_backward(saved_hidden_states) - ctx.forward_function = forward_function - ctx.args = args - return output - pass - - @staticmethod - @torch_amp_custom_bwd - def backward(ctx, dY): - (hidden_states,) = ctx.saved_tensors - hidden_states = hidden_states.to("cuda:0", non_blocking = True).detach() - hidden_states.requires_grad_(True) - with torch.enable_grad(): - (output,) = ctx.forward_function(hidden_states, *ctx.args) - torch.autograd.backward(output, dY) - return (None, hidden_states.grad,) + (None,)*len(ctx.args) - pass -pass - - -@torch._disable_dynamo -def unsloth_offloaded_gradient_checkpoint(function, *args, use_reentrant = None, **kwargs): - 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 -pass - -def unpatch_gradient_checkpointing(): - torch.utils.checkpoint = old_checkpoint -pass - - # ============================================= # Fixes Bitsandbytes to remove missing warnings from transformers.utils.quantization_config import BitsAndBytesConfig, QuantizationMethod @@ -1189,6 +950,7 @@ def patch_gradient_accumulation_fix(Trainer): # Fixes gradient accumulation import inspect if hasattr(Trainer, "get_batch_samples"): + if Trainer.get_batch_samples.__name__ == "_unsloth_get_batch_samples": return if \ not inspect.getsource(Trainer.get_batch_samples).strip()\ .endswith("return batch_samples, num_items_in_batch"): @@ -1215,6 +977,7 @@ def patch_gradient_accumulation_fix(Trainer): pass # Also fix up loss scaling ie negate loss *= self.args.gradient_accumulation_steps + if Trainer.training_step.__name__ == "_unsloth_training_step": return if "num_items_in_batch" not in inspect.signature(Trainer.training_step).parameters: return function = inspect.getsource(Trainer.training_step) @@ -1243,3 +1006,11 @@ def patch_gradient_accumulation_fix(Trainer): exec(function, globals()) Trainer.training_step = _unsloth_training_step pass + + +def patch_tokenizer(model, tokenizer): + model, tokenizer = _patch_tokenizer(model, tokenizer) + if model is not None: + model.config.update({"unsloth_version" : __version__}) + return model, tokenizer +pass diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 45f14c1131..1d9a0c1334 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -339,60 +339,9 @@ class FastGemmaModel(FastLlamaModel): @staticmethod - def post_patch(model): - # Patch model for Gemma - layers = model.model.layers - - # Torch.compile fails on embedding matrix?? - # Workaround randomnly fixes it for torch versions < 2.2 - model.model.embed_tokens = torch.nn.Embedding.from_pretrained(model.model.embed_tokens.weight) - model.config.update({"unsloth_version" : __version__}) - - # We also do this for the lm_head - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - lm_head.weight = model.lm_head.weight - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - model.lm_head = lm_head - - # Gemma has tied weights! This means lm_head == embed_tokens - if model.model.embed_tokens.weight.data_ptr() != model.lm_head.weight.data_ptr(): - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - lm_head.weight = model.model.embed_tokens.weight - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - model.lm_head = lm_head - pass - - # Also patch all dtypes - BnB seems to not allocate the correct type? - # BnB default dtype seems to be float16! - correct_dtype = lm_head.weight.dtype - - for name, module in model.named_modules(): - if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)): - weight = module.weight - quant_state = weight.quant_state - - if type(quant_state) is list: - # BnB seems to have float16 as default! - module.weight.quant_state[2] = correct_dtype # Cast to correct dtype - else: - # https://github.com/TimDettmers/bitsandbytes/pull/763/files - quant_state.dtype = correct_dtype - pass - pass - # Downcast RoPE embedding to correct data type - # RoPE must be done in float32 for Gemma - # if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \ - # and (module.cos_cached.dtype != correct_dtype): - - # module.cos_cached = module.cos_cached.to(correct_dtype) - # module.sin_cached = module.sin_cached.to(correct_dtype) - # pass - # pass - pass + def post_patch(model, tokenizer): + # Gemma does not downcast RoPE + model, tokenizer = patch_model_and_tokenizer(model, tokenizer, downcast_rope = False) # Add 1 to weight # return output * (1 + self.weight) @@ -425,6 +374,6 @@ class FastGemmaModel(FastLlamaModel): for _ in range(3): gc.collect() torch.cuda.empty_cache() - return model + return model, tokenizer pass pass diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index bf40ea8a27..4eb9d64313 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -490,60 +490,9 @@ class FastGemma2Model(FastLlamaModel): @staticmethod - def post_patch(model): - # Patch model for Gemma - layers = model.model.layers - - # Torch.compile fails on embedding matrix?? - # Workaround randomnly fixes it for torch versions < 2.2 - model.model.embed_tokens = torch.nn.Embedding.from_pretrained(model.model.embed_tokens.weight) - model.config.update({"unsloth_version" : __version__}) - - # We also do this for the lm_head - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - lm_head.weight = model.lm_head.weight - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - model.lm_head = lm_head - - # Gemma has tied weights! This means lm_head == embed_tokens - if model.model.embed_tokens.weight.data_ptr() != model.lm_head.weight.data_ptr(): - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - lm_head.weight = model.model.embed_tokens.weight - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - model.lm_head = lm_head - pass - - # Also patch all dtypes - BnB seems to not allocate the correct type? - # BnB default dtype seems to be float16! - correct_dtype = lm_head.weight.dtype - - for name, module in model.named_modules(): - if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)): - weight = module.weight - quant_state = weight.quant_state - - if type(quant_state) is list: - # BnB seems to have float16 as default! - module.weight.quant_state[2] = correct_dtype # Cast to correct dtype - else: - # https://github.com/TimDettmers/bitsandbytes/pull/763/files - quant_state.dtype = correct_dtype - pass - pass - # Downcast RoPE embedding to correct data type - # RoPE must be done in float32 for Gemma - # if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \ - # and (module.cos_cached.dtype != correct_dtype): - - # module.cos_cached = module.cos_cached.to(correct_dtype) - # module.sin_cached = module.sin_cached.to(correct_dtype) - # pass - # pass - pass + def post_patch(model, tokenizer): + # Gemma does not downcast RoPE + model, tokenizer = patch_model_and_tokenizer(model, tokenizer, downcast_rope = False) # Add 1 to weight # return output * (1 + self.weight) @@ -576,6 +525,6 @@ class FastGemma2Model(FastLlamaModel): for _ in range(3): gc.collect() torch.cuda.empty_cache() - return model + return model, tokenizer pass pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index cf05d432c7..3c4d8f3b38 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -57,8 +57,6 @@ from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING from transformers import set_seed as transformers_set_seed from peft import LoraConfig, TaskType, get_peft_model as _get_peft_model from peft import PeftModelForCausalLM -from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit -from peft.tuners.lora import Linear4bit as Peft_Linear4bit from ..save import patch_saving_functions import re, os, inspect, math, sys try: @@ -1518,6 +1516,7 @@ class FastLlamaModel: pass # Return old flag os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" model_patcher.pre_patch() get_statistics() # For debugging - we use a download counter to see if environments are not breaking @@ -1621,7 +1620,7 @@ class FastLlamaModel: ) model, tokenizer = patch_tokenizer(model, tokenizer) - model = model_patcher.post_patch(model) + model, tokenizer = model_patcher.post_patch(model, tokenizer) # Patch up QKV / O and MLP for idx, layer in enumerate(model.model.layers): @@ -1797,93 +1796,15 @@ class FastLlamaModel: internal_model = internal_model.model pass internal_model._saved_temp_tokenizer = tokenizer - - # Also fix torch_dtype - internal_model = model - while hasattr(internal_model, "model"): - if hasattr(internal_model, "config"): - if internal_model.config.torch_dtype == "float32": - internal_model.config.torch_dtype = torch.float32 - elif internal_model.config.torch_dtype == "bfloat16": - internal_model.config.torch_dtype = torch.bfloat16 - elif internal_model.config.torch_dtype == "float16": - internal_model.config.torch_dtype = torch.float16 - pass - pass - internal_model = internal_model.model - pass - if hasattr(internal_model, "config"): - if internal_model.config.torch_dtype == "float32": - internal_model.config.torch_dtype = torch.float32 - elif internal_model.config.torch_dtype == "bfloat16": - internal_model.config.torch_dtype = torch.bfloat16 - elif internal_model.config.torch_dtype == "float16": - internal_model.config.torch_dtype = torch.float16 - pass - pass return model, tokenizer pass @staticmethod - def post_patch(model): - # Patch model - layers = model.model.layers - - # Torch.compile fails on embedding matrix?? - # Workaround randomnly fixes it for torch versions < 2. - model.set_input_embeddings(torch.nn.Embedding.from_pretrained(model.get_input_embeddings().weight)) - model.config.update({"unsloth_version" : __version__}) - - # We also do this for the lm_head - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - lm_head.weight = model.get_output_embeddings().weight - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - model.lm_head = lm_head - - # Also patch all dtypes - BnB seems to not allocate the correct type? - # BnB default dtype seems to be float16! - correct_dtype = lm_head.weight.dtype - - for name, module in model.named_modules(): - if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)): - weight = module.weight - quant_state = weight.quant_state - - if type(quant_state) is list: - # BnB seems to have float16 as default! - module.weight.quant_state[2] = correct_dtype # Cast to correct dtype - else: - # https://github.com/TimDettmers/bitsandbytes/pull/763/files - quant_state.dtype = correct_dtype - pass - pass - # Downcast RoPE embedding to correct data type - if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")): - - if hasattr(module, "cos_cached") and \ - (module.cos_cached.dtype != correct_dtype): - - module.cos_cached = module.cos_cached.to(correct_dtype) - module.sin_cached = module.sin_cached.to(correct_dtype) - - elif hasattr(module, "short_cos_cached") and \ - (module.short_cos_cached.dtype != correct_dtype): - - module.short_cos_cached = module.short_cos_cached.to(correct_dtype) - module.short_sin_cached = module.short_sin_cached.to(correct_dtype) - pass - pass - pass - - # Clear deleted GPU items - for _ in range(3): - gc.collect() - torch.cuda.empty_cache() - return model + def post_patch(model, tokenizer): + model, tokenizer = patch_model_and_tokenizer(model, tokenizer, downcast_rope = True) + return model, tokenizer pass @@ -1910,6 +1831,11 @@ class FastLlamaModel: ): transformers_set_seed(random_state) + if type(r) is not int: + raise TypeError(f"Unsloth: Rank of {str(r)} must be an integer.") + if r <= 0: + raise TypeError(f"Unsloth: Rank of {str(r)} must be larger than 0.") + if isinstance(model, PeftModelForCausalLM): # Check if exactly the same and then pass through! assert(hasattr(model, "peft_config")) diff --git a/unsloth/save.py b/unsloth/save.py index ccda79aeee..b4c6b499cf 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -49,6 +49,7 @@ __all__ = [ keynames = "\n" + "\n".join(os.environ.keys()) IS_COLAB_ENVIRONMENT = "\nCOLAB_" in keynames IS_KAGGLE_ENVIRONMENT = "\nKAGGLE_" in keynames +KAGGLE_TMP = "/tmp" del keynames # Weights @@ -447,13 +448,20 @@ def unsloth_save_model( if push_to_hub and "/" in save_directory: # +1 solves absolute path issues - username = save_directory[:save_directory.find("/")] - new_save_directory = save_directory[save_directory.find("/")+1:] - - logger.warning_once( - f"Unsloth: You are pushing to hub, but you passed your HF username = {username}.\n"\ - f"We shall truncate {save_directory} to {new_save_directory}" - ) + new_save_directory = save_directory + username = new_save_directory[:new_save_directory.find("/")] + new_save_directory = new_save_directory[new_save_directory.find("/")+1:] + if IS_KAGGLE_ENVIRONMENT: + new_save_directory = os.path.join(KAGGLE_TMP, new_save_directory[new_save_directory.find("/")+1:]) + logger.warning_once( + "Unsloth: You are pushing to hub in Kaggle environment.\n"\ + f"To save memory, we shall move {save_directory} to {new_save_directory}" + ) + else: + logger.warning_once( + f"Unsloth: You are pushing to hub, but you passed your HF username = {username}.\n"\ + f"We shall truncate {save_directory} to {new_save_directory}" + ) save_pretrained_settings["save_directory"] = new_save_directory tokenizer_save_settings ["save_directory"] = new_save_directory @@ -507,6 +515,10 @@ def unsloth_save_model( f"{round(max_ram/1024/1024/1024, 2)} out of "\ f"{round(psutil.virtual_memory().total/1024/1024/1024, 2)} RAM for saving.") + # Move temporary_location to /tmp in Kaggle + if IS_KAGGLE_ENVIRONMENT: + temporary_location = os.path.join(KAGGLE_TMP, temporary_location) + # Max directory for disk saving if not os.path.exists(temporary_location): os.makedirs(temporary_location) @@ -708,7 +720,7 @@ def unsloth_save_model( print("Done.") if push_to_hub and hasattr(model, "config"): - print(f"Saved merged model to https://huggingface.co/{username}/{save_directory.lstrip('/')}") + print(f"Saved merged model to https://huggingface.co/{username}/{save_directory.lstrip('/').split('/')[-1]}") pass save_pretrained_settings["state_dict"] = None @@ -1108,14 +1120,17 @@ def save_to_gguf( # Check if quantization succeeded! if not os.path.isfile(final_location): if IS_KAGGLE_ENVIRONMENT: - raise RuntimeError( - f"Unsloth: Quantization failed for {final_location}\n"\ - "You are in a Kaggle environment, which might be the reason this is failing.\n"\ - "Kaggle only provides 20GB of disk space. Merging to 16bit for 7b models use 16GB of space.\n"\ - "This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"\ - "`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"\ - "I suggest you to save the 16bit model first, then use manual llama.cpp conversion." - ) + if not Path(final_location).resolve().is_relative_to(Path('/tmp').resolve()): + raise RuntimeError( + f"Unsloth: Quantization failed for {final_location}\n"\ + "You are in a Kaggle environment, which might be the reason this is failing.\n"\ + "Kaggle only provides 20GB of disk space in the working directory.\n"\ + "Merging to 16bit for 7b models use 16GB of space.\n"\ + "This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"\ + "`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"\ + "You can try saving it to the `/tmp` directory for larger disk space.\n"\ + "I suggest you to save the 16bit model first, then use manual llama.cpp conversion." + ) else: raise RuntimeError( f"Unsloth: Quantization failed for {final_location}\n"\ @@ -1156,14 +1171,17 @@ def save_to_gguf( # Check if quantization succeeded! if not os.path.isfile(final_location): if IS_KAGGLE_ENVIRONMENT: - raise RuntimeError( - f"Unsloth: Quantization failed for {final_location}\n"\ - "You are in a Kaggle environment, which might be the reason this is failing.\n"\ - "Kaggle only provides 20GB of disk space. Merging to 16bit for 7b models use 16GB of space.\n"\ - "This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"\ - "`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"\ - "I suggest you to save the 16bit model first, then use manual llama.cpp conversion." - ) + if not Path(final_location).resolve().is_relative_to(Path('/tmp').resolve()): + raise RuntimeError( + f"Unsloth: Quantization failed for {final_location}\n"\ + "You are in a Kaggle environment, which might be the reason this is failing.\n"\ + "Kaggle only provides 20GB of disk space in the working directory.\n"\ + "Merging to 16bit for 7b models use 16GB of space.\n"\ + "This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"\ + "`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"\ + "You can try saving it to the `/tmp` directory for larger disk space.\n"\ + "I suggest you to save the 16bit model first, then use manual llama.cpp conversion." + ) else: raise RuntimeError( "Unsloth: Quantization failed! You might have to compile llama.cpp yourself, then run this again.\n"\ diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index c05485f902..c639dbf1a0 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -64,6 +64,7 @@ IGNORED_TOKENIZER_NAMES = frozenset( keynames = "\n" + "\n".join(os.environ.keys()) IS_COLAB_ENVIRONMENT = "\nCOLAB_" in keynames IS_KAGGLE_ENVIRONMENT = "\nKAGGLE_" in keynames +KAGGLE_TMP = "/tmp" del keynames @@ -470,8 +471,12 @@ def _load_correct_tokenizer( cache_dir = "huggingface_tokenizers_cache", fix_tokenizer = True, ): - if IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT: + if IS_COLAB_ENVIRONMENT: cache_dir = cache_dir + elif IS_KAGGLE_ENVIRONMENT: + # /tmp of Kaggle seems has a 80GB limit! + # Let's utilize them + cache_dir = os.path.join(KAGGLE_TMP, cache_dir) else: cache_dir = None pass