From 379d8101a5ff73169e345d722856e7f0fedacf81 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Sep 2024 11:48:24 -0700 Subject: [PATCH 1/4] Llama 3.2 (#1058) * Layernorm * Update layernorm.py * Update layernorm.py * Update layernorm.py * Update layernorm.py * Update layernorm.py * Update layernorm.py * Patch layernorm * Update layernorm.py * RMS Layernorm * Update rms_layernorm.py * Causal LM * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update layernorm.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update _utils.py * Update _utils.py * Llama 3.2 --- unsloth/kernels/__init__.py | 17 +- unsloth/kernels/cross_entropy_loss.py | 84 ++++ unsloth/kernels/layernorm.py | 231 ++++++++++ unsloth/kernels/rms_layernorm.py | 89 ++++ unsloth/models/_utils.py | 17 +- unsloth/models/mapper.py | 16 + unsloth/models/vision.py | 592 ++++++++++++++++++++++++++ 7 files changed, 1042 insertions(+), 4 deletions(-) create mode 100644 unsloth/kernels/layernorm.py create mode 100644 unsloth/models/vision.py diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py index cd1d90f262..3e55332c80 100644 --- a/unsloth/kernels/__init__.py +++ b/unsloth/kernels/__init__.py @@ -12,8 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .cross_entropy_loss import fast_cross_entropy_loss -from .rms_layernorm import fast_rms_layernorm +from .cross_entropy_loss import ( + fast_cross_entropy_loss, + patch_llama_for_causal_lm, + unpatch_llama_for_causal_lm, +) +from .rms_layernorm import ( + fast_rms_layernorm, + patch_rms_layernorm, + unpatch_rms_layernorm, +) +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 from .geglu import ( diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 24e8002bec..1fec5d7a85 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -375,3 +375,87 @@ 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, +) +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, + ) + 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 +pass diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py new file mode 100644 index 0000000000..48ade6d5ec --- /dev/null +++ b/unsloth/kernels/layernorm.py @@ -0,0 +1,231 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# Copyright 2024-present Andrej Karpathy & the llm.c team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import triton +import triton.language as tl +import torch +from .utils import calculate_settings + + +@triton.jit +def layernorm_forward( + Y, Y_row_stride, + X, X_row_stride, + W, + b, + r, + mu, + n_cols, eps, + BLOCK_SIZE : tl.constexpr +): + row_idx = tl.program_id(0) + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < n_cols + + Y += row_idx * Y_row_stride + X += row_idx * X_row_stride + r += row_idx + mu += row_idx + + # According to https://pytorch.org/torchtune/stable/_modules/torchtune/modules/layer_norm.html#Fp32LayerNorm, all modules + # are in float32! + X_row = tl.load(X + col_offsets, mask = mask, other = 0).to(tl.float32) + W_row = tl.load(W + col_offsets, mask = mask, other = 0).to(tl.float32) + b_row = tl.load(b + col_offsets, mask = mask, other = 0).to(tl.float32) + + mean_X = tl.sum(X_row, axis = 0) / n_cols + XX = X_row - mean_X + row_var = tl.sum(XX * XX, axis = 0) / n_cols + inv_var = tl.math.rsqrt(row_var + eps) + tl.store (r, inv_var) + tl.store (mu, mean_X) + output = (XX * inv_var) * W_row + b_row + tl.store(Y + col_offsets, output, mask = mask) +pass + + +@triton.jit +def layernorm_backward( + dY, dY_row_stride, + X, X_row_stride, + W, + b, + r, + mu, + n_cols, eps, + BLOCK_SIZE : tl.constexpr +): + # Approximately follows https://github.com/karpathy/llm.c/blob/master/doc/layernorm/layernorm.md + row_idx = tl.program_id(0) + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < n_cols + + dY += row_idx * dY_row_stride + X += row_idx * X_row_stride + r += row_idx + mu += row_idx + + # According to https://pytorch.org/torchtune/stable/_modules/torchtune/modules/layer_norm.html#Fp32LayerNorm, all modules + # are in float32! + dY_row = tl.load(dY + col_offsets, mask = mask, other = 0).to(tl.float32) + X_row = tl.load(X + col_offsets, mask = mask, other = 0).to(tl.float32) + W_row = tl.load(W + col_offsets, mask = mask, other = 0).to(tl.float32) + b_row = tl.load(b + col_offsets, mask = mask, other = 0).to(tl.float32) + + inv_var = tl.load(r) .to(tl.float32) + mean = tl.load(mu).to(tl.float32) + normed = (X_row - mean) * inv_var + dY_W = dY_row * W_row + dX_row = dY_W - tl.sum(dY_W, axis = 0) / n_cols - normed * tl.sum(dY_W * normed, axis = 0) / n_cols + dX_row = dX_row * inv_var + tl.store(dY + col_offsets, dX_row, mask = mask) +pass + + +class Fast_Layernorm(torch.autograd.Function): + @staticmethod + def forward(ctx, X, W, b, eps): + shape = X.shape + dim = shape[-1] + X = X.view(-1, dim) + n_rows, n_cols = X.shape + BLOCK_SIZE, num_warps = calculate_settings(n_cols) + + Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0") + r = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") + mu = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") + + layernorm_forward[(n_rows,)]( + Y, Y.stride(0), + X, X.stride(0), + W, + b, + r, + mu, + n_cols, eps, + BLOCK_SIZE = BLOCK_SIZE, + num_warps = num_warps, + ) + ctx.eps = eps + ctx.BLOCK_SIZE = BLOCK_SIZE + ctx.num_warps = num_warps + ctx.save_for_backward(X, W, b, r, mu) + return Y.view(*shape) + pass + + @staticmethod + def backward(ctx, dY): + shape = dY.shape + dim = shape[-1] + dY = dY.view(-1, dim) + X, W, b, r, mu = ctx.saved_tensors + n_rows, n_cols = dY.shape + + layernorm_backward[(n_rows,)]( + dY, dY.stride(0), + X, X .stride(0), + W, + b, + r, + mu, + n_cols, ctx.eps, + BLOCK_SIZE = ctx.BLOCK_SIZE, + num_warps = ctx.num_warps, + ) + dX = dY.view(*shape) + return dX, None, None, None, None + pass +pass + + +def fast_layernorm(layernorm, X): + assert(layernorm.elementwise_affine is True) + W = layernorm.weight + bias = layernorm.bias + eps = layernorm.variance_epsilon if \ + hasattr(layernorm, "variance_epsilon") \ + else layernorm.eps + out = Fast_Layernorm.apply(X, W, bias, eps) + return out +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, + bsz = 21, random_state = 3407, seqlen = 3341, +): + from torch.nn import LayerNorm + layernorm = LayerNorm((dim,), eps = eps, device = "cuda", dtype = dtype) + torch.cuda.manual_seed(random_state) + torch.manual_seed(random_state) + torch.nn.init.uniform_(layernorm.weight) + torch.nn.init.uniform_(layernorm.bias) + X = torch.randn((bsz, seqlen, dim), dtype = dtype, device = "cuda") + XX = X.clone() + X .requires_grad_(True) + XX.requires_grad_(True) + Y = layernorm(X) + YY = torch.randn((bsz, seqlen, dim), dtype = dtype, device = "cuda", requires_grad = True) + Y.backward(YY) + correct_grad = X.grad.clone() + # from unsloth.kernels import fast_layernorm + Y = fast_layernorm(layernorm, XX) + Y.backward(YY) + assert(torch.dist(correct_grad, XX.grad).item() <= 0.1) +pass + + +def testing_suite_layernorm(): + for dim in [512, 1024, 2048]: + for dtype in [torch.float16, torch.bfloat16]: + with torch.autocast(device_type = "cuda", dtype = dtype): + for seqlen in [3341, 2048, 349]: + for random_state in [3407, 42]: + test_layernorm( + dim = dim, + eps = 1e-5, + dtype = dtype, + bsz = 21, + random_state = random_state, + seqlen = seqlen, + ) + pass + pass + pass + pass + pass +pass diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index ac5beb5ab1..13faf08d6a 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -192,3 +192,92 @@ def fast_rms_layernorm(layernorm, X, gemma = False): out = Fast_RMS_Layernorm.apply(X, W, eps, gemma) return out pass + + +from transformers.models.llama.modeling_llama import LlamaRMSNorm +class Unsloth_LlamaRMSNorm(LlamaRMSNorm): + def forward(self, X): + return fast_rms_layernorm(self, X, gemma = False) + pass +pass + +try: + from transformers.models.mllama.modeling_mllama import MllamaTextRMSNorm + class Unsloth_MllamaTextRMSNorm(MllamaTextRMSNorm): + def forward(self, X): + return fast_rms_layernorm(self, X, gemma = False) + pass + pass +except: + pass +pass + +def patch_rms_layernorm(): + import transformers.models.llama.modeling_llama + transformers.models.llama.modeling_llama.LlamaRMSNorm = Unsloth_LlamaRMSNorm + try: + import transformers.models.mllama.modeling_mllama + transformers.models.mllama.modeling_mllama.MllamaTextRMSNorm = Unsloth_MllamaTextRMSNorm + except: + pass + return +pass + + +def unpatch_rms_layernorm(): + import transformers.models.llama.modeling_llama + transformers.models.llama.modeling_llama.LlamaRMSNorm = LlamaRMSNorm + try: + import transformers.models.mllama.modeling_mllama + transformers.models.mllama.modeling_mllama.MllamaTextRMSNorm = MllamaTextRMSNorm + except: + pass + return + return +pass + + +def test_rms_layernorm( + dim = 1024, eps = 1e-5, dtype = torch.float16, + bsz = 21, random_state = 3407, seqlen = 3341, +): + from transformers.models.llama.modeling_llama import LlamaRMSNorm + layernorm = LlamaRMSNorm((dim,), eps = eps).to("cuda") + torch.cuda.manual_seed(random_state) + torch.manual_seed(random_state) + torch.nn.init.uniform_(layernorm.weight) + X = torch.randn((bsz, seqlen, dim), dtype = dtype, device = "cuda") + XX = X.clone() + X .requires_grad_(True) + XX.requires_grad_(True) + Y = layernorm(X) + YY = torch.randn((bsz, seqlen, dim), dtype = dtype, device = "cuda", requires_grad = True) + Y.backward(YY) + correct_grad = X.grad.clone() + # from unsloth.kernels import fast_rms_layernorm + Y = fast_rms_layernorm(layernorm, XX) + Y.backward(YY) + assert(torch.amax(correct_grad - XX.grad).item() <= 0.05) +pass + + +def testing_suite_layernorm(): + for dim in [512, 1024, 2048]: + for dtype in [torch.float16, torch.bfloat16]: + with torch.autocast(device_type = "cuda", dtype = dtype): + for seqlen in [3341, 2048, 349]: + for random_state in [3407, 42]: + test_rms_layernorm( + dim = dim, + eps = 1e-5, + dtype = dtype, + bsz = 21, + random_state = random_state, + seqlen = seqlen, + ) + pass + pass + pass + pass + pass +pass diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index f868c855bc..af7e1eb293 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.9.post2" +__version__ = "2024.9.post3" __all__ = [ "prepare_model_for_kbit_training", @@ -41,6 +41,8 @@ __all__ = [ "torch_amp_custom_bwd", "accelerate_old_send_to_device", "accelerate_new_send_to_device", + "patch_gradient_checkpointing", + "unpatch_gradient_checkpointing", ] import torch @@ -791,7 +793,7 @@ class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function): 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 + hidden_states.requires_grad_(True) with torch.enable_grad(): (output,) = ctx.forward_function(hidden_states, *ctx.args) torch.autograd.backward(output, dY) @@ -806,6 +808,17 @@ def unsloth_offloaded_gradient_checkpoint(function, *args, use_reentrant = None, 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 diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index 50436a7a4f..7f27437904 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -400,6 +400,22 @@ __INT_TO_FLOAT_MAPPER = \ "unsloth/Qwen2.5-Coder-7B-Instruct", "Qwen/Qwen2.5-Coder-7B-Instruct", ), + "unsloth/Llama-3.2-1B-bnb-4bit" : ( + "unsloth/Llama-3.2-1B", + "meta-llama/Llama-3.2-1B", + ), + "unsloth/Llama-3.2-3B-bnb-4bit" : ( + "unsloth/Llama-3.2-3B", + "meta-llama/Llama-3.2-3B", + ), + "unsloth/Llama-3.2-1B-Instruct-bnb-4bit" : ( + "unsloth/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-1B-Instruct", + ), + "unsloth/Llama-3.2-3B-Instruct-bnb-4bit" : ( + "unsloth/Llama-3.2-3B-Instruct", + "meta-llama/Llama-3.2-3B-Instruct", + ), } INT_TO_FLOAT_MAPPER = {} diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py new file mode 100644 index 0000000000..988fda4938 --- /dev/null +++ b/unsloth/models/vision.py @@ -0,0 +1,592 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .llama import * +from ..kernels import patch_layernorm, unpatch_layernorm +from ..kernels import patch_rms_layernorm, unpatch_rms_layernorm +from ..kernels import patch_llama_for_causal_lm, unpatch_llama_for_causal_lm +from ._utils import patch_gradient_checkpointing + +from transformers import AutoProcessor, AutoModelForVision2Seq + + +class FastVisionModel: + + def pre_patch(self): + patch_gradient_checkpointing() + patch_layernorm() + patch_rms_layernorm() + patch_llama_for_causal_lm() + pass + + def post_unpatch(self): + unpatch_layernorm() + unpatch_rms_layernorm() + unpatch_llama_for_causal_lm() + pass + + + @staticmethod + def from_pretrained( + model_name = "llava-hf/llava-1.5-7b-hf", + max_seq_length = None, + dtype = None, + load_in_4bit = True, + token = None, + device_map = "sequential", + rope_scaling = None, + trust_remote_code = False, + **kwargs, + ): + if trust_remote_code: + print( + "Unsloth: WARNING `trust_remote_code` is True.\n"\ + "Are you certain you want to do remote code execution?" + ) + pass + if token is None: token = get_token() + if model_patcher is None: model_patcher = FastLlamaModel + SUPPORTS_BFLOAT16 = is_bfloat16_supported() + gpu_stats = torch.cuda.get_device_properties(0) + max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) + + statistics = \ + f"==((====))== Unsloth {__version__}: Fast {model_patcher.__name__[4:-5]} patching. Transformers = {transformers_version}.\n"\ + f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB. Platform = {platform_system}.\n"\ + f"O^O/ \_/ \\ Pytorch: {torch.__version__}. CUDA = {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit = {torch.version.cuda}.\n"\ + f"\ / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\n"\ + f' "-____-" Free Apache license: http://github.com/unslothai/unsloth' + print(statistics) + + # Warn about fast transfers + old_hf_transfer = os.environ.get("HF_HUB_ENABLE_HF_TRANSFER", "0") + if os.environ.get("HF_HUB_ENABLE_HF_TRANSFER", "0") == "1": + print("Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!") + pass + # Return old flag + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer + + get_statistics() # For debugging - we use a download counter to see if environments are not breaking + + if dtype is None: + dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 + elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16: + logger.warning_once("Device does not support bfloat16. Will change to float16.") + dtype = torch.float16 + + assert(dtype == torch.float16 or dtype == torch.bfloat16 or dtype == torch.float32) + + # RoPE Scaling + model_config = AutoConfig.from_pretrained(model_name, token = token) + model_max_seq_length = model_config.max_position_embeddings + + # Check if RoPE Scaling is even allowed + model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__] + has_rope_scaling = False + try: + with open(inspect.getfile(model_function), "r") as file: + has_rope_scaling = "self.config.rope_scaling" in file.read() + except: pass + has_rope_scaling = True + + # If max_seq_length is not specified, use maximum fron config + if max_seq_length is None: + max_seq_length = model_max_seq_length + pass + + if (rope_scaling is None) and (max_seq_length > model_max_seq_length): + + rope_scaling = max_seq_length / model_max_seq_length + + logger.warning_once( + f"Unsloth: {model_name} can only handle sequence lengths of at most "\ + f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "\ + f"{round(rope_scaling, 3)}, it can be magically be extended to "\ + f"{max_seq_length}!" + ) + + # Warn RoPE scaling isn't allowed + if not has_rope_scaling: + raise RuntimeError( + "However, {model_name} doesn't support RoPE Scaling!\n"\ + "Please file a feature request at https://github.com/unslothai/unsloth." + ) + pass + + rope_scaling = {"type": "linear", "factor": rope_scaling,} + + # Add to kwargs + kwargs["rope_scaling"] = rope_scaling + pass + # We currently only support NVIDIA GPUs - AMD / Intel is a work in progress! + pre_check = check_nvidia() + + bnb_config = None + if load_in_4bit: + bnb_config = BitsAndBytesConfig( + load_in_4bit = True, + bnb_4bit_use_double_quant = True, + bnb_4bit_quant_type = "nf4", + bnb_4bit_compute_dtype = dtype, + ) + pass + + # https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/discussions/12 + # RoPE Scaling's max_position_embeddings must be updated + max_position_embeddings = max(max_seq_length, model_max_seq_length) + kwargs.pop("attn_implementation", None); # No need since we auto call it + + # Cannot be None, since HF now checks for the config + if load_in_4bit: kwargs["quantization_config"] = bnb_config + + self.pre_patch() + model = AutoModelForVision2Seq.from_pretrained( + model_name, + device_map = device_map, + torch_dtype = dtype, + # quantization_config = bnb_config, + token = token, + max_position_embeddings = max_position_embeddings, + trust_remote_code = trust_remote_code, + attn_implementation = "eager", + **kwargs, + ) + self.post_unpatch() + + # Return old flag + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer + # We currently only support NVIDIA GPUs - AMD / Intel is a work in progress! + post_check = check_nvidia() + + # Counteract saved tokenizers + tokenizer = AutoProcessor.from_pretrained( + model_name, + ) + model = FastVisionModel.post_patch(model) + + # Patch Trainer + from transformers.trainer import Trainer + try: + if Trainer._inner_training_loop.__name__ != "_fast_inner_training_loop": + inner_training_loop = inspect.getsource(Trainer._inner_training_loop) + Trainer._original_training_loop = inner_training_loop + else: + inner_training_loop = Trainer._original_training_loop + except: + raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!') + pass + + if ((post_check - pre_check) >= 1).sum() > 1: + raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!') + + import transformers.trainer + items_in_trainer = dir(transformers.trainer) + good_items = [] + for item in items_in_trainer: + # TODO: Support Deepspeed + if item.startswith(("deepspeed", "xm", "met", "smp")): continue + if item in inner_training_loop: good_items.append(item) + pass + exec("from transformers.trainer import (" + ", ".join(x for x in good_items) + ")", globals()) + + start = re.search('logger\.info\([\"\'].+?Running training', inner_training_loop).span(0)[0] + end = inner_training_loop.find("\n\n", start) + original_debug = inner_training_loop[start:end] + spaces = re.search('\n([\s\t]{1,})', original_debug).group(0)[1:] + front_spaces = re.match('([\s\t]{1,})', inner_training_loop).group(0) + + debug_info = """debug_info = \\ + f"==((====))== Unsloth - 2x faster free finetuning | Num GPUs = {args.world_size}\\n"\\ + f" \\\\\\ /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,}\\n"\\ + f"O^O/ \\_/ \\ Batch size per device = {self._train_batch_size:,} | Gradient Accumulation steps = {args.gradient_accumulation_steps}\\n"\\ + f"\\ / Total batch size = {total_train_batch_size:,} | Total steps = {max_steps:,}\\n"\\ + f' "-____-" Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}' + logger.warning(debug_info) + import subprocess, re, gc, numpy as np + a = np.array([0,]) + try: + a = subprocess.check_output('nvidia-smi --query-gpu=memory.used --format=csv', shell = True) + a = re.findall(rb'([\\d]{1,})[\\s]{1,}M', a) + a = np.array([int(x.decode('utf-8'))/1024 for x in a]) + except: + if not torch.cuda.is_available(): + raise RuntimeError('Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!') + if ((a - PRE_CHECK) >= 1).sum() > 1: + raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!') + for _ in range(3): + gc.collect() + torch.cuda.empty_cache()""" + + debug_info = debug_info.split('\n') + debug_info = "\n".join([debug_info[0]] + [spaces + x[8:] for x in debug_info[1:]]) + inner_training_loop = inner_training_loop.replace(original_debug, debug_info) + + debug_info = """n_total_devices = total_train_batch_size // \\ + args.gradient_accumulation_steps // self._train_batch_size + if n_total_devices > 1: + logger.warning_once('Unsloth currently does not support multi GPU setups - but we are working on it!') + debug_info =""" + debug_info = debug_info.split('\n') + debug_info = "\n".join([debug_info[0]] + [spaces + x[8:] for x in debug_info[1:]]) + inner_training_loop = inner_training_loop.replace("debug_info =", debug_info, 1) + + front_spaces = re.match(r"[\t\s]{1,}", inner_training_loop).group(0) + inner_training_loop = re.sub(r"^" + front_spaces, "", inner_training_loop, flags = re.MULTILINE) + inner_training_loop = inner_training_loop.replace( + "train_dataloader = tpu_spmd_dataloader(train_dataloader)", + "raise RuntimeError('Unsloth: TPUs are not yet supported!')" + ) + inner_training_loop = inner_training_loop.replace( + "self.accelerator.free_memory()", + "self.accelerator.free_memory()\n" + \ + front_spaces + "if self.is_deepspeed_enabled:"\ + "raise RuntimeError('Unsloth: Deepspeed is not yet supported!')\n", 1, + ) + + check_batches = """train_dataloader = self.get_train_dataloader() + ga = args.gradient_accumulation_steps + bsz = self._train_batch_size + total_batches = bsz * ga * args.world_size + n_total_devices = total_batches // ga // bsz + if n_total_devices > 1: + logger.warning_once('Unsloth currently does not support multi GPU setups - but we are working on it!') + divisor = n_total_devices / 1 + bsz = self._train_batch_size = max(int(bsz / divisor), 1) + if total_batches // ga // bsz > 1: + divisor = n_total_devices / 1 + ga = args.gradient_accumulation_steps = max(int(ga / divisor), 1)""" + check_batches = check_batches.split('\n') + check_batches = "\n".join([check_batches[0]] + [front_spaces + x[8:] for x in check_batches[1:]]) + inner_training_loop = inner_training_loop.replace( + "train_dataloader = self.get_train_dataloader()", + check_batches, 1, + ) + inner_training_loop = inner_training_loop.replace( + "_inner_training_loop", + "_fast_inner_training_loop", 1, + ) + exec(inner_training_loop, globals()) + + Trainer._inner_training_loop = _fast_inner_training_loop + inner_training_loop = inner_training_loop.replace( + "is_torch_tpu_available()", + "False", + ) + if "n_total_devices >" not in inner_training_loop: + raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!') + pass + inner_training_loop = inner_training_loop.replace( + "is_sagemaker_mp_enabled()", + "False", + ) + exec(inner_training_loop, globals()) + Trainer._inner_training_loop = _fast_inner_training_loop + + # Save max_seq_length + model.max_seq_length = max_position_embeddings + internal_model = model + while hasattr(internal_model, "model"): + internal_model.max_seq_length = max_position_embeddings + internal_model = internal_model.model + pass + internal_model.max_seq_length = max_position_embeddings + + # Fix up config for transformers uploading PEFT + # Not necessary anymore since we require transformers>=4.37! + if False: + name = model.config._name_or_path + if name.startswith("unsloth/") and name.endswith("-bnb-4bit"): + name = name[:len(name) - len("-bnb-4bit")] + model.config.update({"_name_or_path" : name}) + pass + pass + + # Log Unsloth version for future fastpaths for inference + model.config.update({"unsloth_version" : __version__}) + + # Add save modules + patch_saving_functions(model) + Trainer._inner_training_loop = _fast_inner_training_loop + + # 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 + lm_head = model.get_output_embeddings().weight + + # 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 + pass + + # Clear deleted GPU items + for _ in range(3): + gc.collect() + torch.cuda.empty_cache() + return model + pass + + + @staticmethod + def get_peft_model( + model, + r = 16, + target_modules = "all-linear", + lora_alpha = 16, + lora_dropout = 0, + bias = "none", + layers_to_transform = None, + layers_pattern = None, + use_gradient_checkpointing = True, + random_state = 3407, + max_seq_length = 2048, # not used anymore + use_rslora = False, + modules_to_save = None, + init_lora_weights = True, + loftq_config = {}, + temporary_location = "_unsloth_temporary_saved_buffers", + **kwargs, + ): + transformers_set_seed(random_state) + + # Get LoRA + arguments = dict( + r = r, + lora_alpha = lora_alpha, + target_modules = target_modules, + lora_dropout = lora_dropout, + bias = bias, + layers_to_transform = layers_to_transform, + init_lora_weights = init_lora_weights, + # loftq_config = loftq_config, + # use_rslora = use_rslora, + modules_to_save = modules_to_save, + **kwargs, + ) + + lora_config = LoraConfig(**arguments) + + model = _get_peft_model(model, lora_config) + + model = FastVisionModel.patch_peft_model(model, use_gradient_checkpointing) + + # Clear deleted GPU items + for _ in range(3): + gc.collect() + torch.cuda.empty_cache() + pass + + return model + pass + + + @staticmethod + def patch_peft_model( + model, + use_gradient_checkpointing = True, + ): + + model = prepare_model_for_kbit_training( + model, + use_gradient_checkpointing = use_gradient_checkpointing, + use_reentrant = True, + ) + + # Fix up config for transformers uploading PEFT + for active_adapter in model.peft_config.keys(): + # Not necessary since we requires transformers >= 4.37 + if False: + name = model.peft_config[active_adapter].base_model_name_or_path + if name.startswith("unsloth/") and name.endswith("-bnb-4bit"): + name = name[:len(name) - len("-bnb-4bit")] + model.peft_config[active_adapter].base_model_name_or_path = name + pass + # Add revision to enable future fast inference paths + # [TODO] Bugs out!see https://github.com/unslothai/unsloth/issues/492 + # model.peft_config[active_adapter].revision = f"unsloth" + pass + + from transformers.trainer import Trainer + if Trainer._inner_training_loop.__name__ != "_fast_inner_training_loop": + raise RuntimeError( + '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!\n'\ + 'We do have a separate beta version, which you can contact us about!\n'\ + 'Thank you for your understanding and we appreciate it immensely!' + ) + pass + + logger.warning_once( + f"Unsloth {__version__} patched {len(model.model.model.layers)} layers with "\ + f"{n_qkv} QKV layers, {n_o} O layers and {n_mlp} MLP layers.", + ) + patch_saving_functions(model) + + # Patch cross entropy loss labels + # Fixes https://github.com/unslothai/unsloth/issues/10 + max_seq_length = model.max_seq_length + extra_ignored_labels = torch.full((max_seq_length, 1), -100, device = "cuda:0") + model.model.extra_ignored_labels = extra_ignored_labels + internal_model = model + while hasattr(internal_model, "model"): + internal_model.max_seq_length = max_seq_length + internal_model = internal_model.model + pass + internal_model.max_seq_length = max_seq_length + + # Patch tokenizer to pad to the right + internal_model = model + while hasattr(internal_model, "model"): + if hasattr(internal_model, "_saved_temp_tokenizer"): + internal_model._saved_temp_tokenizer.padding_side = "right" + pass + internal_model = internal_model.model + pass + if hasattr(internal_model, "_saved_temp_tokenizer"): + internal_model._saved_temp_tokenizer.padding_side = "right" + pass + + # Clear deleted GPU items + for _ in range(3): + gc.collect() + torch.cuda.empty_cache() + pass + return model + pass + + + @staticmethod + def for_inference(model): + # if model.config.model_type == "qwen2": + # FastLlamaModel.for_training(model) + # return + # pass + + internal_model = model + internal_model.gradient_checkpointing = False + internal_model.training = False + + while hasattr(internal_model, "model"): + internal_model = internal_model.model + internal_model.gradient_checkpointing = False + internal_model.training = False + pass + if hasattr(internal_model, "training"): + internal_model.training = False + pass + + # Also check if lm_head / embeddings are trained + internal_model = model + while not hasattr(internal_model, "lm_head"): + internal_model = internal_model.model + pass + lm_head = internal_model.lm_head.weight + device_type = lm_head.device.type + dtype = model.config.torch_dtype + + if type(dtype) is str: + if dtype == "float16": dtype = torch.float16 + elif dtype == "bfloat16": dtype = torch.bfloat16 + pass + + # Also disable training for embeddings for NEFTune + if hasattr(model, "get_input_embeddings"): + embeddings = model.get_input_embeddings() + if hasattr(embeddings, "training"): embeddings.training = False + pass + if hasattr(model, "get_output_embeddings"): + embeddings = model.get_output_embeddings() + if hasattr(embeddings, "training"): embeddings.training = False + pass + + return model + pass + + + @staticmethod + def for_training(model, use_gradient_checkpointing = True): + internal_model = model + internal_model.gradient_checkpointing = use_gradient_checkpointing + internal_model.training = True + + # Delete all fast inference loras + for param in model.parameters(): + if hasattr(param, "_fast_lora"): + del param._fast_lora + pass + + while hasattr(internal_model, "model"): + internal_model = internal_model.model + internal_model.gradient_checkpointing = use_gradient_checkpointing + internal_model.training = True + pass + if hasattr(internal_model, "training"): + internal_model.training = True + pass + + # Also re-enable training for embeddings for NEFTune + if hasattr(model, "get_input_embeddings"): + embeddings = model.get_input_embeddings() + if hasattr(embeddings, "training"): embeddings.training = True + pass + if hasattr(model, "get_output_embeddings"): + embeddings = model.get_output_embeddings() + if hasattr(embeddings, "training"): embeddings.training = True + pass + + return model + pass +pass From 063a2b42b02f9991f723b2c3f8637378cd3bc28d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Sep 2024 12:18:43 -0700 Subject: [PATCH 2/4] Llama 3.2 --- unsloth/models/llama.py | 10 ++++++++-- unsloth/models/loader.py | 22 ++++++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index a245330108..aedc935d75 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -62,7 +62,10 @@ from peft.tuners.lora import Linear4bit as Peft_Linear4bit from ..save import patch_saving_functions import re, os, inspect, math, sys from huggingface_hub.utils._token import get_token - +from transformers import __version__ as transformers_version +from packaging.version import Version +transformers_version = Version(transformers_version) +SUPPORTS_LLAMA32 = transformers_version >= Version("4.46") def original_apply_qkv(self, X): Q = self.q_proj(X) @@ -1383,7 +1386,10 @@ def _wrap_fast_inference(generate, device_type, dtype, model): pass # For newer HF - kwargs["cache_implementation"] = "dynamic" + if SUPPORTS_LLAMA32: + kwargs["cache_implementation"] = "hybrid" + else: + kwargs["cache_implementation"] = "dynamic" # For num_logits_to_keep kwargs["num_logits_to_keep"] = 1 diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 13710eeda1..8afb3a3a82 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -31,6 +31,7 @@ SUPPORTS_FOURBIT = transformers_version >= Version("4.37") SUPPORTS_GEMMA = transformers_version >= Version("4.38") SUPPORTS_GEMMA2 = transformers_version >= Version("4.42") SUPPORTS_LLAMA31 = transformers_version >= Version("4.43.2") +SUPPORTS_LLAMA32 = transformers_version >= Version("4.46") if SUPPORTS_GEMMA: from .gemma import FastGemmaModel if SUPPORTS_GEMMA2: @@ -137,6 +138,17 @@ def get_model_name(model_name, load_in_4bit = True): pass +LLAMA32_MODEL_NAMES = set(( + "unsloth/Llama-3.2-1B-bnb-4bit", + "unsloth/Llama-3.2-3B-bnb-4bit", + "unsloth/Llama-3.2-1B", + "unsloth/Llama-3.2-3B", + "unsloth/Llama-3.2-1B-Instruct-bnb-4bit", + "unsloth/Llama-3.2-3B-Instruct-bnb-4bit", + "unsloth/Llama-3.2-1B-Instruct", + "unsloth/Llama-3.2-3B-Instruct", +)) + class FastLanguageModel(FastLlamaModel): @staticmethod def from_pretrained( @@ -242,7 +254,17 @@ class FastLanguageModel(FastLlamaModel): f'Try `pip install --upgrade "transformers>=4.43.2"`\n'\ f"to obtain the latest transformers build, then restart this session."\ ) + + elif model_name.lower() in LLAMA32_MODEL_NAMES and not SUPPORTS_LLAMA32: + raise ImportError( + f"Unsloth: Your transformers version of {transformers_version} does not support Llama 3.2.\n"\ + f"The minimum required version is 4.46\n"\ + f'Try `pip install --upgrade "transformers>=4.46"`\n'\ + f"to obtain the latest transformers build, then restart this session."\ + ) + dispatch_model = FastLlamaModel + elif model_type == "mistral": dispatch_model = FastMistralModel elif model_type == "gemma": if not SUPPORTS_GEMMA: From 96e92c9282efccfe64c2ac61f5371abed677753b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Sep 2024 12:35:38 -0700 Subject: [PATCH 3/4] Fix version --- unsloth/models/llama.py | 2 +- unsloth/models/loader.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index aedc935d75..f5dc02704e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -65,7 +65,7 @@ from huggingface_hub.utils._token import get_token from transformers import __version__ as transformers_version from packaging.version import Version transformers_version = Version(transformers_version) -SUPPORTS_LLAMA32 = transformers_version >= Version("4.46") +SUPPORTS_LLAMA32 = transformers_version > Version("4.45.0") def original_apply_qkv(self, X): Q = self.q_proj(X) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 8afb3a3a82..ee1f680ace 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -31,7 +31,7 @@ SUPPORTS_FOURBIT = transformers_version >= Version("4.37") SUPPORTS_GEMMA = transformers_version >= Version("4.38") SUPPORTS_GEMMA2 = transformers_version >= Version("4.42") SUPPORTS_LLAMA31 = transformers_version >= Version("4.43.2") -SUPPORTS_LLAMA32 = transformers_version >= Version("4.46") +SUPPORTS_LLAMA32 = transformers_version > Version("4.45.0") if SUPPORTS_GEMMA: from .gemma import FastGemmaModel if SUPPORTS_GEMMA2: @@ -264,7 +264,7 @@ class FastLanguageModel(FastLlamaModel): ) dispatch_model = FastLlamaModel - + elif model_type == "mistral": dispatch_model = FastMistralModel elif model_type == "gemma": if not SUPPORTS_GEMMA: From 5c7cb059237d5245b290f066de1dfdd8ed9d918e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Sep 2024 12:42:24 -0700 Subject: [PATCH 4/4] Update llama.py --- unsloth/models/llama.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index f5dc02704e..bae6d5b80d 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1387,7 +1387,8 @@ def _wrap_fast_inference(generate, device_type, dtype, model): # For newer HF if SUPPORTS_LLAMA32: - kwargs["cache_implementation"] = "hybrid" + # kwargs["cache_implementation"] = "hybrid" + pass else: kwargs["cache_implementation"] = "dynamic" # For num_logits_to_keep