1 line
146 KiB
JSON
1 line
146 KiB
JSON
{"repo": "biodatlab/medexact-lamar", "files": {"unsloth/models/llama.py": {"9692400ad88999b8065609115b6eb554cfa0a114": {"t": 1785674498.5904007, "n": 144115, "h": "57b3b631a4dfebb6", "c": "# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nimport gc\nimport math\nimport functools\nfrom typing import Optional, Tuple, List, Union\n\nfrom ._utils import *\nfrom ._utils import apply_unsloth_gradient_checkpointing\nfrom ._utils import __version__, importlib_version\nfrom ._utils import move_to_device\nfrom ._utils import (\n _get_inference_mode_context_manager,\n _prepare_model_for_qat,\n is_bfloat16_supported,\n get_quant_type,\n)\nfrom .loader_utils import _get_fp8_mode_and_check_settings\nfrom ..utils.packing import (\n get_packed_info_from_kwargs,\n mask_packed_sequence_boundaries,\n)\nfrom ..utils.attention_dispatch import (\n AttentionConfig,\n AttentionContext,\n run_attention,\n SDPA,\n select_attention_backend,\n)\nfrom torch.nn.functional import scaled_dot_product_attention\nfrom transformers import __version__ as transformers_version\nfrom unsloth_zoo.utils import Version, _get_dtype\nfrom unsloth_zoo.hf_utils import (\n dtype_from_config,\n add_dtype_kwargs,\n fix_lora_auto_mapping,\n)\nfrom unsloth_zoo.peft_utils import SKIP_QUANTIZATION_MODULES\nfrom ..device_type import (\n is_hip,\n get_device_type,\n DEVICE_TYPE,\n DEVICE_TYPE_TORCH,\n DEVICE_COUNT,\n ALLOW_PREQUANTIZED_MODELS,\n)\n\ntransformers_version = Version(transformers_version)\n# Transformers moved rotary embeddings out of all attention layers\nIS_ATTENTION_REFACTOR = transformers_version > Version(\"4.47.1\")\ntry:\n from transformers.modeling_layers import GradientCheckpointingLayer\nexcept:\n GradientCheckpointingLayer = type(None)\n\nfrom transformers.models.llama.modeling_llama import (\n logger,\n BaseModelOutputWithPast,\n CausalLMOutputWithPast,\n)\nfrom transformers.modeling_attn_mask_utils import (\n _prepare_4d_causal_attention_mask_for_sdpa,\n)\nfrom ..kernels import *\nfrom ..tokenizer_utils import *\nfrom .vision import FastBaseModel\n\n# Final patching code\nfrom transformers.models.llama.modeling_llama import (\n LlamaAttention,\n LlamaDecoderLayer,\n LlamaModel,\n LlamaForCausalLM,\n)\n\n# For Pytorch 2.1.1\ntry:\n from transformers.models.llama.modeling_llama import (\n LlamaSdpaAttention,\n LlamaFlashAttention2,\n )\nexcept:\n LlamaSdpaAttention = LlamaAttention\n LlamaFlashAttention2 = LlamaAttention\n\nfrom transformers import (\n AutoTokenizer,\n AutoModelForCausalLM,\n AutoModelForSequenceClassification,\n BitsAndBytesConfig,\n AutoConfig,\n)\nfrom transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING\nfrom transformers import set_seed as transformers_set_seed\nfrom peft import LoraConfig, TaskType, get_peft_model as _get_peft_model\nfrom peft import PeftModelForCausalLM, PeftModelForSequenceClassification\nfrom ..save import patch_saving_functions\nimport re, os, inspect, math, sys\nimport types\n\ntry:\n from huggingface_hub.utils import get_token\nexcept:\n # Old HF Hub versions <= 0.0.25\n from huggingface_hub.utils._token import get_token\nfrom triton import __version__ as triton_version\n\nHAS_XFORMERS = xformers is not None\nBlockDiagonalCausalMask = (\n xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None\n)\n\nif DEVICE_TYPE == \"xpu\":\n clean_gpu_cache = torch.xpu.empty_cache\n get_current_device = torch.xpu.current_device\nelse:\n clean_gpu_cache = torch.cuda.empty_cache\n get_current_device = torch.cuda.current_device\n\n\ndef original_apply_qkv(self, X):\n Q = self.q_proj(X)\n K = self.k_proj(X)\n V = self.v_proj(X)\n return Q, K, V\n\n\ndef original_apply_o(self, X):\n O = self.o_proj(X)\n return O\n\n\nfrom math import sqrt as math_sqrt\n\nKV_CACHE_INCREMENT = 512 # KV Cache update size\ntorch_nn_functional_softmax = torch.nn.functional.softmax\n# SDPA has GQA internally\nSDPA_HAS_GQA = \"enable_gqa\" in scaled_dot_product_attention.__doc__\n\nfrom peft.utils.other import ModulesToSaveWrapper\n\n\ndef _offload_frozen_module_for_training(\n module: ModulesToSaveWrapper,\n device_type: str,\n offload_device: Optional[str] = \"cpu\",\n) -> None:\n \"\"\"\n Offload frozen module to CPU and configure trainable copy for mixed precision training.\n\n This function optimizes memory usage by:\n 1. Moving the trainable copy to the target device with appropriate precision\n 2. Optionally offloading the original frozen module to CPU/disk to free VRAM\n 3. Converting float16 to float32 for compatibility with certain GPUs (e.g., Tesla T4)\n\n Args:\n module: The module to configure. Must be a ModulesToSaveWrapper with a\n `modules_to_save` attribute containing trainable and original modules.\n device_type: Target device string for training (e.g., \"cuda:0\", \"xpu:0\")\n offload_device: Device to offload frozen parameters (default: \"cpu\").\n If None, the original frozen module remains on its current device.\n Note: Currently only \"cpu\" is supported; disk offloading is planned.\n\n Returns:\n None (modifies module in-place)\n\n Note:\n - Float16 weights are automatically promoted to float32 for GPU compatibility\n - When offload_device is specified, frozen parameters are moved to free VRAM\n - Future versions will support disk-based offloading for even larger models\n\n See Also:\n - https://github.com/unslothai/unsloth/pull/1200 (Tesla T4 float32 requirement)\n \"\"\"\n # Early return with explicit None if module doesn't support mixed precision training\n if not hasattr(module, \"modules_to_save\"):\n return None\n\n new_dtype = module.modules_to_save.default.weight.dtype\n if new_dtype == torch.float16:\n # See https://github.com/unslothai/unsloth/pull/1200\n # Tesla T4 must use float32 and not float16\n new_dtype = torch.float32\n\n module.modules_to_save.default.to(\n device = device_type, dtype = new_dtype, non_blocking = True\n )\n module.modules_to_save.default.requires_grad_(True)\n\n # [TODO] Move old module to CPU - should be disk!\n if offload_device is not None:\n module.original_module.to(device = offload_device, non_blocking = True)\n module.original_module.requires_grad_(False)\n\n\n# Fix new HF's inference code\ndef _fast_prepare_inputs_for_generation(\n self,\n input_ids,\n attention_mask = None,\n inputs_embeds = None,\n **kwargs,\n):\n past_key_values = kwargs.get(\"past_key_values\", None)\n original_attention_mask = attention_mask\n\n # Handle inputs_embeds - only use on FIRST generation step (no cache)\n # This fixes GitHub issue #3798: inputs_embeds was ignored\n use_inputs_embeds = inputs_embeds is not None and past_key_values is None\n\n if input_ids is not None and input_ids.numel() > 0:\n bs, seq_length = input_ids.shape\n device = input_ids.device\n elif inputs_embeds is not None:\n bs, seq_length, _ = inputs_embeds.shape\n device = inputs_embeds.device\n else:\n bs, seq_length = 1, 0\n device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n if past_key_values is not None:\n # Check for uninitialized DynamicCache\n if len(past_key_values) == 0:\n past_key_values = None\n kwargs[\"past_key_values\"] = None\n use_inputs_embeds = inputs_embeds is not None\n # New since 4.56\n elif (\n hasattr(past_key_values, \"get_seq_length\")\n and past_key_values.get_seq_length() == 0\n ):\n past_key_values = None\n kwargs[\"past_key_values\"] = None\n use_inputs_embeds = inputs_embeds is not None\n else:\n if input_ids is not None and input_ids.numel() > 0:\n bs = input_ids.shape[0]\n input_ids = input_ids[:, [-1]]\n device = input_ids.device\n seq_length = 1\n elif inputs_embeds is not None:\n bs, seq_length, _ = inputs_embeds.shape\n device = inputs_embeds.device\n else:\n bs, seq_length = 1, 0\n device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n if hasattr(past_key_values, \"get_seq_length\"):\n past_len = int(past_key_values.get_seq_length())\n else:\n # legacy tuple cache: (layer, (K,V))\n past_len = int(past_key_values[0][0].shape[-2])\n\n max_cache_len = None\n if hasattr(past_key_values, \"get_max_cache_shape\"):\n m = past_key_values.get_max_cache_shape()\n max_cache_len = int(m) if m is not None and m > 0 else None\n elif hasattr(past_key_values, \"get_max_length\"):\n m = past_key_values.get_max_length()\n max_cache_len = int(m) if m is not None else None\n\n # ensure cache_position\n cache_position = kwargs.get(\"cache_position\", None)\n if cache_position is None:\n kwargs[\"cache_position\"] = torch.arange(\n past_len,\n past_len + seq_length,\n device = device,\n dtype = torch.long,\n )\n else:\n if (\n hasattr(cache_position, \"device\")\n and cache_position.device != device\n ):\n kwargs[\"cache_position\"] = cache_position.to(device)\n\n # Get to the base model\n base_model = self\n if hasattr(base_model, \"base_model_prefix\"):\n base_model = getattr(base_model, base_model.base_model_prefix)\n\n if hasattr(\n base_model, \"_prepare_4d_causal_attention_mask_with_cache_position\"\n ):\n if not hasattr(base_model, \"_unsloth_mask_needs_device\"):\n\n def _check_needs_device(fn) -> bool:\n try:\n sig = inspect.signature(inspect.unwrap(fn))\n return \"device\" in sig.parameters\n except:\n # transformers <= 4.51.3 includes device arg but > 4.51.3 does not\n return transformers_version < Version(\"4.52.0\")\n\n base_model._unsloth_mask_needs_device = _check_needs_device(\n base_model._prepare_4d_causal_attention_mask_with_cache_position\n )\n\n if max_cache_len is not None:\n target_length = max_cache_len\n elif (\n original_attention_mask is not None\n and original_attention_mask.dim() == 2\n ):\n target_length = original_attention_mask.shape[-1]\n else:\n target_length = past_len + seq_length\n\n mask_kwargs = {\n \"sequence_length\": seq_length,\n \"target_length\": target_length,\n \"dtype\": self.dtype,\n \"cache_position\": kwargs[\"cache_position\"],\n \"batch_size\": bs,\n \"config\": self.config,\n \"past_key_values\": past_key_values,\n }\n if base_model._unsloth_mask_needs_device:\n mask_kwargs[\"device\"] = device\n\n attention_mask = (\n base_model._prepare_4d_causal_attention_mask_with_cache_position(\n attention_mask,\n **mask_kwargs,\n )\n )\n else:\n if transformers_version <= Version(\"4.52.4\"):\n logger.warning_once(\n f\"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method \"\n \"defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're \"\n \"writing code, see Llama for an example implementation. If you're a user, please report this \"\n \"issue on GitHub.\"\n )\n\n if kwargs.get(\"position_ids\", None) is None:\n if original_attention_mask is not None and original_attention_mask.dim() == 2:\n position_ids = original_attention_mask.long().cumsum(-1) - 1\n position_ids.masked_fill_(original_attention_mask == 0, 1)\n position_ids = position_ids[:, -seq_length:]\n kwargs[\"position_ids\"] = position_ids\n elif kwargs.get(\"cache_position\", None) is not None:\n cp = kwargs[\"cache_position\"]\n if cp.dim() == 1:\n cp = cp.unsqueeze(0).expand(bs, -1)\n kwargs[\"position_ids\"] = cp\n\n result = {\n \"attention_mask\": attention_mask,\n **kwargs,\n }\n if use_inputs_embeds:\n result[\"inputs_embeds\"] = inputs_embeds\n result[\"input_ids\"] = None\n else:\n result[\"input_ids\"] = input_ids\n return result\n\n\ndef fix_prepare_inputs_for_generation(module):\n # Fix prepare_inputs_for_generation\n if hasattr(module, \"prepare_inputs_for_generation\"):\n module.prepare_inputs_for_generation = _fast_prepare_inputs_for_generation\n\n\ntorch_matmul = torch.matmul\n\n\ndef LlamaAttention_fast_forward_inference(\n self,\n hidden_states: torch.Tensor,\n past_key_value: Optional[Tuple[torch.Tensor]],\n position_ids,\n do_prefill = False,\n attention_mask = None,\n rotary_seq_len = None,\n):\n \"\"\"\n https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406\n Fast inference using KV cache.\n QK^T can be computed in 4 chunks\n\n [Q, q] @ [K, k].T where q, k are the new tokens.\n [QK^T, Qk^T]\n [qK^T, qk^T]\n\n Since the attention mask wipes Qk^T, we just get\n [QK^T, 0]\n [qK^T, qk^T]\n\n Since softmax is row-wise, we get\n softmax([QK^T, 0])\n softmax([qK^T, qk^T])\n\n We then multiply by [V]\n [v]\n softmax([QK^T, 0]) [softmax(QK^T)V] *\n softmax([qK^T, qk^T]) [softmax([qK^T, qk^T]) @ [V, v]]\n\n But notice * [softmax(QK^T)V] is just the last attention.\n We just need to compute the last final row.\n\n This means we can pass in a row of Q, but we need to\n remember K and V, which are called the KV cache.\n \"\"\"\n Xn = hidden_states\n bsz, _, hd = hidden_states.size()\n K1, V1 = past_key_value\n dtype = Xn.dtype\n\n n_heads = self.config.num_attention_heads\n n_groups = self.num_key_value_groups\n n_kv_heads = self.config.num_key_value_heads\n head_dim = self.head_dim\n # assert(n_kv_heads * n_groups == n_heads)\n\n hidden_size = self.config.hidden_size\n attention_size = n_heads * head_dim\n seq_len = K1.shape[-2]\n kv_seq_len = seq_len + 1\n\n # Prefill phase\n # if not hasattr(self, \"paged_attention\"):\n device = hidden_states.device\n if do_prefill:\n self.paged_attention = torch.empty(\n (KV_CACHE_INCREMENT + seq_len + 1, 2, bsz, n_kv_heads, head_dim),\n dtype = dtype,\n device = device,\n )\n self.paged_attention_K = self.paged_attention[:, 0]\n self.paged_attention_V = self.paged_attention[:, 1]\n self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3)\n self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3)\n self.temp_QA = torch.empty(\n (2, bsz, 1, attention_size), dtype = dtype, device = device\n )\n self.temp_KV = torch.empty(\n (2, bsz, 1, n_kv_heads * head_dim), dtype = dtype, device = device\n )\n self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)\n\n # Mistral Nemo 12b has weird dimensions\n if attention_size != hidden_size:\n self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)\n else:\n self.temp_O = self.temp_QA[1][:, :, :hidden_size]\n\n self.attention = torch.empty(\n (bsz, n_heads, 1, KV_CACHE_INCREMENT + seq_len), dtype = dtype, device = device\n )\n self.scalar = 1.0 / math_sqrt(self.head_dim)\n self.half_head_dim = head_dim // 2\n elif kv_seq_len >= self.paged_attention.shape[0]:\n self.paged_attention.resize_(\n (\n self.paged_attention.shape[0] + KV_CACHE_INCREMENT,\n 2,\n bsz,\n n_kv_heads,\n head_dim,\n )\n )\n self.paged_attention_K = self.paged_attention[:, 0]\n self.paged_attention_V = self.paged_attention[:, 1]\n self.attention.resize_(\n (bsz, n_heads, 1, self.attention.shape[-1] + KV_CACHE_INCREMENT)\n )\n\n Qn = fast_linear_forward(self.q_proj, Xn, out = self.temp_QA[0])\n Kn = fast_linear_forward(self.k_proj, Xn, out = self.temp_KV[0])\n Vn = fast_linear_forward(self.v_proj, Xn, out = self.temp_KV[1])\n Qn = Qn.view(bsz, 1, n_heads, head_dim).transpose(1, 2)\n Kn = Kn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2)\n Vn = Vn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2)\n\n # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)\n # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)\n\n # Need to do it prior 2 steps before hitting full on short KV cache\n # or else error\n # ensure correct shape\n if position_ids.dim() == 1:\n position_ids = position_ids[:, None]\n # Transformers 5.x generate() accumulates position_ids as [batch, full_seq_len]\n # across decode steps. In single-token inference we only need the last position.\n if position_ids.shape[-1] > 1:\n position_ids = position_ids[:, -1:]\n position_ids = position_ids.to(Qn.device)\n\n if rotary_seq_len is None:\n rotary_seq_len = max(kv_seq_len, int(position_ids.max().item()) + 1)\n self.rotary_emb.extend_rope_embedding(Vn, rotary_seq_len + 1) # +1 slack\n cos, sin = self.rotary_emb.get_cached(rotary_seq_len, Qn.device.index or 0)\n\n cos = cos[position_ids].unsqueeze(1).to(device = Qn.device, dtype = Qn.dtype)\n sin = sin[position_ids].unsqueeze(1).to(device = Qn.device, dtype = Qn.dtype)\n\n h = self.half_head_dim\n\n RH_Q = self.RH_Q\n RH_Q[:, :, :, :h] = Qn[:, :, :, h:]\n RH_Q[:, :, :, h:] = Qn[:, :, :, :h]\n RH_Q[:, :, :, :h].neg_() # torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h])\n Qn *= cos\n Qn.addcmul_(RH_Q, sin)\n\n RH_K = RH_Q[\n :, :n_kv_heads, :, :\n ] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = \"cuda:0\")\n RH_K[:, :, :, :h] = Kn[:, :, :, h:]\n RH_K[:, :, :, h:] = Kn[:, :, :, :h]\n RH_K[:, :, :, :h].neg_() # torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h])\n Kn *= cos\n Kn.addcmul_(RH_K, sin)\n\n # New KV cache\n # Kn = torch.cat([K1, Kn], dim = 2)\n # Vn = torch.cat([V1, Vn], dim = 2)\n self.paged_attention_K[seq_len] = Kn.permute(2, 0, 1, 3)\n self.paged_attention_V[seq_len] = Vn.permute(2, 0, 1, 3)\n Kn = self.paged_attention_K[:kv_seq_len].permute(1, 2, 0, 3)\n Vn = self.paged_attention_V[:kv_seq_len].permute(1, 2, 0, 3)\n\n # Handle sliding windows\n sliding_window = getattr(self.config, \"sliding_window\", None)\n if sliding_window is not None and kv_seq_len > sliding_window:\n # From https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py#L193\n start = kv_seq_len - sliding_window\n Knn = Kn[:, :, start:, :] # .contiguous()\n Vnn = Vn[:, :, start:, :] # .contiguous()\n if attention_mask is not None:\n attention_mask = attention_mask[..., start:]\n else:\n Knn, Vnn = Kn, Vn\n\n # Grouped query attention\n _, _, cached_len, _ = Knn.shape\n if bsz == 1 or ((not SDPA_HAS_GQA) and n_groups != 1):\n Knn = Knn[:, :, None, :, :].expand(\n bsz, n_kv_heads, n_groups, cached_len, head_dim\n )\n Vnn = Vnn[:, :, None, :, :].expand(\n bsz, n_kv_heads, n_groups, cached_len, head_dim\n )\n Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)\n Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim)\n\n # when qlen==vlen and attn_mask is None, we should use causal attention\n Q_len = Qn.shape[-2]\n K_len = Knn.shape[-2]\n if attention_mask is None and Q_len == K_len:\n is_causal = True\n else:\n is_causal = False\n # Attention\n if bsz == 1:\n Qn *= self.scalar # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963\n # It seems like doing (Q * scalar) @ K is better than (Q @ K) * scalar to stop overflows\n A = torch_matmul(\n Qn, Knn.transpose(2, 3), out = self.attention[:, :, :, :cached_len]\n )\n A[:] = torch_nn_functional_softmax(\n A, dim = -1, dtype = torch.float32\n ) # .to(A.dtype)\n A = torch_matmul(A, Vnn, out = Qn)\n # --- attention_mask fixup for SDPA if user passes 2D padding mask\n else:\n if attention_mask is not None and attention_mask.dim() == 2:\n attention_mask = attention_mask[:, None, None, :].to(torch.bool)\n # is it more appropriate to use _prepare_4d_causal_attention_mask_for_sdpa?\n elif (\n attention_mask is not None\n and attention_mask.dim() == 4\n and attention_mask.dtype != torch.bool\n ):\n # Decode is more stable with boolean keep masks than additive bf16 masks.\n attention_mask = attention_mask.eq(0)\n\n if SDPA_HAS_GQA:\n A = scaled_dot_product_attention(\n Qn,\n Knn,\n Vnn,\n attn_mask = attention_mask,\n is_causal = is_causal,\n enable_gqa = True,\n )\n else:\n A = scaled_dot_product_attention(\n Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = is_causal\n )\n A = A.transpose(1, 2)\n A = A.reshape(bsz, 1, attention_size)\n A = fast_linear_forward(self.o_proj, A, out = self.temp_O)\n return A, (Kn, Vn)\n\n\ntorch_nn_functional_silu = torch.nn.functional.silu\n\n\ndef fast_swiglu_inference(\n self, X, temp_gate = None, temp_up = None, gate_multiplier = None, down_multiplier = None\n):\n # gate = self.gate_proj(X)\n # up = self.up_proj(X)\n bsz, _, hd = X.shape\n # mlp_size = self.config.intermediate_size\n # temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = \"cuda:0\")\n\n gate = fast_linear_forward(self.gate_proj, X, out = temp_gate)\n\n if gate_multiplier is not None:\n gate *= gate_multiplier\n\n up = fast_linear_forward(self.up_proj, X, out = temp_up)\n\n gate = torch_nn_functional_silu(gate, inplace = True)\n gate *= up\n\n # X = self.down_proj(gate)\n down = fast_linear_forward(self.down_proj, gate, out = up[:, :, :hd])\n\n if down_multiplier is not None:\n down *= down_multiplier\n\n return down\n\n\ntorch_square = torch.square\ntorch_mean = torch.mean\n\n\ndef fast_rms_layernorm_inference(self, X, XX = None, XX2 = None, variance = None):\n old_dtype = X.dtype\n if XX is None:\n XX = X.to(torch.float32)\n variance = XX.square().mean(-1, keepdim = True)\n else:\n XX.copy_(X)\n torch_mean(torch_square(XX, out = XX2), -1, keepdim = True, out = variance)\n variance += self.variance_epsilon\n XX *= variance.rsqrt_()\n\n if XX is None:\n X = XX.to(old_dtype)\n else:\n X.copy_(XX)\n\n X *= self.weight\n return X\n\n\ndef fast_rms_layernorm_inference_gemma(self, X, out_weight = None):\n XX = X.to(torch.float32)\n variance = XX.square().mean(-1, keepdim = True)\n variance += self.variance_epsilon\n XX *= variance.rsqrt_()\n\n if out_weight is None:\n out_weight = self.weight + 1.0\n else:\n out_weight[:] = self.weight\n out_weight += 1.0\n\n XX *= out_weight\n return XX.to(X.dtype)\n\n\n# Normal layernorm with mean removal\n@torch.compile(fullgraph = False, dynamic = True, options = torch_compile_options)\ndef fast_layernorm_compiled(layernorm, X):\n old_dtype = X.dtype\n X = X.float()\n mean = X.mean(-1, keepdim = True)\n Xbar = X - mean\n X = (\n Xbar\n * torch.rsqrt(Xbar.square().mean(-1, keepdim = True) + layernorm.variance_epsilon)\n * layernorm.weight.float()\n )\n return X.to(old_dtype)\n\n\n# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L320\ndef LlamaAttention_fast_forward(\n self,\n hidden_states: torch.Tensor,\n causal_mask: Optional[BlockDiagonalCausalMask] = None,\n attention_mask: Optional[torch.Tensor] = None,\n position_ids: Optional[torch.LongTensor] = None,\n past_key_value: Optional[Tuple[torch.Tensor]] = None,\n output_attentions: bool = False,\n use_cache: bool = False,\n padding_mask: Optional[torch.LongTensor] = None,\n position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,\n *args,\n **kwargs,\n) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:\n # Clear inference\n if hasattr(self, \"paged_attention\"):\n del self.paged_attention_K\n del self.paged_attention_V\n del self.paged_attention\n del self.temp_QA\n del self.temp_KV\n del self.RH_Q\n del self.attention\n bsz, q_len, _ = hidden_states.size()\n\n n_heads = self.config.num_attention_heads\n n_groups = self.num_key_value_groups\n n_kv_heads = self.config.num_key_value_heads\n head_dim = self.head_dim\n assert n_kv_heads * n_groups == n_heads\n\n Q, K, V = self.apply_qkv(self, hidden_states)\n Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)\n K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)\n V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)\n seq_info = get_packed_info_from_kwargs(kwargs, Q.device)\n\n kv_seq_len = K.shape[-2]\n if past_key_value is not None:\n kv_seq_len += past_key_value[0].shape[-2]\n\n if position_embeddings and kv_seq_len <= position_embeddings[0].shape[0]:\n cos, sin = position_embeddings\n else:\n rotary_emb = self.rotary_emb\n rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)\n cos, sin = rotary_emb.get_cached(kv_seq_len, Q.device.index)\n cos = cos.to(device = Q.device, dtype = Q.dtype)\n sin = sin.to(device = Q.device, dtype = Q.dtype)\n\n rope_position_ids = position_ids\n if rope_position_ids is None and seq_info is not None:\n rope_position_ids = kwargs.get(\"position_ids\")\n\n # Q, K = (\n # fast_rope_embedding(Q, K, cos, sin)\n # if rope_position_ids is None\n # else inplace_rope_embedding(Q, K, cos, sin, rope_position_ids)\n # )\n Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)\n\n if past_key_value is not None:\n K = torch.cat([past_key_value[0], K], dim = 2)\n V = torch.cat([past_key_value[1], V], dim = 2)\n past_key_value = (K, V) if use_cache else None\n\n # Attention module\n use_varlen = seq_info is not None and past_key_value is None\n backend = (\n SDPA if attention_mask is not None else select_attention_backend(use_varlen)\n )\n\n # should dropout be hardcoded to 0.0?\n config = AttentionConfig(\n backend = backend,\n n_kv_heads = n_kv_heads,\n n_groups = n_groups,\n flash_dense_kwargs = {\"causal\": True},\n flash_varlen_kwargs = {\"dropout_p\": 0.0, \"causal\": True},\n )\n context = AttentionContext(\n bsz = bsz,\n q_len = q_len,\n kv_seq_len = kv_seq_len,\n n_heads = n_heads,\n head_dim = head_dim,\n requires_grad = hidden_states.requires_grad,\n seq_info = seq_info,\n attention_mask = attention_mask,\n causal_mask = causal_mask,\n )\n\n A = run_attention(config = config, context = context, Q = Q, K = K, V = V)\n attn_output = A.reshape(bsz, q_len, n_heads * head_dim)\n attn_output = self.apply_o(self, attn_output)\n attn_weights = None\n return attn_output, attn_weights, past_key_value\n\n\n# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L590\ndef LlamaDecoderLayer_fast_forward(\n self,\n hidden_states: torch.Tensor,\n causal_mask = None,\n attention_mask: Optional[torch.Tensor] = None,\n position_ids: Optional[torch.LongTensor] = None,\n past_key_value: Optional[Tuple[torch.Tensor]] = None,\n output_attentions: Optional[bool] = False,\n use_cache: Optional[bool] = False,\n padding_mask: Optional[torch.LongTensor] = None,\n position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,\n *args,\n **kwargs,\n) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:\n \"\"\"\n Args:\n hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`\n attention_mask (`torch.FloatTensor`, *optional*): attention mask of size\n `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.\n output_attentions (`bool`, *optional*):\n Whether or not to return the attentions tensors of all attention layers. See `attentions` under\n returned tensors for more detail.\n use_cache (`bool`, *optional*):\n If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding\n (see `past_key_values`).\n past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states\n \"\"\"\n if use_cache and hasattr(self, \"_flag_for_generation\"):\n residual = hidden_states\n hidden_states = fast_rms_layernorm_inference(\n self.input_layernorm, hidden_states\n )\n hidden_states, self_attn_weights, present_key_value = self.self_attn(\n hidden_states = hidden_states,\n causal_mask = causal_mask,\n attention_mask = attention_mask,\n position_ids = position_ids,\n past_key_value = past_key_value,\n output_attentions = output_attentions,\n use_cache = use_cache,\n padding_mask = padding_mask,\n position_embeddings = position_embeddings,\n **kwargs,\n )\n hidden_states += residual\n\n # Fully Connected\n residual = hidden_states\n hidden_states = fast_rms_layernorm_inference(\n self.post_attention_layernorm, hidden_states\n )\n hidden_states = fast_swiglu_inference(self.mlp, hidden_states)\n hidden_states += residual\n else:\n residual = hidden_states\n hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states)\n hidden_states, self_attn_weights, present_key_value = self.self_attn(\n hidden_states = hidden_states,\n causal_mask = causal_mask,\n attention_mask = attention_mask,\n position_ids = position_ids,\n past_key_value = past_key_value,\n output_attentions = output_attentions,\n use_cache = use_cache,\n padding_mask = padding_mask,\n position_embeddings = position_embeddings,\n **kwargs,\n )\n hidden_states = residual + hidden_states\n\n # Fully Connected\n residual = hidden_states\n hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states)\n hidden_states = self.mlp(hidden_states)\n hidden_states = residual + hidden_states\n\n outputs = (hidden_states,)\n if output_attentions:\n outputs += (self_attn_weights,)\n if use_cache:\n outputs += (present_key_value,)\n return outputs\n\n\n# https://github.com/unslothai/unsloth/issues/404#issuecomment-2323473452\n__DTYPE_MAP = {\n \"float32\": torch.float32,\n torch.float32: torch.float32,\n \"float16\": torch.float16,\n torch.float16: torch.float16,\n \"bfloat16\": torch.bfloat16,\n torch.bfloat16: torch.bfloat16,\n}\n\n\n# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L825\ndef LlamaModel_fast_forward(\n self,\n input_ids: Optional[torch.LongTensor] = None,\n causal_mask: Optional[BlockDiagonalCausalMask] = None,\n attention_mask: Optional[torch.Tensor] = None,\n position_ids: Optional[torch.LongTensor] = None,\n past_key_values: Optional[List[torch.FloatTensor]] = None,\n inputs_embeds: Optional[torch.FloatTensor] = None,\n use_cache: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n return_dict: Optional[bool] = None,\n *args,\n **kwargs,\n) -> Union[Tuple, BaseModelOutputWithPast]:\n output_attentions = (\n output_attentions\n if output_attentions is not None\n else self.config.output_attentions\n )\n assert output_attentions is False\n output_hidden_states = (\n output_hidden_states\n if output_hidden_states is not None\n else self.config.output_hidden_states\n )\n use_cache = use_cache if use_cache is not None else self.config.use_cache\n\n return_dict = (\n return_dict if return_dict is not None else self.config.use_return_dict\n )\n\n # retrieve input_ids and inputs_embeds\n if input_ids is not None and inputs_embeds is not None:\n raise ValueError(\n \"Unsloth: You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time\"\n )\n elif input_ids is not None:\n batch_size, seq_length = input_ids.shape\n elif inputs_embeds is not None:\n batch_size, seq_length, _ = inputs_embeds.shape\n else:\n raise ValueError(\n \"Unsloth: You have to specify either decoder_input_ids or decoder_inputs_embeds\"\n )\n\n seq_length_with_past = seq_length\n\n # Fix out of bounds tokenization unless we were given packed metadata\n allow_overlength = getattr(self, \"_unsloth_allow_packed_overlength\", False) or (\n \"packed_seq_lengths\" in kwargs\n )\n if hasattr(self, \"max_seq_length\") and not allow_overlength:\n if seq_length > self.max_seq_length:\n shape = input_ids.shape if input_ids is not None else inputs_embeds.shape\n logger.warning_once(\n f\"Unsloth: Input IDs of shape {shape} with length {seq_length} > the model's max sequence length of {self.max_seq_length}.\\n\"\n \"We shall truncate it ourselves. It's imperative if you correct this issue first.\"\n )\n if input_ids is not None:\n input_ids = input_ids[:, : self.max_seq_length]\n elif inputs_embeds is not None:\n inputs_embeds = inputs_embeds[:, : self.max_seq_length, :]\n if (\n attention_mask is not None\n and attention_mask.shape[-1] > self.max_seq_length\n ):\n attention_mask = attention_mask[:, : self.max_seq_length]\n\n past_key_values_length = 0\n\n if past_key_values is not None:\n past_key_values_length = past_key_values[0][0].shape[2]\n seq_length_with_past = seq_length_with_past + past_key_values_length\n\n # We already handle KV cache position_ids ourselves.\n if False: # (past_key_values_length != 0):\n position_ids = torch.arange(\n past_key_values_length,\n seq_length + past_key_values_length,\n dtype = torch.int32,\n device = f\"{DEVICE_TYPE_TORCH}:0\",\n )\n position_ids = position_ids.unsqueeze(0).view(-1, seq_length)\n elif position_ids is not None:\n position_ids = position_ids.view(-1, seq_length).to(torch.int32) # .long()\n else:\n position_ids = None\n\n if position_ids is not None:\n if position_ids.shape[0] != batch_size:\n position_ids = position_ids.repeat((batch_size, 1))\n\n # Embed positions\n if inputs_embeds is None:\n inputs_embeds = self.embed_tokens(input_ids)\n\n inputs_embeds = inputs_embeds.to(_get_dtype(dtype_from_config(self.config)))\n\n # Normalized from Gemma\n IS_GEMMA = self.config.model_type.startswith(\"gemma\")\n IS_GEMMA2 = self.config.model_type.startswith(\"gemma2\")\n IS_COHERE = self.config.model_type.startswith(\"cohere\")\n IS_GRANITE = self.config.model_type.startswith(\"granite\")\n IS_FALCON_H1 = self.config.model_type.startswith(\"falcon_h1\")\n\n train_embed_tokens = self.embed_tokens.weight.requires_grad\n\n if IS_GEMMA:\n # Match Gemma exactly by casting to bfloat16 / float16\n # inputs_embeds *= math_sqrt(self.config.hidden_size)\n # Ie 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32\n # & 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32\n normalizer = torch.tensor(\n math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype\n )\n\n if train_embed_tokens:\n # Careful we must not do an inplace op!\n inputs_embeds = inputs_embeds * normalizer\n else:\n inputs_requires_grad = inputs_embeds.requires_grad\n if not inputs_embeds.is_leaf:\n inputs_embeds = inputs_embeds.detach()\n inputs_requires_grad = True\n elif inputs_requires_grad:\n inputs_embeds.requires_grad_(False)\n inputs_embeds *= normalizer\n # inputs_embeds *= math_sqrt(self.config.hidden_size)\n if inputs_requires_grad:\n inputs_embeds.requires_grad_(True)\n\n # Fix up attention mask by setting elements to 0\n # Specifically for DPO\n if (\n getattr(self, \"_has_no_labels\", False) is True\n and (attention_mask is not None)\n and attention_mask.ndim == 2\n and (past_key_values is None)\n and (not train_embed_tokens)\n and self.training\n ):\n # Careful for inference the attention_mask is size (1, kv_seq_len)\n # Whilst the input_embeds is size (1, 1, 4096)\n inputs_requires_grad = inputs_embeds.requires_grad\n if not inputs_embeds.is_leaf:\n inputs_embeds = inputs_embeds.detach()\n inputs_requires_grad = True\n elif inputs_requires_grad:\n inputs_embeds.requires_grad_(False)\n attention_mask = attention_mask[:, : self.max_seq_length] # Must resize!\n inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)\n if inputs_requires_grad:\n inputs_embeds.requires_grad_(True)\n\n # Ignore attention_mask\n if attention_mask is None:\n padding_mask = None\n elif self.training:\n attention_mask = None\n padding_mask = None\n else:\n # if 0 in attention_mask:\n # padding_mask = attention_mask\n # else:\n padding_mask = None\n\n attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(\n attention_mask,\n (batch_size, seq_length),\n inputs_embeds,\n past_key_values_length,\n sliding_window = getattr(self.config, \"sliding_window\", None),\n )\n # Must NOT convert to bool - weirdly this causes stuff to error out!\n # if attention_mask is not None:\n # attention_mask = attention_mask.to(torch.bool)\n\n hidden_states = inputs_embeds\n if IS_GRANITE or IS_FALCON_H1: # granite has embedding multiplier\n hidden_states = self.config.embedding_multiplier * hidden_states\n\n if past_key_values is None and self.training:\n use_cache = False\n # if use_cache:\n # logger.warning_once(\n # \"Unsloth: `use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`\"\n # )\n # use_cache = False\n\n # decoder layers\n all_hidden_states = () if output_hidden_states else None\n all_self_attns = () if output_attentions else None\n next_decoder_cache = () if use_cache else None\n\n # Gradient checkpointing methods (ie sqrt)\n if hasattr(self, \"_gradient_checkpointing_boundaries\"):\n boundaries = self._gradient_checkpointing_boundaries\n else:\n boundaries = None\n\n # Check checkpointing method\n gradient_checkpointing = False\n\n if self.gradient_checkpointing and self.training and not use_cache:\n gradient_checkpointing = True\n\n # Gemma2 has alternating SWA and global attn\n use_static_mask = True\n dynamic_SWA_mask = None\n dynamic_GA_mask = None\n if IS_GEMMA2:\n if HAS_FLASH_ATTENTION_SOFTCAPPING and attention_mask is None:\n self.SWA_mask = True\n self.GA_mask = False\n elif attention_mask is not None:\n # Fixes https://github.com/unslothai/unsloth/issues/853\n # Unsloth needs a 2D mask, not a [2, 1, n, n] mask!\n\n # https://github.com/pytorch/pytorch/issues/103749\n # Need to convert to float and not using bool\n # attention_mask = (1.0 - attention_mask.float()) * torch.finfo(inputs_embeds.dtype).min\n dynamic_SWA_mask = _prepare_4d_causal_attention_mask_for_sdpa(\n attention_mask,\n (batch_size, seq_length),\n inputs_embeds,\n past_key_values_length,\n sliding_window = self.config.sliding_window,\n )\n dynamic_GA_mask = _prepare_4d_causal_attention_mask_for_sdpa(\n attention_mask,\n (batch_size, seq_length),\n inputs_embeds,\n past_key_values_length,\n sliding_window = None,\n )\n use_static_mask = False\n\n elif not hasattr(self, \"SWA_mask\"):\n if HAS_FLEX_ATTENTION:\n # Use Flex Attention instead!\n self.SWA_mask = create_flex_attention_sliding_window_mask(\n self.max_seq_length, self.config.sliding_window\n )\n self.GA_mask = create_flex_attention_causal_mask(self.max_seq_length)\n else:\n n = self.max_seq_length # self.config.max_position_embeddings\n # masked_fill is making stuff slower!\n # self. GA_mask = create_boolean_mask(n = n, sliding_window = 0)\n # self.SWA_mask = create_boolean_mask(n = n, sliding_window = self.config.sliding_window)\n from transformers.modeling_attn_mask_utils import AttentionMaskConverter\n\n self.SWA_mask = (\n AttentionMaskConverter(\n is_causal = True,\n sliding_window = self.config.sliding_window,\n )\n .to_causal_4d(\n 1,\n n,\n n,\n dtype = inputs_embeds.dtype,\n device = DEVICE_TYPE_TORCH,\n )\n .squeeze(0)\n .squeeze(0)\n )\n\n self.GA_mask = (\n AttentionMaskConverter(\n is_causal = True,\n )\n .to_causal_4d(\n 1,\n n,\n n,\n dtype = inputs_embeds.dtype,\n device = DEVICE_TYPE_TORCH,\n )\n .squeeze(0)\n .squeeze(0)\n )\n pass\n\n if (\n IS_ATTENTION_REFACTOR\n and (\n hasattr(self, \"rotary_emb\")\n or not hasattr(self.layers[0].self_attn, \"rotary_emb\")\n )\n ) or IS_GRANITE:\n # Transformers main has made it mandatory to pass position_embeddings\n # https://github.com/huggingface/transformers/pull/34858\n # Also, transformers 4.45.0 supports granite but with the attention refactor (it always had the refactor)\n # unsloth's check for granite too has \"version >= 4.45.0 (rightly so)\".\n # so let granite always use the attention refactor implementation.\n\n self.rotary_emb.extend_rope_embedding(\n hidden_states, self.config.max_position_embeddings\n )\n position_embeddings = self.rotary_emb.get_cached(\n self.config.max_position_embeddings, hidden_states.device.index\n )\n else:\n position_embeddings = None\n\n # Go through every layer!\n for idx, decoder_layer in enumerate(self.layers):\n if output_hidden_states:\n all_hidden_states += (hidden_states,)\n past_key_value = past_key_values[idx] if past_key_values is not None else None\n\n mask = causal_mask\n if IS_GEMMA2:\n use_sliding_window = idx % 2 == 0\n if use_sliding_window:\n mask = self.SWA_mask if use_static_mask else dynamic_SWA_mask\n else:\n mask = self.GA_mask if use_static_mask else dynamic_GA_mask\n kwargs[\"use_sliding_window\"] = use_sliding_window\n\n if gradient_checkpointing and not isinstance(\n decoder_layer, GradientCheckpointingLayer\n ):\n\n def create_custom_forward(module):\n def custom_forward(*inputs):\n return module(\n *inputs,\n past_key_value,\n output_attentions,\n padding_mask = padding_mask,\n position_embeddings = position_embeddings,\n **kwargs,\n )\n\n return custom_forward\n\n layer_outputs = torch.utils.checkpoint.checkpoint(\n create_custom_forward(decoder_layer),\n hidden_states,\n mask,\n attention_mask,\n position_ids,\n use_reentrant = True,\n preserve_rng_state = False,\n )\n hidden_states = layer_outputs[0]\n\n else:\n layer_outputs = decoder_layer(\n hidden_states,\n causal_mask = mask,\n attention_mask = attention_mask,\n position_ids = position_ids,\n past_key_value = past_key_value,\n output_attentions = output_attentions,\n use_cache = use_cache,\n padding_mask = padding_mask,\n position_embeddings = position_embeddings,\n **kwargs,\n )\n hidden_states = layer_outputs[0]\n\n if use_cache:\n next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)\n if output_attentions:\n all_self_attns += (layer_outputs[1],)\n\n # Final layernorm\n if use_cache:\n if IS_FALCON_H1:\n hidden_states = fast_rms_layernorm_inference(\n self.final_layernorm, hidden_states\n )\n else:\n hidden_states = (\n fast_rms_layernorm_inference_gemma\n if IS_GEMMA\n else fast_rms_layernorm_inference\n )(self.norm, hidden_states)\n elif IS_COHERE:\n hidden_states = self.norm(hidden_states)\n elif IS_FALCON_H1:\n hidden_states = fast_rms_layernorm(\n self.final_layernorm, hidden_states, gemma = IS_GEMMA\n )\n else:\n hidden_states = fast_rms_layernorm(self.norm, hidden_states, gemma = IS_GEMMA)\n\n if output_hidden_states:\n all_hidden_states += (hidden_states,)\n next_cache = next_decoder_cache if use_cache else None\n\n if not return_dict:\n return tuple(\n v\n for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]\n if v is not None\n )\n return BaseModelOutputWithPast(\n last_hidden_state = hidden_states,\n past_key_values = next_cache,\n hidden_states = all_hidden_states,\n attentions = all_self_attns,\n )\n\n\n# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L825\ndef _LlamaModel_fast_forward_inference(\n attention_fast_forward_inference = LlamaAttention_fast_forward_inference,\n mlp_fast_forward_inference = fast_swiglu_inference,\n):\n # This makes the attention and MLP customisable.\n # Now for models like qwen3 or cohere which use custom attention operations, we can use this function\n def LlamaModel_fast_forward_inference_custom(\n self,\n input_ids,\n past_key_values,\n position_ids,\n attention_mask = None,\n **kwargs,\n ):\n input_ids = input_ids[:, : self.max_seq_length]\n bsz, q_len = input_ids.shape\n hd = self.config.hidden_size\n mlp_size = self.config.intermediate_size\n\n X = self.model.embed_tokens(input_ids)\n X = X.to(_get_dtype(dtype_from_config(self.config)))\n bsz, q_len, hd = X.shape\n assert q_len == 1\n # Get saved buffers to reduce memory movement\n residual = torch.empty(\n (bsz, q_len, hd), dtype = torch.float32, device = f\"{DEVICE_TYPE_TORCH}:0\"\n )\n _XX = torch.empty(\n (2, bsz, q_len, hd), dtype = torch.float32, device = f\"{DEVICE_TYPE_TORCH}:0\"\n )\n XX, XX2 = _XX[0], _XX[1]\n variance = torch.empty(\n (bsz, q_len, 1), dtype = torch.float32, device = f\"{DEVICE_TYPE_TORCH}:0\"\n )\n temp_mlp = torch.empty(\n (2, bsz, 1, mlp_size), dtype = X.dtype, device = f\"{DEVICE_TYPE_TORCH}:0\"\n )\n temp_gates, temp_ups = (\n tuple(temp_mlp[0].to(torch.device(x)) for x in range(DEVICE_COUNT)),\n tuple(temp_mlp[1].to(torch.device(x)) for x in range(DEVICE_COUNT)),\n )\n\n seq_len = past_key_values[0][0].shape[-2]\n kv_seq_len = seq_len + 1\n if attention_mask is not None:\n attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(\n attention_mask,\n (bsz, q_len),\n X,\n seq_len,\n sliding_window = getattr(self.config, \"sliding_window\", None),\n )\n # Pre-convert to bool once for all layers (avoids per-layer .eq(0))\n if attention_mask is not None and attention_mask.dtype != torch.bool:\n attention_mask = attention_mask.eq(0)\n else:\n attention_mask = None\n\n # Compute rotary_seq_len once to avoid per-layer GPU-CPU sync from .item()\n rotary_seq_len = max(kv_seq_len, int(position_ids.max().item()) + 1)\n\n next_decoder_cache = []\n\n for idx, decoder_layer in enumerate(self.model.layers):\n device_index = getattr(decoder_layer, \"_per_layer_device_index\", 0)\n X, residual, position_ids = move_to_device(\n device_index, X, residual, position_ids\n )\n residual.copy_(X) # residual = X\n X = fast_rms_layernorm_inference(\n decoder_layer.input_layernorm,\n X,\n XX = XX,\n XX2 = XX2,\n variance = variance,\n )\n X, present_key_value = attention_fast_forward_inference(\n decoder_layer.self_attn,\n hidden_states = X,\n past_key_value = past_key_values[idx],\n position_ids = position_ids,\n attention_mask = attention_mask,\n do_prefill = not hasattr(decoder_layer.self_attn, \"paged_attention\"),\n rotary_seq_len = rotary_seq_len,\n )\n X += residual\n\n residual.copy_(X) # residual = X\n X = fast_rms_layernorm_inference(\n decoder_layer.post_attention_layernorm,\n X,\n XX = XX,\n XX2 = XX2,\n variance = variance,\n )\n X = mlp_fast_forward_inference(\n decoder_layer.mlp,\n X,\n temp_gate = temp_gates[device_index],\n temp_up = temp_ups[device_index],\n )\n X += residual\n\n next_decoder_cache.append(present_key_value)\n X = fast_rms_layernorm_inference(\n self.model.norm,\n X,\n XX = XX,\n XX2 = XX2,\n variance = variance,\n )\n\n return BaseModelOutputWithPast(\n last_hidden_state = X,\n past_key_values = next_decoder_cache,\n hidden_states = [],\n attentions = [],\n )\n\n return LlamaModel_fast_forward_inference_custom\n\n\n# For ensuring backwards compatibility, we create LlamaModel_fast_forward_inference that is consumed by other models\nLlamaModel_fast_forward_inference = _LlamaModel_fast_forward_inference()\n\n\ndef CausalLM_fast_forward(fast_forward_inference):\n def _CausalLM_fast_forward(\n self,\n input_ids: torch.LongTensor = None,\n causal_mask: Optional[BlockDiagonalCausalMask] = None,\n attention_mask: Optional[torch.Tensor] = None,\n position_ids: Optional[torch.LongTensor] = None,\n past_key_values: Optional[List[torch.FloatTensor]] = None,\n inputs_embeds: Optional[torch.FloatTensor] = None,\n labels: Optional[torch.LongTensor] = None,\n use_cache: Optional[bool] = None,\n output_attentions: Optional[bool] = None,\n output_hidden_states: Optional[bool] = None,\n return_dict: Optional[bool] = None,\n num_logits_to_keep: Optional[int] = 0,\n logits_to_keep: Optional[int] = 0,\n *args,\n **kwargs,\n ) -> Union[Tuple, CausalLMOutputWithPast]:\n if past_key_values is not None:\n outputs = fast_forward_inference(\n self,\n input_ids,\n past_key_values,\n position_ids = position_ids,\n attention_mask = attention_mask,\n **kwargs,\n )\n else:\n causal_mask = (\n xformers.attn_bias.LowerTriangularMask() if HAS_XFORMERS else None\n )\n\n output_attentions = (\n output_attentions\n if output_attentions is not None\n else self.config.output_attentions\n )\n output_hidden_states = (\n output_hidden_states\n if output_hidden_states is not None\n else self.config.output_hidden_states\n )\n return_dict = (\n return_dict if return_dict is not None else self.config.use_return_dict\n )\n # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)\n self.model._has_no_labels = labels is None\n outputs = self.model(\n input_ids = input_ids,\n causal_mask = causal_mask,\n attention_mask = attention_mask,\n position_ids = position_ids,\n past_key_values = past_key_values,\n inputs_embeds = inputs_embeds,\n use_cache = use_cache,\n output_attentions = output_attentions,\n output_hidden_states = output_hidden_states,\n return_dict = return_dict,\n **kwargs,\n )\n hidden_states = outputs[0]\n\n bsz, q_len, hd = hidden_states.shape\n lm_head = self.lm_head.weight\n lm_head_device = lm_head.device\n\n logit_softcapping = getattr(self.config, \"final_logit_softcapping\", 0)\n logit_scaling = getattr(self.config, \"logit_scale\", 0)\n dtype = lm_head.dtype\n num_logits_to_keep = max(num_logits_to_keep, logits_to_keep)\n\n # Move items to same device as lm_head\n hidden_states = hidden_states.to(lm_head_device)\n if labels is not None:\n labels = labels.to(lm_head_device)\n\n # Output last hidden states without logits if asked\n if os.environ.get(\"UNSLOTH_RETURN_HIDDEN_STATES\", \"0\") == \"1\":\n if num_logits_to_keep != 0:\n hidden_states = hidden_states[:, -num_logits_to_keep:, :]\n return CausalLMOutputWithPast(\n loss = None,\n logits = hidden_states,\n past_key_values = outputs.past_key_values,\n hidden_states = outputs.hidden_states,\n attentions = outputs.attentions,\n )\n\n if bsz == 1 and q_len == 1:\n logits = torch.mv(lm_head, hidden_states.ravel().to(dtype))\n logits = logits.unsqueeze(0).unsqueeze(0)\n elif num_logits_to_keep != 0:\n logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :].to(dtype))\n else:\n RETURN_LOGITS = os.environ.get(\"UNSLOTH_RETURN_LOGITS\", \"0\") == \"1\"\n # < 1024 Normal Unsloth uses less VRAM!\n if bsz * q_len <= 1024 and not RETURN_LOGITS:\n # Use unsloth_fused_ce_loss which actually calculates the best chunk size to reduce VRAM usage\n RETURN_LOGITS = False\n\n if not RETURN_LOGITS and labels is not None:\n n_items = kwargs.get(\"num_items_in_batch\", None)\n if n_items is None:\n n_items = kwargs.get(\"n_items\", None)\n\n if self.config.model_type == \"falcon_h1\":\n hidden_states = hidden_states * self.config.lm_head_multiplier\n\n ### DISABLED since T4 breaks\n # OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.\n # loss = fused_linear_cross_entropy(\n # hidden_states = hidden_states,\n # lm_weight = lm_head,\n # labels = labels,\n # num_items_in_batch = n_items,\n # logit_softcapping = logit_softcapping,\n # )\n loss = unsloth_fused_ce_loss(\n trainer = None,\n hidden_states = hidden_states,\n lm_head_weight = lm_head,\n lm_head_bias = None,\n labels = labels,\n mask = None,\n n_items = n_items,\n scaling = getattr(self, \"accelerator_scaler\", None),\n target_gb = None,\n torch_compile = True,\n logit_softcapping = logit_softcapping,\n )\n if not return_dict:\n output = (logits,) + outputs[1:]\n return (loss,) + output if loss is not None else output\n\n output = CausalLMOutputWithPast(\n loss = loss,\n logits = EMPTY_LOGITS,\n past_key_values = outputs.past_key_values,\n hidden_states = outputs.hidden_states,\n attentions = outputs.attentions,\n )\n return output\n pass\n logits = self.lm_head(hidden_states.to(dtype))\n\n logits = logits.to(_get_dtype(dtype_from_config(self.config)))\n loss = None\n logit_softcapping = getattr(self.config, \"final_logit_softcapping\", 0)\n logit_scaling = getattr(self.config, \"logit_scale\", 0)\n if self.config.model_type == \"granite\":\n # granite uses logit_scaling as key and they divide by the scale unlike cohere\n # notice that for granite, logits_scale is 16 and for cohere it is 0.125 (aka 1/8) in their respective configs\n # granite: https://github.com/huggingface/transformers/blob/4d1d0f29a493098e6bc6b904b82e29cb331827f5/src/transformers/models/granite/modeling_granite.py#L1103\n # cohere: https://github.com/huggingface/transformers/blob/4d1d0f29a493098e6bc6b904b82e29cb331827f5/src/transformers/models/cohere/modeling_cohere.py#L1176\n logit_scaling = 1 / getattr(self.config, \"logits_scaling\", 1)\n elif self.config.model_type == \"falcon_h1\":\n logit_scaling = self.config.lm_head_multiplier\n\n if labels is not None:\n shift_logits = logits\n # if not hasattr(self, \"extra_ignored_labels\"):\n # # Fixes https://github.com/unslothai/unsloth/issues/10\n # self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = \"cuda:0\")\n # pass\n shift_labels = torch.empty_like(labels)\n shift_labels[..., :-1] = labels[..., 1:]\n shift_labels[..., -1] = -100\n mask_packed_sequence_boundaries(\n shift_labels,\n kwargs.get(\"packed_seq_lengths\"),\n )\n # shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))\n n_items = kwargs.get(\"num_items_in_batch\", None)\n if n_items is None:\n n_items = kwargs.get(\"n_items\", None)\n loss = fast_cross_entropy_loss(\n logits = shift_logits,\n labels = shift_labels,\n logit_softcapping = logit_softcapping,\n logit_scaling = logit_scaling,\n n_items = n_items,\n )\n else:\n if logit_scaling != 0:\n if logits.requires_grad:\n logits = logit_scaling * logits\n else:\n logits *= logit_scaling\n if logit_softcapping != 0:\n if logits.requires_grad:\n logits = (1.0 / logit_softcapping) * logits\n logits = torch.tanh(logits)\n logits = logit_softcapping * logits\n else:\n logits *= 1.0 / logit_softcapping\n logits.tanh_()\n logits *= logit_softcapping\n\n if not return_dict:\n output = (logits,) + outputs[1:]\n return (loss,) + output if loss is not None else output\n return CausalLMOutputWithPast(\n loss = loss,\n logits = logits,\n past_key_values = outputs.past_key_values,\n hidden_states = outputs.hidden_states,\n attentions = outputs.attentions,\n )\n\n return _CausalLM_fast_forward\n\n\n@torch._disable_dynamo\ndef PeftModel_fast_forward(\n self,\n input_ids = None,\n causal_mask = None,\n attention_mask = None,\n inputs_embeds = None,\n labels = None,\n output_attentions = None,\n output_hidden_states = None,\n return_dict = None,\n task_ids = None,\n num_logits_to_keep = 0,\n logits_to_keep = 0,\n **kwargs,\n):\n is_classification = \"Classification\" in str(type(self.base_model.model))\n if is_classification:\n return self.base_model(\n input_ids = input_ids,\n attention_mask = attention_mask,\n inputs_embeds = inputs_embeds,\n labels = labels,\n output_attentions = output_attentions,\n output_hidden_states = output_hidden_states,\n return_dict = return_dict,\n **kwargs,\n )\n else:\n return self.base_model(\n input_ids = input_ids,\n causal_mask = causal_mask,\n attention_mask = attention_mask,\n inputs_embeds = inputs_embeds,\n labels = labels,\n output_attentions = output_attentions,\n output_hidden_states = output_hidden_states,\n return_dict = return_dict,\n num_logits_to_keep = num_logits_to_keep,\n logits_to_keep = logits_to_keep,\n **kwargs,\n )\n\n\ndef _get_rope_theta(config, default = 10000.0):\n \"\"\"Get rope_theta from config, handling both transformers 4.x and 5.x.\"\"\"\n try:\n return config.rope_theta\n except (AttributeError, KeyError):\n pass\n rp = getattr(config, \"rope_parameters\", None)\n if isinstance(rp, dict):\n return rp.get(\"rope_theta\", default)\n return default\n\n\n# Solves https://github.com/unslothai/unsloth/issues/168\n# Static KV Cache was introduced in 4.38.0, causing training to be much slower.\n# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.\n# https://github.com/huggingface/transformers/pull/27931\n# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py\nclass LlamaRotaryEmbedding(torch.nn.Module):\n # Fixes https://github.com/huggingface/transformers/pull/28837\n # https://github.com/microsoft/DeepSpeed/issues/4932\n # The precision of RoPE buffers is not correct, so we cast to int64.\n def __init__(\n self,\n dim = None,\n max_position_embeddings = 2048,\n base = 10000,\n device = None,\n config = None, # [TODO] Hack to pass in config - need to remove later\n ):\n super().__init__()\n if config is not None:\n # [TODO] Hack to pass in config - need to remove later\n base = _get_rope_theta(config, default = base)\n partial_rotary_factor = (\n config.partial_rotary_factor\n if hasattr(config, \"partial_rotary_factor\")\n else 1.0\n )\n dim = getattr(config, \"head_dim\", None)\n if dim is None:\n dim = int((config.hidden_size // config.num_attention_heads))\n device = DEVICE_TYPE_TORCH\n max_position_embeddings = config.max_position_embeddings\n\n self.dim = dim\n self.max_position_embeddings = max_position_embeddings\n self.base = base\n # Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this\n self.current_rope_size = min(4 * 8192, self.max_position_embeddings)\n self.multi_gpu_cos_cached = [None] * DEVICE_COUNT\n self.multi_gpu_sin_cached = [None] * DEVICE_COUNT\n\n # Normal Llama-3 RoPE\n inv_freq = 1.0 / (\n self.base\n ** (\n torch.arange(0, self.dim, 2, dtype = torch.int64, device = \"cpu\").float()\n / self.dim\n )\n )\n inv_freq = self._apply_inv_freq_scaling(inv_freq)\n self.register_buffer(\"inv_freq\", inv_freq, persistent = False)\n\n # Build here to make `torch.jit.trace` work.\n for device_idx in range(DEVICE_COUNT):\n self._set_cos_sin_cache(\n seq_len = self.current_rope_size,\n device = torch.device(device_idx),\n dtype = torch.get_default_dtype(),\n )\n\n # dummy so that patch_utils doesn't fail for now\n self.cos_cached = torch.empty(\n 1, device = get_current_device(), dtype = torch.get_default_dtype()\n )\n self.sin_cached = torch.empty(\n 1, device = get_current_device(), dtype = torch.get_default_dtype()\n )\n\n def _apply_inv_freq_scaling(self, inv_freq):\n \"\"\"Override to apply custom inv_freq scaling (e.g., extended RoPE).\"\"\"\n return inv_freq\n\n def _apply_time_scaling(self, t):\n \"\"\"Override to apply custom time scaling (e.g., linear scaling).\"\"\"\n return t\n\n def _set_cos_sin_cache(self, seq_len, device, dtype):\n # Note: on the original Llama codebase, these tensors are created on the target device (and not on CPU) and\n # in FP32. They are applied (multiplied) in FP32 as well.\n self.current_rope_size = seq_len\n t = torch.arange(\n self.current_rope_size, device = self.inv_freq.device, dtype = torch.int64\n ).float()\n t = self._apply_time_scaling(t)\n\n freqs = torch.outer(t, self.inv_freq)\n # Different from paper, but it uses a different permutation in order to obtain the same calculation\n emb = torch.cat((freqs, freqs), dim = -1)\n cos = emb.cos().to(dtype = dtype, device = device, non_blocking = True)\n sin = emb.sin().to(dtype = dtype, device = device, non_blocking = True)\n self.multi_gpu_cos_cached[device.index] = cos\n self.multi_gpu_sin_cached[device.index] = sin\n return cos, sin\n\n def forward(self, x, position_ids = None, seq_len = None):\n # x: [bs, num_attention_heads, seq_len, head_size]\n if seq_len is not None and seq_len > self.current_rope_size:\n self._set_cos_sin_cache(seq_len = seq_len, device = x.device, dtype = x.dtype)\n\n device_index = x.device.index\n return (\n self.multi_gpu_cos_cached[device_index][:seq_len],\n self.multi_gpu_sin_cached[device_index][:seq_len],\n )\n\n def get_cached(self, seq_len = None, device_index = None):\n if device_index is None:\n device_index = get_current_device()\n return self.multi_gpu_cos_cached[device_index], self.multi_gpu_sin_cached[\n device_index\n ]\n\n def extend_rope_embedding(self, x, seq_len):\n if seq_len <= self.current_rope_size:\n return\n # Iteratively grow by increments of 8192\n self.current_rope_size = ((seq_len // 8192) + ((seq_len % 8192) != 0)) * 8192\n for device_idx in range(DEVICE_COUNT):\n self._set_cos_sin_cache(\n self.current_rope_size, device = torch.device(device_idx), dtype = x.dtype\n )\n\n\nclass LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):\n \"\"\"LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev\"\"\"\n\n # Fixes https://github.com/huggingface/transformers/pull/28837\n # https://github.com/microsoft/DeepSpeed/issues/4932\n # The precision of RoPE buffers is not correct, so we cast to int64.\n def __init__(\n self,\n dim = None,\n max_position_embeddings = 2048,\n base = 10000,\n device = None,\n scaling_factor = 1.0,\n config = None, # [TODO] Hack to pass in config - need to remove later\n ):\n self.scaling_factor = scaling_factor\n super().__init__(\n dim = dim,\n max_position_embeddings = max_position_embeddings,\n base = base,\n device = device,\n config = config,\n )\n\n def _apply_time_scaling(self, t):\n \"\"\"Apply linear scaling to time indices.\"\"\"\n return t / self.scaling_factor\n\n\n# See https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/rotary_embedding.py#L736\n# For Llama 3.1\nclass LlamaExtendedRotaryEmbedding(LlamaRotaryEmbedding):\n def __init__(\n self,\n dim = None,\n max_position_embeddings = 2048,\n base = 10000,\n device = None,\n config = None, # [TODO] Hack to pass in config - need to remove later\n ):\n super().__init__(\n dim = dim,\n max_position_embeddings = max_position_embeddings,\n base = base,\n device = device,\n config = config,\n )\n\n # From https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py#L41\n def _apply_inv_freq_scaling(self, freqs: torch.Tensor):\n # Values obtained from grid search\n scale_factor = 8\n low_freq_factor = 1\n high_freq_factor = 4\n old_context_len = 8192 # original llama3 length\n\n low_freq_wavelen = old_context_len / low_freq_factor\n high_freq_wavelen = old_context_len / high_freq_factor\n new_freqs = []\n for freq in freqs:\n wavelen = 2 * math.pi / freq\n if wavelen < high_freq_wavelen:\n new_freqs.append(freq)\n elif wavelen > low_freq_wavelen:\n new_freqs.append(freq / scale_factor)\n else:\n assert low_freq_wavelen != high_freq_wavelen\n smooth = (old_context_len / wavelen - low_freq_factor) / (\n high_freq_factor - low_freq_factor\n )\n new_freqs.append((1 - smooth) * freq / scale_factor + smooth * freq)\n return torch.tensor(new_freqs, dtype = freqs.dtype, device = freqs.device)\n\n\nclass LongRopeRotaryEmbedding(torch.nn.Module):\n # For Phi 3.5 128K https://huggingface.co/microsoft/Phi-3.5-mini-instruct/blob/main/modeling_phi3.py\n def __init__(\n self,\n dim = None,\n max_position_embeddings = 131072,\n original_max_position_embeddings = 4096,\n base = 10000,\n short_factor = None,\n long_factor = None,\n device = None,\n config = None, # [TODO] Hack to pass in config - need to remove later\n ):\n super().__init__()\n assert short_factor is not None\n assert long_factor is not None\n assert type(original_max_position_embeddings) is int\n\n if config is not None:\n # [TODO] Hack to pass in config - need to remove later\n base = _get_rope_theta(config, default = base)\n partial_rotary_factor = (\n config.partial_rotary_factor\n if hasattr(config, \"partial_rotary_factor\")\n else 1.0\n )\n dim = int((config.hidden_size // config.num_attention_heads))\n device = DEVICE_TYPE_TORCH\n max_position_embeddings = config.max_position_embeddings\n\n self.dim = dim\n self.max_position_embeddings = max_position_embeddings\n self.original_max_position_embeddings = original_max_position_embeddings\n self.base = base\n # Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this\n self.current_rope_size = min(\n original_max_position_embeddings, self.max_position_embeddings\n )\n self.multi_gpu_short_cos_cached = [None] * DEVICE_COUNT\n self.multi_gpu_short_sin_cached = [None] * DEVICE_COUNT\n self.multi_gpu_long_cos_cached = [None] * DEVICE_COUNT\n self.multi_gpu_long_sin_cached = [None] * DEVICE_COUNT\n\n # Long RoPE similar to RoPE except short sequences have 1 cos / sin\n # and long sequences have another cos / sin\n inv_freq_shape = (\n torch.arange(0, self.dim, 2, dtype = torch.int64, device = \"cpu\").float()\n / self.dim\n )\n short_factor = torch.tensor(short_factor, device = \"cpu\", dtype = torch.float32)\n long_factor = torch.tensor(long_factor, device = \"cpu\", dtype = torch.float32)\n short_inv_freq = 1.0 / (short_factor * self.base**inv_freq_shape)\n long_inv_freq = 1.0 / (long_factor * self.base**inv_freq_shape)\n\n # Phi-3 Scale factor\n scale = self.max_position_embeddings / self.original_max_position_embeddings\n if scale <= 1.0:\n scaling_factor = 1.0\n else:\n scaling_factor = math.sqrt(\n 1 + math.log(scale) / math.log(self.original_max_position_embeddings)\n )\n self.scaling_factor = scaling_factor\n\n # Short and long inv_freq\n self.register_buffer(\"short_inv_freq\", short_inv_freq, persistent = False)\n self.register_buffer(\"long_inv_freq\", long_inv_freq, persistent = False)\n\n # Build here to make `torch.jit.trace` work.\n # Initialize short sequences cache for all devices\n dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16\n t = torch.arange(\n original_max_position_embeddings,\n device = self.short_inv_freq.device,\n dtype = torch.int64,\n ).float()\n freqs = torch.outer(t, self.short_inv_freq)\n emb = torch.cat((freqs, freqs), dim = -1)\n\n for device_idx in range(DEVICE_COUNT):\n device_obj = torch.device(device_idx)\n cos_cached = (emb.cos() * self.scaling_factor).to(\n dtype = dtype, device = device_obj, non_blocking = True\n )\n sin_cached = (emb.sin() * self.scaling_factor).to(\n dtype = dtype, device = device_obj, non_blocking = True\n )\n self.multi_gpu_short_cos_cached[device_idx] = cos_cached\n self.multi_gpu_short_sin_cached[device_idx] = sin_cached\n\n # dummy so that patch_utils doesn't fail for now\n self.short_cos_cached = torch.empty(\n 1, device = get_current_device(), dtype = torch.get_default_dtype()\n )\n self.short_sin_cached = torch.empty(\n 1, device = get_current_device(), dtype = torch.get_default_dtype()\n )\n self.long_cos_cached = torch.empty(\n 1, device = get_current_device(), dtype = torch.get_default_dtype()\n )\n self.long_sin_cached = torch.empty(\n 1, device = get_current_device(), dtype = torch.get_default_dtype()\n )\n\n def _set_cos_sin_cache(self, seq_len, device, dtype):\n # Note: on the original Llama codebase, these tensors are created on the target device (and not on CPU) and\n # in FP32. They are applied (multiplied) in FP32 as well.\n self.current_rope_size = seq_len\n\n t = torch.arange(\n self.current_rope_size, device = self.long_inv_freq.device, dtype = torch.int64\n ).float()\n # Long sequences\n freqs = torch.outer(t, self.long_inv_freq)\n emb = torch.cat((freqs, freqs), dim = -1)\n cos_cached = (emb.cos() * self.scaling_factor).to(\n dtype = dtype, device = device, non_blocking = True\n )\n sin_cached = (emb.sin() * self.scaling_factor).to(\n dtype = dtype, device = device, non_blocking = True\n )\n self.multi_gpu_long_cos_cached[device.index] = cos_cached\n self.multi_gpu_long_sin_cached[device.index] = sin_cached\n return cos_cached, sin_cached\n\n def forward(self, x, position_ids = None, seq_len = None):\n # x: [bs, num_attention_heads, seq_len, head_size]\n if seq_len is not None and seq_len > self.current_rope_size:\n self._set_cos_sin_cache(seq_len = seq_len, device = x.device, dtype = x.dtype)\n\n device_index = x.device.index\n\n if seq_len is not None and seq_len < self.original_max_position_embeddings:\n return (\n self.multi_gpu_short_cos_cached[device_index][:seq_len],\n self.multi_gpu_short_sin_cached[device_index][:seq_len],\n )\n else:\n return (\n self.multi_gpu_long_cos_cached[device_index][:seq_len],\n self.multi_gpu_long_sin_cached[device_index][:seq_len],\n )\n\n def get_cached(self, seq_len = None, device_index = None):\n if device_index is None:\n device_index = get_current_device()\n if seq_len is not None and seq_len < self.original_max_position_embeddings:\n return self.multi_gpu_short_cos_cached[\n device_index\n ], self.multi_gpu_short_sin_cached[device_index]\n return self.multi_gpu_long_cos_cached[\n device_index\n ], self.multi_gpu_long_sin_cached[device_index]\n\n def extend_rope_embedding(self, x, seq_len):\n if seq_len <= self.current_rope_size:\n return\n # Iteratively grow by increments of 8192\n self.current_rope_size = ((seq_len // 8192) + ((seq_len % 8192) != 0)) * 8192\n for device_idx in range(DEVICE_COUNT):\n self._set_cos_sin_cache(\n self.current_rope_size, device = torch.device(device_idx), dtype = x.dtype\n )\n\n\ndef unsloth_fast_generate(\n self,\n *args,\n **kwargs,\n):\n # If the model starts out in training mode, restore training mode after generation\n restore_training_mode = self.training\n\n FastLlamaModel.for_inference(self)\n\n # Unpack BatchEncoding passed as input_ids for backwards compatibility.\n # Old notebooks do model.generate(input_ids=tokenizer(...)) where the tokenizer\n # output is a BatchEncoding (dict-like). Transformers v5 generate() calls\n # .shape on it directly and crashes. Unpack into separate kwargs so both\n # v4 and v5 work transparently.\n _maybe_encoding = kwargs.get(\"input_ids\", None)\n if (\n _maybe_encoding is not None\n and not isinstance(_maybe_encoding, torch.Tensor)\n and hasattr(_maybe_encoding, \"items\")\n ):\n batch_data = kwargs.pop(\"input_ids\")\n for key, val in batch_data.items():\n kwargs.setdefault(key, val)\n\n dtype = _get_dtype(dtype_from_config(self.config))\n\n if hasattr(self, \"config\") and hasattr(self.config, \"max_position_embeddings\"):\n if (\n \"input_ids\" in kwargs\n and kwargs[\"input_ids\"] is not None\n and \"max_new_tokens\" in kwargs\n ):\n _ids = kwargs[\"input_ids\"]\n if hasattr(_ids, \"shape\") and (\n _ids.shape[-1] + kwargs[\"max_new_tokens\"]\n > self.config.max_position_embeddings\n ):\n raise ValueError(\n f\"Unsloth: input length {_ids.shape[-1]} + max_new_tokens {kwargs['max_new_tokens']} exceeds the maximum sequence length of {self.config.max_position_embeddings}!\\n\"\n \"You will need to do long context extension by increasing the `max_seq_length` in `FastLanguageModel.from_pretrained`.\"\n )\n\n # Must patch accelerate for Xformers\n # if accelerate_new_send_to_device is not None:\n # import accelerate.utils.operations\n # accelerate.utils.operations.send_to_device = accelerate_new_send_to_device\n # pass\n\n # For newer HF\n kwargs[\"cache_implementation\"] = \"dynamic\"\n # For num_logits_to_keep\n num_logits_to_keep = kwargs.get(\"num_logits_to_keep\", None)\n logits_to_keep = kwargs.get(\"logits_to_keep\", None)\n if num_logits_to_keep is None and logits_to_keep is None:\n kwargs[\"num_logits_to_keep\"] = 1\n\n # Remove token_type_ids\n kwargs.pop(\"token_type_ids\", None)\n\n # Check pad_token\n model_eos_token_id = getattr(self.config, \"eos_token_id\", None)\n if model_eos_token_id is not None and hasattr(model_eos_token_id, \"__iter__\"):\n model_eos_token_id = model_eos_token_id[0]\n\n kwargs[\"pad_token_id\"] = kwargs.pop(\"pad_token_id\", model_eos_token_id)\n\n # Mixed precision autocast\n with (\n _get_inference_mode_context_manager(self),\n torch.autocast(device_type = DEVICE_TYPE_TORCH, dtype = dtype),\n ):\n output = self._old_generate(*args, **kwargs)\n\n # Return accelerate back\n # if accelerate_new_send_to_device is not None:\n # accelerate.utils.operations.send_to_device = accelerate_old_send_to_device\n # pass\n\n if restore_training_mode:\n FastLlamaModel.for_training(self)\n\n return output\n\n\nclass FastLlamaModel:\n @staticmethod\n def _prepare_for_qat(model, qat_scheme):\n model = _prepare_model_for_qat(model, qat_scheme)\n return model\n\n @staticmethod\n def pre_patch():\n init_name, function = patch_llama_rope_scaling(\n model_name = \"llama\",\n rope_module = LlamaRotaryEmbedding,\n scaled_rope_module = LlamaLinearScalingRotaryEmbedding,\n extended_rope_module = LlamaExtendedRotaryEmbedding,\n attention_module = LlamaAttention,\n longrope_module = LongRopeRotaryEmbedding,\n )\n if init_name is not None:\n exec(function, globals())\n LlamaAttention.__init__ = eval(init_name)\n LlamaAttention.forward = LlamaAttention_fast_forward\n LlamaSdpaAttention.forward = LlamaAttention_fast_forward\n LlamaFlashAttention2.forward = LlamaAttention_fast_forward\n LlamaDecoderLayer.forward = LlamaDecoderLayer_fast_forward\n LlamaModel.forward = LlamaModel_fast_forward\n LlamaForCausalLM.forward = CausalLM_fast_forward(\n LlamaModel_fast_forward_inference\n )\n PeftModelForCausalLM.forward = PeftModel_fast_forward\n fix_prepare_inputs_for_generation(LlamaForCausalLM)\n\n # Solves https://github.com/unslothai/unsloth/issues/168\n # Static KV Cache was introduced in 4.38.0, causing training to be much slower.\n # Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.\n # https://github.com/huggingface/transformers/pull/27931\n # https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py\n import transformers.models.llama.modeling_llama\n\n transformers.models.llama.modeling_llama.LlamaRotaryEmbedding = (\n LlamaRotaryEmbedding\n )\n transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding = (\n LlamaLinearScalingRotaryEmbedding\n )\n return\n\n @staticmethod\n def from_pretrained(\n model_name = \"unsloth/llama-3-8b-bnb-4bit\",\n max_seq_length = None,\n dtype = None,\n load_in_4bit = True,\n token = None,\n device_map = \"sequential\",\n rope_scaling = None,\n fix_tokenizer = True,\n model_patcher = None,\n tokenizer_name = None,\n trust_remote_code = False,\n revision = None,\n fast_inference = False, # uses vLLM\n gpu_memory_utilization = 0.5,\n float8_kv_cache = False,\n random_state = 3407,\n max_lora_rank = 16,\n disable_log_stats = False,\n unsloth_vllm_standby = False,\n num_labels = None,\n qat_scheme = None,\n load_in_fp8 = False, # fp8 LoRA (True, False, 'block')\n **kwargs,\n ):\n os.environ[\"UNSLOTH_USE_NEW_MODEL\"] = \"0\"\n if trust_remote_code:\n if fast_inference:\n raise NotImplementedError(\n \"Unsloth: Fast inference does not support `trust_remote_code` yet.\"\n )\n print(\n \"Unsloth: WARNING `trust_remote_code` is True.\\n\"\n \"Are you certain you want to do remote code execution?\"\n )\n if fast_inference:\n if not is_vLLM_available():\n print(\"Unsloth: vLLM is not installed! Will use Unsloth inference!\")\n fast_inference = False\n if DEVICE_TYPE == \"cuda\":\n major_version, minor_version = torch.cuda.get_device_capability()\n if major_version < 7:\n print(\n \"Unsloth: vLLM does not work on older GPUs - will switch to Unsloth inference!\"\n )\n fast_inference = False\n elif DEVICE_TYPE == \"hip\":\n fast_inference = True\n if (\n unsloth_vllm_standby\n and os.environ.get(\"UNSLOTH_VLLM_STANDBY\", \"0\") == \"0\"\n ):\n raise RuntimeError(\n \"Unsloth: `unsloth_vllm_standby` is True, but environment variable `UNSLOTH_VLLM_STANDBY` is not set to 1!\"\n )\n\n token = hf_login(token)\n if model_patcher is None:\n model_patcher = FastLlamaModel\n SUPPORTS_BFLOAT16 = is_bfloat16_supported()\n\n if DEVICE_TYPE == \"cuda\":\n gpu_stats = torch.cuda.get_device_properties(0)\n gpu_stats_name = (\n gpu_stats.name + \". \" if gpu_stats.name != \"\" else \"NVIDIA GPU Device. \"\n )\n gpu_version = torch.version.cuda\n gpu_stats_snippet = f\"CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {gpu_version}.\"\n try:\n vllm_version = f\" vLLM: {importlib_version('vllm')}.\"\n except:\n vllm_version = \"\"\n elif DEVICE_TYPE == \"hip\":\n gpu_stats = torch.cuda.get_device_properties(0)\n gpu_stats_name = resolve_hip_gpu_stats_name(gpu_stats)\n gpu_version = torch.version.hip\n gpu_stats_snippet = f\"ROCm Toolkit: {gpu_version}.\"\n try:\n vllm_version = f\" vLLM: {importlib_version('vllm')}.\"\n except:\n vllm_version = \"\"\n elif DEVICE_TYPE == \"xpu\":\n gpu_stats = torch.xpu.get_device_properties(0)\n gpu_stats_name = (\n gpu_stats.name + \". \" if gpu_stats.name != \"\" else \"Intel XPU Device. \"\n )\n gpu_version = torch.version.xpu\n gpu_stats_snippet = f\"Intel Toolkit: {gpu_version}.\"\n try:\n vllm_version = f\" vLLM: {importlib_version('vllm')}.\"\n except:\n vllm_version = \"\"\n else:\n raise ValueError(f\"Unsloth: Unsupported device type: {DEVICE_TYPE}\")\n\n max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)\n\n statistics = (\n f\"==((====))== Unsloth {__version__}: Fast {model_patcher.__name__[4:-5]} patching. Transformers: {transformers_version}.{vllm_version}\\n\"\n f\" {chr(92)}{chr(92)} /| {gpu_stats_name}Num GPUs = {DEVICE_COUNT}. Max memory: {max_memory} GB. Platform: {platform_system}.\\n\"\n f\"O^O/ {chr(92)}_/ {chr(92)} Torch: {torch.__version__}. {gpu_stats_snippet} Triton: {triton_version}\\n\"\n f\"{chr(92)} / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\\n\"\n f' \"-____-\" Free license: http://github.com/unslothai/unsloth'\n )\n\n print(statistics)\n\n # Warn about fast transfers\n if \"HF_HUB_ENABLE_HF_TRANSFER\" in os.environ:\n old_hf_transfer = os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"]\n if old_hf_transfer in (\"False\", \"false\"):\n old_hf_transfer = \"0\"\n if old_hf_transfer in (\"True\", \"true\"):\n old_hf_transfer = \"1\"\n else:\n old_hf_transfer = \"0\"\n if old_hf_transfer == \"1\":\n print(\n \"Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\"\n )\n if old_hf_transfer != \"0\":\n os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = \"1\"\n\n model_patcher.pre_patch()\n # For debugging - we use a download counter to see if environments are not breaking or if HF is down\n get_statistics(kwargs.get(\"local_files_only\", False))\n\n if dtype is None:\n dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16\n elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:\n logger.warning_once(\n \"Device does not support bfloat16. Will change to float16.\"\n )\n dtype = torch.float16\n # elif dtype == torch.float16 and SUPPORTS_BFLOAT16:\n # logger.warning_once(\"Device supports bfloat16 but you selected float16. Will change to bfloat16.\")\n # dtype = torch.bfloat16\n\n assert (\n dtype == torch.float16 or dtype == torch.bfloat16 or dtype == torch.float32\n )\n\n # RoPE Scaling\n model_config = AutoConfig.from_pretrained(\n model_name,\n token = token,\n attn_implementation = \"sdpa\",\n )\n model_config.model_name = model_name\n model_max_seq_length = model_config.max_position_embeddings\n\n verify_fp8_support_if_applicable(model_config)\n\n # Check if RoPE Scaling is even allowed\n model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__]\n IS_FALCON_H1 = model_config.model_type.startswith(\"falcon_h1\")\n\n preferred_attn_impl = determine_attention_implementation(\n model_function, model_config\n )\n\n has_rope_scaling = False\n try:\n with open(inspect.getfile(model_function), \"r\", encoding = \"utf-8\") as file:\n has_rope_scaling = \"self.config.rope_scaling\" in file.read()\n except:\n pass\n has_rope_scaling = True\n\n # If max_seq_length is not specified, use maximum from config\n if max_seq_length is None:\n max_seq_length = model_max_seq_length\n\n if (rope_scaling is None) and (max_seq_length > model_max_seq_length):\n rope_scaling = max_seq_length / model_max_seq_length\n\n if fast_inference:\n raise NotImplementedError(\n \"Unsloth: Fast inference does not yet work with RoPE Scaling.\"\n )\n\n logger.warning_once(\n f\"Unsloth: {model_name} can only handle sequence lengths of at most \"\n f\"{model_max_seq_length}.\\nBut with kaiokendev's RoPE scaling of \"\n f\"{round(rope_scaling, 3)}, it can be magically be extended to \"\n f\"{max_seq_length}!\"\n )\n\n # Warn RoPE scaling isn't allowed\n if not has_rope_scaling:\n raise RuntimeError(\n f\"However, {model_name} doesn't support RoPE Scaling!\\n\"\n \"Please file a feature request at https://github.com/unslothai/unsloth.\"\n )\n\n rope_scaling = {\n \"type\": \"linear\",\n \"factor\": rope_scaling,\n }\n\n # Add to kwargs\n kwargs[\"rope_scaling\"] = rope_scaling\n\n bnb_config = None\n if load_in_4bit:\n llm_int8_skip_modules = SKIP_QUANTIZATION_MODULES.copy()\n if IS_FALCON_H1:\n # we cannot quantize out_proj layer due to mamba kernels: https://github.com/tiiuae/Falcon-H1/issues/13#issuecomment-2918671274\n llm_int8_skip_modules.append(\"out_proj\")\n bnb_config = BitsAndBytesConfig(\n load_in_4bit = True,\n bnb_4bit_use_double_quant = True,\n bnb_4bit_quant_type = \"nf4\",\n bnb_4bit_compute_dtype = dtype,\n llm_int8_skip_modules = llm_int8_skip_modules,\n )\n\n # https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/discussions/12\n # RoPE Scaling's max_position_embeddings must be updated\n max_position_embeddings = max(max_seq_length, model_max_seq_length)\n kwargs.pop(\"attn_implementation\", None) # No need since we auto call it\n\n # Cannot be None, since HF now checks for the config\n if load_in_4bit:\n kwargs[\"quantization_config\"] = bnb_config\n\n kwargs = add_dtype_kwargs(dtype, kwargs)\n\n raise_handler = RaiseUninitialized()\n if num_labels is not None:\n # Transformers 5.x @strict config classes reject unexpected kwargs\n # like num_labels and max_position_embeddings. Set on the config\n # object directly and pass config= instead.\n model_config.num_labels = num_labels\n if max_position_embeddings is not None:\n model_config.max_position_embeddings = max_position_embeddings\n # Pop config-level attrs that would be rejected by @strict model init\n for _cfg_key in (\"id2label\", \"label2id\", \"rope_scaling\"):\n _cfg_val = kwargs.pop(_cfg_key, None)\n if _cfg_val is not None:\n setattr(model_config, _cfg_key, _cfg_val)\n model = AutoModelForSequenceClassification.from_pretrained(\n model_name,\n config = model_config,\n device_map = device_map,\n # torch_dtype = dtype, # transformers changed torch_dtype to dtype\n # quantization_config = bnb_config,\n token = token,\n trust_remote_code = trust_remote_code,\n attn_implementation = preferred_attn_impl,\n **kwargs,\n )\n elif not fast_inference:\n model = AutoModelForCausalLM.from_pretrained(\n model_name,\n device_map = device_map,\n # torch_dtype = dtype, # transformers changed torch_dtype to dtype\n # quantization_config = bnb_config,\n token = token,\n max_position_embeddings = max_position_embeddings,\n trust_remote_code = trust_remote_code,\n attn_implementation = preferred_attn_impl,\n **kwargs,\n )\n model.fast_generate = make_fast_generate_wrapper(model.generate)\n model.fast_generate_batches = None\n else:\n from unsloth_zoo.vllm_utils import (\n load_vllm,\n get_vllm_state_dict,\n convert_vllm_to_huggingface,\n generate_batches,\n )\n\n fp8_mode = None\n if load_in_fp8 != False:\n fp8_mode = _get_fp8_mode_and_check_settings(\n load_in_fp8,\n fast_inference,\n )\n\n allowed_args = inspect.getfullargspec(load_vllm).args\n load_vllm_kwargs = dict(\n model_name = model_name,\n config = model_config,\n gpu_memory_utilization = gpu_memory_utilization,\n max_seq_length = max_seq_length,\n dtype = dtype,\n float8_kv_cache = float8_kv_cache,\n enable_lora = True,\n max_lora_rank = max_lora_rank,\n disable_log_stats = disable_log_stats,\n use_bitsandbytes = load_in_4bit,\n unsloth_vllm_standby = unsloth_vllm_standby,\n fp8_mode = fp8_mode,\n )\n for allowed_arg in allowed_args:\n if allowed_arg not in load_vllm_kwargs and allowed_arg in kwargs:\n load_vllm_kwargs[allowed_arg] = kwargs[allowed_arg]\n pass\n\n # Load vLLM first\n llm = load_vllm(**load_vllm_kwargs)\n\n # Convert to HF format\n _, quant_state_dict = get_vllm_state_dict(\n llm,\n config = model_config,\n load_in_fp8 = load_in_fp8,\n )\n model = convert_vllm_to_huggingface(\n quant_state_dict, model_config, dtype, bnb_config\n )\n model.vllm_engine = llm\n model.fast_generate = model.vllm_engine.generate\n model.fast_generate_batches = functools.partial(\n generate_batches, model.vllm_engine\n )\n raise_handler.remove()\n # Return old flag\n os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = old_hf_transfer\n\n # Counteract saved tokenizers\n tokenizer_name = model_name if tokenizer_name is None else tokenizer_name\n tokenizer = load_correct_tokenizer(\n tokenizer_name = tokenizer_name,\n model_max_length = max_position_embeddings,\n padding_side = \"right\",\n token = token,\n trust_remote_code = trust_remote_code,\n fix_tokenizer = fix_tokenizer,\n )\n\n model, tokenizer = patch_tokenizer(model, tokenizer)\n model, tokenizer = model_patcher.post_patch(\n model, tokenizer, correct_dtype = dtype\n )\n\n # Patch up QKV / O and MLP\n for idx, layer in enumerate(model.model.layers):\n layer.self_attn.apply_qkv = original_apply_qkv\n layer.self_attn.apply_o = original_apply_o\n\n # Patch Trainer\n from transformers.trainer import Trainer\n\n try:\n if Trainer._inner_training_loop.__name__ != \"_fast_inner_training_loop\":\n inner_training_loop = inspect.getsource(Trainer._inner_training_loop)\n Trainer._original_training_loop = inner_training_loop\n else:\n inner_training_loop = Trainer._original_training_loop\n except:\n raise RuntimeError(\"Unsloth: Unsuccessfully patched inner_training_loop\")\n\n import transformers.trainer\n\n items_in_trainer = dir(transformers.trainer)\n good_items = []\n for item in items_in_trainer:\n if item in inner_training_loop:\n good_items.append(item)\n exec(\n \"from transformers.trainer import (\"\n + \", \".join(x for x in good_items)\n + \")\",\n globals(),\n )\n\n start = re.search(\n r\"logger\\.info\\([\\\"\\'].+?Running training\", inner_training_loop\n ).span(0)[0]\n end = inner_training_loop.find(\"\\n\\n\", start)\n original_debug = inner_training_loop[start:end]\n spaces = re.search(r\"\\n([\\s\\t]{1,})\", original_debug).group(0)[1:]\n front_spaces = re.match(r\"([\\s\\t]{1,})\", inner_training_loop).group(0)\n\n # Cannot use \\\\ since it will cause a SyntaxWarning in Python 3.12\n # Instead use chr(92) == \\\\\n debug_info = \"\"\"debug_info = \\\\\n f\"==((====))== Unsloth - 2x faster free finetuning | Num GPUs used = {len(set(p.device for p in model.parameters()))}\\\\n\"\\\\\n f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\n f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\n f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size:,}\\\\n\"\\\\\n f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_param_count(model)*100:.2f}% trained)'\n logger.warning(debug_info)\n import gc\n for _ in range(3):\n gc.collect()\n if DEVICE_TYPE == \"xpu\":\n torch.xpu.empty_cache()\n else:\n torch.cuda.empty_cache()\"\"\"\n\n debug_info = debug_info.split(\"\\n\")\n debug_info = \"\\n\".join(\n [debug_info[0]] + [spaces + x[8:] for x in debug_info[1:]]\n )\n inner_training_loop = inner_training_loop.replace(original_debug, debug_info)\n\n debug_info = \"\"\"n_total_devices = total_train_batch_size // \\\\\n args.gradient_accumulation_steps // self._train_batch_size\n if n_total_devices > 1:\n logger.warning_once('Unsloth is running with multi GPUs - the effective batch size is multiplied by ' + str(n_total_devices))\n debug_info =\"\"\"\n debug_info = debug_info.split(\"\\n\")\n debug_info = \"\\n\".join(\n [debug_info[0]] + [spaces + x[8:] for x in debug_info[1:]]\n )\n inner_training_loop = inner_training_loop.replace(\"debug_info =\", debug_info, 1)\n\n front_spaces = re.match(r\"[\\t\\s]{1,}\", inner_training_loop).group(0)\n inner_training_loop = re.sub(\n r\"^\" + front_spaces, \"\", inner_training_loop, flags = re.MULTILINE\n )\n inner_training_loop = inner_training_loop.replace(\n \"train_dataloader = tpu_spmd_dataloader(train_dataloader)\",\n \"raise RuntimeError('Unsloth: TPUs are not yet supported!')\",\n )\n inner_training_loop = inner_training_loop.replace(\n \"_inner_training_loop\",\n \"_fast_inner_training_loop\",\n 1,\n )\n inner_training_loop = inner_training_loop.replace(\n \"is_torch_tpu_available()\",\n \"False\",\n )\n exec(inner_training_loop, globals())\n Trainer._inner_training_loop = _fast_inner_training_loop\n\n # Save max_seq_length\n model.max_seq_length = max_seq_length\n m = model\n while hasattr(m, \"model\"):\n m.max_seq_length = max_seq_length\n m = m.model\n m.max_seq_length = max_seq_length\n # Save to modules as well\n for module in model.modules():\n module.max_seq_length = max_seq_length\n\n # We check the tokenizer first for errors\n if fix_tokenizer:\n tokenizer = check_tokenizer(\n model = model,\n tokenizer = tokenizer,\n model_name = model_name,\n model_max_length = max_position_embeddings,\n padding_side = \"right\",\n token = token,\n )\n patch_saving_functions(tokenizer)\n\n # Fix up config for transformers uploading PEFT\n # Not necessary anymore since we require transformers>=4.37!\n if False:\n name = model.config._name_or_path\n if name.startswith(\"unsloth/\") and name.endswith(\"-bnb-4bit\"):\n name = name[: len(name) - len(\"-bnb-4bit\")]\n model.config.update({\"_name_or_path\": name})\n\n # Log Unsloth version for future fastpaths for inference\n model.config.update({\"unsloth_version\": __version__})\n\n # Add save modules\n patch_saving_functions(model)\n Trainer._inner_training_loop = _fast_inner_training_loop\n\n # Fix gradient accumulation\n patch_gradient_accumulation_fix(Trainer)\n\n # Save tokenizer for inference purposes\n tokenizer.padding_side = \"left\" # Force inference\n internal_model = model\n while hasattr(internal_model, \"model\"):\n internal_model._saved_temp_tokenizer = tokenizer\n # Also set is_loaded_in_8bit to disable incorrect DDP\n internal_model.is_loaded_in_8bit = True\n\n internal_model = internal_model.model\n internal_model._saved_temp_tokenizer = tokenizer\n # Also set is_loaded_in_8bit to disable incorrect DDP\n internal_model.is_loaded_in_8bit = True\n\n # For transformers > 4.47.1, we need to add rotary_emb to all attention layers\n if IS_ATTENTION_REFACTOR or hasattr(model.model, \"rotary_emb\"):\n rotary_emb = model.model.rotary_emb\n for layer in model.model.layers:\n layer.self_attn.rotary_emb = rotary_emb\n\n # Add for_inference and for_training\n model.for_training = functools.partial(FastLlamaModel.for_training, model)\n model.for_inference = functools.partial(FastLlamaModel.for_inference, model)\n m = model\n while hasattr(m, \"model\"):\n m.for_training = functools.partial(FastBaseModel.for_training, m)\n m.for_inference = functools.partial(FastBaseModel.for_inference, m)\n m = m.model\n\n # Patch generate\n is_classification = \"Classification\" in str(type(model))\n if not is_classification and model.generate.__name__ != \"unsloth_fast_generate\":\n model._old_generate = model.generate\n unsloth_fast_generate.__doc__ = model._old_generate.__doc__\n model.generate = types.MethodType(unsloth_fast_generate, model)\n # Set weight[padding_idx] = 0 for embeddings that are NOT tied with the\n # lm_head. When weights are tied, zeroing the padding row also zeros\n # the corresponding lm_head row, forcing logit = 0 for the pad token.\n # This is higher than the (negative) logits for real tokens in models\n # like Gemma, causing the decoder to emit <pad> and produce gibberish.\n # Skip entirely if eos_token == pad_token to avoid zeroing EOS embedding.\n eos_token_id = (\n getattr(tokenizer, \"eos_token_id\", None) if tokenizer is not None else None\n )\n pad_token_id = (\n getattr(tokenizer, \"pad_token_id\", None) if tokenizer is not None else None\n )\n if tokenizer is not None and eos_token_id != pad_token_id:\n lm_head = getattr(model, \"lm_head\", None)\n lm_head_weight = (\n getattr(lm_head, \"weight\", None) if lm_head is not None else None\n )\n with torch.no_grad():\n for name, module in model.named_modules():\n if type(module) is torch.nn.Embedding:\n if (\n getattr(module, \"weight\", None) is not None\n and getattr(module, \"padding_idx\", None) is not None\n ):\n if module.padding_idx < module.weight.shape[0]:\n # Skip if tied to lm_head\n if (\n lm_head_weight is not None\n and module.weight.data_ptr()\n == lm_head_weight.data_ptr()\n ):\n continue\n module.weight[module.padding_idx] = 0\n return model, tokenizer\n\n @staticmethod\n def post_patch(model, tokenizer, correct_dtype = None):\n model, tokenizer = patch_model_and_tokenizer(\n model, tokenizer, downcast_rope = True, correct_dtype = correct_dtype\n )\n return model, tokenizer\n\n @staticmethod\n def get_peft_model(\n model,\n r = 16,\n target_modules = [\n \"q_proj\",\n \"k_proj\",\n \"v_proj\",\n \"o_proj\",\n \"gate_proj\",\n \"up_proj\",\n \"down_proj\",\n ],\n lora_alpha = 16,\n lora_dropout = 0.0,\n bias = \"none\",\n layers_to_transform = None,\n layers_pattern = None,\n use_gradient_checkpointing = \"unsloth\",\n random_state = 3407,\n max_seq_length = 2048, # not used anymore\n use_rslora = False,\n modules_to_save = None,\n init_lora_weights = True,\n loftq_config = {},\n temporary_location = \"_unsloth_temporary_saved_buffers\",\n qat_scheme = None,\n target_parameters = None, # For MoE expert layers (nn.Parameter)\n ensure_weight_tying = False,\n **kwargs,\n ):\n if os.environ.get(\"UNSLOTH_USE_NEW_MODEL\", \"0\") == \"1\":\n # Check for other PEFT args in kwargs\n for peft_arg, flag in (\n (\"finetune_vision_layers\", False),\n (\"finetune_language_layers\", True),\n (\"finetune_attention_modules\", True),\n (\"finetune_mlp_modules\", True),\n ):\n if peft_arg not in kwargs:\n kwargs[peft_arg] = flag\n return FastBaseModel.get_peft_model(\n model = model,\n r = r,\n target_modules = target_modules,\n lora_alpha = lora_alpha,\n lora_dropout = lora_dropout,\n bias = bias,\n layers_to_transform = layers_to_transform,\n layers_pattern = layers_pattern,\n use_gradient_checkpointing = use_gradient_checkpointing,\n random_state = random_state,\n max_seq_length = max_seq_length,\n use_rslora = use_rslora,\n modules_to_save = modules_to_save,\n init_lora_weights = init_lora_weights,\n loftq_config = loftq_config,\n temporary_location = temporary_location,\n target_parameters = target_parameters,\n ensure_weight_tying = ensure_weight_tying,\n **kwargs,\n )\n if os.environ.get(\"UNSLOTH_ENABLE_FULL_FINETUNING\", \"0\") == \"1\":\n print(\n \"Unsloth: Full finetuning is enabled, so .get_peft_model has no effect\"\n )\n return model\n transformers_set_seed(random_state)\n\n # Apply gradient checkpointing with smart heuristics\n max_seq = getattr(model, \"max_seq_length\", 512)\n dtype = model.get_input_embeddings().weight.dtype\n use_gradient_checkpointing = apply_unsloth_gradient_checkpointing(\n use_gradient_checkpointing, max_seq, dtype\n )\n\n if type(r) is not int:\n raise TypeError(f\"Unsloth: Rank of {str(r)} must be an integer.\")\n if r <= 0:\n raise TypeError(f\"Unsloth: Rank of {str(r)} must be larger than 0.\")\n\n if isinstance(model, PeftModelForCausalLM) or isinstance(\n model, PeftModelForSequenceClassification\n ):\n # Check if exactly the same and then pass through!\n assert hasattr(model, \"peft_config\")\n\n peft_config = model.peft_config[\"default\"].to_dict()\n check_parameters = [\n \"r\",\n \"lora_alpha\",\n \"lora_dropout\",\n \"bias\",\n \"layers_to_transform\",\n \"layers_pattern\",\n \"use_rslora\",\n \"init_lora_weights\",\n ]\n check_all = True\n for param in check_parameters:\n check_all = check_all and (peft_config[param] == eval(param))\n\n # Check save_modules\n old_target_modules = list(peft_config[\"target_modules\"])\n modules_to_save = peft_config[\"modules_to_save\"]\n if modules_to_save is None:\n modules_to_save = {}\n modules_to_save = list(modules_to_save)\n old_target_modules += modules_to_save\n\n # Combine all\n new_target_modules = list(target_modules) + list(\n modules_to_save if modules_to_save is not None else []\n )\n\n # Now check!\n new_target_modules = set(new_target_modules)\n check_all = check_all and (\n len(set(old_target_modules) ^ new_target_modules) == 0\n )\n\n check_all = check_all and (\n (loftq_config == {} or loftq_config is None)\n and (\n peft_config[\"loftq_config\"] == {}\n or peft_config[\"loftq_config\"] is None\n )\n )\n\n if check_all:\n # Simply pass through!\n logger.warning(\n \"Unsloth: Already have LoRA adapters! We shall skip this step.\"\n )\n\n # Offload!\n # [TODO] First offload lm_head and embed_tokens to CPU (should be disk!!)\n if \"embed_tokens\" in new_target_modules:\n print(\n \"Unsloth: Training embed_tokens in mixed precision to save VRAM\"\n )\n\n _offload_frozen_module_for_training(\n model.get_input_embeddings(), DEVICE_TYPE_TORCH\n )\n\n if \"lm_head\" in new_target_modules:\n print(\"Unsloth: Training lm_head in mixed precision to save VRAM\")\n\n _offload_frozen_module_for_training(\n model.get_output_embeddings(), DEVICE_TYPE_TORCH\n )\n\n return model\n else:\n raise TypeError(\n \"Unsloth: Your model already has LoRA adapters. Your new parameters are different.\"\n )\n\n if loftq_config is None:\n loftq_config = {}\n\n signature = str(inspect.signature(LoraConfig))\n SUPPORTS_LOFTQ = \"loftq_config\" in signature\n SUPPORTS_RSLORA = \"use_rslora\" in signature\n\n if lora_dropout != 0:\n logger.warning_once(\n f\"Unsloth: Dropout = 0 is supported for fast patching. You are using dropout = {lora_dropout}.\\n\"\n f\"Unsloth will patch all other layers, except LoRA matrices, causing a performance hit.\"\n )\n\n if bias != \"none\":\n logger.warning_once(\n f\"Unsloth: bias = `none` is supported for fast patching. You are using bias = {bias}.\\n\"\n f\"Unsloth will patch all other layers, except LoRA matrices, causing a performance hit.\"\n )\n\n if not (\n type(init_lora_weights) is bool\n or init_lora_weights == \"gaussian\"\n or init_lora_weights == \"loftq\"\n or init_lora_weights == \"corda\"\n ):\n raise ValueError(\n 'Unsloth: `init_lora_weights` must be either [True, False, \"gaussian\", \"loftq\", \"corda\"].'\n )\n\n if init_lora_weights == \"loftq\":\n if not SUPPORTS_LOFTQ:\n import peft\n\n raise RuntimeError(\n f\"Unsloth: Your PEFT version of {peft.__version__} does not support LoftQ init.\\n\"\n \"Please install PEFT 0.7.2 or higher.\\n\"\n \"You can also install from source: `pip install git+https://github.com/huggingface/peft.git\"\n )\n\n if loftq_config == {}:\n from peft import LoftQConfig\n\n logger.warning_once(\n \"Unsloth: init_lora_weights = `loftq` is set, but `loftq_config` is None.\\n\"\n \"We shall use `loftq_config = LoftQConfig(loftq_bits = 4, loftq_iter = 1)`.\"\n )\n loftq_config = LoftQConfig(loftq_bits = 4, loftq_iter = 1)\n\n if hasattr(model.config, \"quantization_config\"):\n raise ValueError(\n \"Unsloth: You are using `loftq` init, yet `load_in_4bit = True` was set.\\n\"\n \"Reload your model without any quantization by setting `load_in_4bit = False`.\"\n )\n\n assert type(use_rslora) is bool\n if use_rslora:\n if not SUPPORTS_RSLORA:\n # We manually check for PEFT\n import peft\n\n raise RuntimeError(\n f\"Unsloth: Your PEFT version of {peft.__version__} does not support `use_rslora`.\\n\"\n \"Please install PEFT 0.7.2 or higher.\\n\"\n \"You can also install from source: `pip install git+https://github.com/huggingface/peft.git\"\n )\n\n accepted_modules = frozenset(\n (\n \"lm_head\",\n \"q_proj\",\n \"k_proj\",\n \"v_proj\",\n \"o_proj\",\n \"gate_proj\",\n \"up_proj\",\n \"down_proj\",\n ),\n )\n model.config.update({\"unsloth_version\": __version__})\n\n if type(modules_to_save) is tuple:\n modules_to_save = list(modules_to_save)\n\n train_lm_head = False\n train_embed_tokens = False\n final_modules = []\n for module in target_modules:\n if module == \"embed_tokens\":\n # logger.warning_once(\n # \"Unsloth: `embed_tokens` should be placed in `modules_to_save` and not `target_modules`. \"\\\n # \"Luckily, we shall do it for you!\"\n # )\n train_embed_tokens = True\n if modules_to_save is None:\n modules_to_save = [\"embed_tokens\"]\n else:\n modules_to_save.append(\"embed_tokens\")\n\n else:\n try:\n assert module in accepted_modules\n final_modules.append(module)\n except AssertionError as e:\n final_modules.append(module)\n print(\n \"Unsloth: You added custom modules, but Unsloth hasn't optimized for this.\\n\"\n \"Beware - your finetuning might be noticeably slower!\"\n )\n pass\n\n # Check if we added new tokens!\n if hasattr(model, \"_need_to_train_embeddings\"):\n # Check if embed_tokens/lm_head are already being trained\n # (either as LoRA targets in final_modules or via modules_to_save)\n _embed_already_trained = (\n train_embed_tokens or \"embed_tokens\" in final_modules\n )\n _lm_head_already_trained = train_lm_head or \"lm_head\" in final_modules\n if not _lm_head_already_trained or not _embed_already_trained:\n print(\n \"Unsloth: You added new tokens but did not specify if you wanted to \"\n \"train the lm_head and embed_tokens.\\nWe must turn it on for you.\"\n )\n\n # Only add to modules_to_save if not already a LoRA target\n if not _embed_already_trained:\n train_embed_tokens = True\n if modules_to_save is None:\n modules_to_save = [\"embed_tokens\"]\n elif \"embed_tokens\" not in modules_to_save:\n modules_to_save.append(\"embed_tokens\")\n\n if not _lm_head_already_trained:\n train_lm_head = True\n if modules_to_save is None:\n modules_to_save = [\"lm_head\"]\n elif \"lm_head\" not in modules_to_save:\n modules_to_save.append(\"lm_head\")\n\n # Check for Llama-3\n # if hasattr(model._saved_temp_tokenizer, \"_using_llama3_template\"):\n # if not train_embed_tokens and not train_lm_head:\n # raise RuntimeError(\"\")\n\n # First fix untrained tokens\n # Wrong - can cause reserved tokens to pop out!!\n # if train_embed_tokens or train_lm_head:\n # fix_untrained_tokens(model, eps = 1e-16)\n # pass\n\n # Check modules_to_save\n if modules_to_save is not None:\n for module in modules_to_save:\n if module == \"lm_head\":\n train_lm_head = True\n elif module == \"embed_tokens\":\n train_embed_tokens = True\n else:\n raise TypeError(\n f\"Unsloth: Module = {module} is not allowed. Only 'lm_head' and 'embed_tokens' is allowed.\"\n )\n if isinstance(modules_to_save, (tuple, list)):\n modules_to_save = list(set(modules_to_save))\n\n vllm_engine = None\n if hasattr(model, \"vllm_engine\"):\n # Fast inference!\n vllm_engine = model.vllm_engine\n vllm_fast_generate = model.fast_generate\n vllm_fast_generate_batches = model.fast_generate_batches\n\n if modules_to_save is not None:\n raise NotImplementedError(\n \"Unsloth: Currently fast inference does not work with training embeddings or lm_head.\"\n )\n\n if bias != \"none\":\n raise NotImplementedError(\n \"Unsloth: Currently fast inference does not work with using biases for LoRA.\"\n )\n\n # Does not get lora yet, so get name from model, not base model\n is_classification = \"Classification\" in str(type(model))\n\n # Auto-detect MoE models and populate target_parameters for expert layers\n if target_parameters is None:\n target_parameters = get_moe_target_parameters(model, target_modules)\n\n arguments = dict(\n r = r,\n lora_alpha = lora_alpha,\n target_modules = final_modules,\n lora_dropout = lora_dropout,\n bias = bias,\n task_type = TaskType.CAUSAL_LM if not is_classification else TaskType.SEQ_CLS,\n layers_to_transform = layers_to_transform,\n init_lora_weights = init_lora_weights,\n loftq_config = loftq_config,\n use_rslora = use_rslora,\n modules_to_save = modules_to_save,\n target_parameters = target_parameters,\n ensure_weight_tying = ensure_weight_tying,\n **kwargs,\n )\n if not SUPPORTS_LOFTQ:\n del arguments[\"loftq_config\"]\n if not SUPPORTS_RSLORA:\n del arguments[\"use_rslora\"]\n\n _saved_temp_tokenizer = model._saved_temp_tokenizer\n\n lora_config = LoraConfig(**arguments)\n # First offload lm_head and embed_tokens to disk\n input_embeddings_device = model.get_input_embeddings().weight.device\n if is_classification:\n output_embeddings_device = model.score.weight.device\n else:\n output_embeddings_device = model.get_output_embeddings().weight.device\n\n if use_gradient_checkpointing == \"unsloth\":\n if train_embed_tokens:\n print(\"Unsloth: Offloading input_embeddings to disk to save VRAM\")\n offload_input_embeddings(model, temporary_location)\n\n # Remove old items to save VRAM\n for _ in range(3):\n gc.collect()\n clean_gpu_cache()\n\n if train_lm_head:\n print(\"Unsloth: Offloading output_embeddings to disk to save VRAM\")\n offload_output_embeddings(model, temporary_location)\n\n # Remove old items to save VRAM\n for _ in range(3):\n gc.collect()\n clean_gpu_cache()\n\n model = _get_peft_model(model, lora_config)\n # Fix LoraConfig.auto_mapping is None\n fix_lora_auto_mapping(model)\n\n # Apply QAT + LoRA if specified\n if qat_scheme is not None:\n print(\"Unsloth: Applying QAT to mitigate quantization degradation\")\n model = FastLlamaModel._prepare_for_qat(model, qat_scheme)\n\n model._saved_temp_tokenizer = _saved_temp_tokenizer\n\n model = FastLlamaModel.patch_peft_model(model, use_gradient_checkpointing)\n\n if ensure_weight_tying:\n try:\n input_embeddings = model.get_input_embeddings()\n output_embeddings = model.get_output_embeddings()\n\n if input_embeddings is not None and output_embeddings is not None:\n\n def _retie_parameter(target_module, source_module):\n if not hasattr(source_module, \"weight\"):\n return\n weight = source_module.weight\n # Remove existing registration to avoid \"attribute already exists\"\n if \"weight\" in getattr(target_module, \"_parameters\", {}):\n target_module._parameters.pop(\"weight\")\n if hasattr(target_module, \"weight\"):\n try:\n delattr(target_module, \"weight\")\n except Exception as exc:\n logger.warning_once(\n f\"Unsloth: Could not delete existing weight attr during retie on \"\n f\"{type(target_module).__name__}: {exc}\"\n )\n target_module.register_parameter(\"weight\", weight)\n\n # Tie trainable copies created by ModulesToSaveWrapper first (these are used in forward)\n if hasattr(input_embeddings, \"modules_to_save\") and hasattr(\n output_embeddings, \"modules_to_save\"\n ):\n if hasattr(\n input_embeddings.modules_to_save, \"default\"\n ) and hasattr(output_embeddings.modules_to_save, \"default\"):\n _retie_parameter(\n output_embeddings.modules_to_save.default,\n input_embeddings.modules_to_save.default,\n )\n\n # Tie original_module references as well if present\n if hasattr(input_embeddings, \"original_module\") and hasattr(\n output_embeddings, \"original_module\"\n ):\n _retie_parameter(\n output_embeddings.original_module,\n input_embeddings.original_module,\n )\n except Exception as e:\n logger.warning_once(\n f\"Unsloth: Failed to ensure weight tying between embeddings and lm_head: {e}\"\n )\n\n if train_embed_tokens:\n print(\"Unsloth: Training embed_tokens in mixed precision to save VRAM\")\n assert hasattr(model.get_input_embeddings(), \"modules_to_save\")\n\n _offload_frozen_module_for_training(\n model.get_input_embeddings(), DEVICE_TYPE_TORCH, offload_device = None\n )\n\n if train_lm_head:\n print(\"Unsloth: Training lm_head in mixed precision to save VRAM\")\n assert hasattr(model.get_output_embeddings(), \"modules_to_save\")\n\n _offload_frozen_module_for_training(\n model.get_output_embeddings(), DEVICE_TYPE_TORCH, offload_device = None\n )\n\n # Patch tokenizer to pad to the right\n internal_model = model\n while hasattr(internal_model, \"model\"):\n if hasattr(internal_model, \"_saved_temp_tokenizer\"):\n internal_model._saved_temp_tokenizer.padding_side = \"right\"\n # Also set is_loaded_in_8bit to disable incorrect DDP\n internal_model.is_loaded_in_8bit = True\n internal_model = internal_model.model\n if hasattr(internal_model, \"_saved_temp_tokenizer\"):\n internal_model._saved_temp_tokenizer.padding_side = \"right\"\n # Also set is_loaded_in_8bit to disable incorrect DDP\n internal_model.is_loaded_in_8bit = True\n\n # Clear deleted GPU items\n for _ in range(3):\n gc.collect()\n clean_gpu_cache()\n\n patch_peft_fast_inference(model)\n\n # Add for_inference and for_training\n model.for_training = functools.partial(FastLlamaModel.for_training, model)\n model.for_inference = functools.partial(FastLlamaModel.for_inference, model)\n m = model\n while hasattr(m, \"model\"):\n m.for_training = functools.partial(FastBaseModel.for_training, m)\n m.for_inference = functools.partial(FastBaseModel.for_inference, m)\n m = m.model\n return model\n\n @staticmethod\n def patch_peft_model(\n model,\n use_gradient_checkpointing = \"unsloth\",\n ):\n if os.environ.get(\"UNSLOTH_USE_NEW_MODEL\", \"0\") == \"1\":\n return FastBaseModel.patch_peft_model(\n model = model,\n use_gradient_checkpointing = use_gradient_checkpointing,\n )\n if not isinstance(model, PeftModelForCausalLM) and not isinstance(\n model, PeftModelForSequenceClassification\n ):\n raise TypeError(\n \"Unsloth: Your model needs to call `.get_peft_model` first!\"\n )\n\n # Get activation function\n model_type = model.config.model_type\n\n if model_type == \"llama\":\n apply_lora_mlp = apply_lora_mlp_swiglu\n elif model_type == \"mistral\":\n apply_lora_mlp = apply_lora_mlp_swiglu\n elif model_type == \"qwen2\":\n apply_lora_mlp = apply_lora_mlp_swiglu\n elif model_type == \"gemma\":\n apply_lora_mlp = apply_lora_mlp_geglu_approx\n elif model_type == \"gemma2\":\n apply_lora_mlp = apply_lora_mlp_geglu_approx\n elif model_type == \"cohere\":\n apply_lora_mlp = apply_lora_mlp_swiglu\n elif model_type == \"granite\":\n apply_lora_mlp = apply_lora_mlp_swiglu\n elif model_type == \"qwen3\":\n apply_lora_mlp = apply_lora_mlp_swiglu\n elif model_type == \"falcon_h1\":\n apply_lora_mlp = apply_lora_mlp_swiglu\n elif model_type == \"qwen3moe\":\n apply_lora_mlp = apply_lora_mlp_swiglu\n else:\n raise NotImplementedError(f\"Unsloth: {model_type} is not yet implemented!\")\n\n model = prepare_model_for_kbit_training(\n model,\n use_gradient_checkpointing = use_gradient_checkpointing,\n use_reentrant = True,\n )\n\n # Fix up config for transformers uploading PEFT\n for active_adapter in model.peft_config.keys():\n # Not necessary since we requires transformers >= 4.37\n if False:\n name = model.peft_config[active_adapter].base_model_name_or_path\n if name.startswith(\"unsloth/\") and name.endswith(\"-bnb-4bit\"):\n name = name[: len(name) - len(\"-bnb-4bit\")]\n model.peft_config[active_adapter].base_model_name_or_path = name\n pass\n # Add revision to enable future fast inference paths\n # [TODO] Bugs out!see https://github.com/unslothai/unsloth/issues/492\n # model.peft_config[active_adapter].revision = f\"unsloth\"\n\n from transformers.trainer import Trainer\n\n if Trainer._inner_training_loop.__name__ != \"_fast_inner_training_loop\":\n raise RuntimeError(\n \"Unsloth: Unsuccessfully patched Trainer! Please file a bug report!\"\n )\n\n # Fix loftq issues\n # loftq_config must not = None, but rather {}\n all_configs = model.peft_config\n for key, current_config in all_configs.items():\n if (\n hasattr(current_config, \"loftq_config\")\n and current_config.loftq_config is None\n ):\n new_args = current_config.__dict__\n new_args[\"loftq_config\"] = {}\n current_config = current_config.__class__(**new_args)\n all_configs[key] = current_config\n\n # Do patching\n n_mlp = 0\n n_qkv = 0\n n_o = 0\n\n active_adapter = (\n model.active_adapters[0]\n if hasattr(model, \"active_adapters\")\n else model.active_adapter\n )\n\n # Get dropout and bias\n lora_dropout = model.peft_config[active_adapter].lora_dropout\n bias = model.peft_config[active_adapter].bias\n\n # We also do not inplace edit QKV for Cohere!\n _apply_lora_mlp = (\n functools.partial(apply_lora_mlp, inplace = False)\n if model_type == \"cohere\"\n else apply_lora_mlp\n )\n\n if lora_dropout == 0 and bias == \"none\":\n for idx, layer in enumerate(model.model.model.layers):\n if model_type != \"falcon_h1\":\n # LoRAMLP.apply doesn't have functionality for gate and down multipliers yet.\n # Don't patch falcon h1 for the time being.\n\n # MLP patching\n mlp_module = layer.mlp\n gate_proj = mlp_module.gate_proj\n up_proj = mlp_module.up_proj\n down_proj = mlp_module.down_proj\n\n if (\n hasattr(gate_proj, \"lora_A\")\n and hasattr(up_proj, \"lora_A\")\n and hasattr(down_proj, \"lora_A\")\n and (getattr(gate_proj, \"base_layer\", gate_proj).bias is None)\n and (getattr(up_proj, \"base_layer\", up_proj).bias is None)\n and (getattr(down_proj, \"base_layer\", down_proj).bias is None)\n and (\n len(getattr(gate_proj, \"lora_magnitude_vector\", []) or [])\n == 0\n )\n and (\n len(getattr(up_proj, \"lora_magnitude_vector\", []) or [])\n == 0\n )\n and (\n len(getattr(down_proj, \"lora_magnitude_vector\", []) or [])\n == 0\n )\n ):\n # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module\n if hasattr(mlp_module, \"_unsloth_forward\"):\n # then we've patched the mlp to use TiledMLP\n mlp_module._unsloth_forward = types.MethodType(\n _apply_lora_mlp, mlp_module\n )\n else:\n mlp_module.forward = types.MethodType(\n _apply_lora_mlp, mlp_module\n )\n n_mlp += 1\n else:\n logger.warning_once(\n \"Not an error, but Unsloth cannot patch MLP layers with our manual autograd engine since either LoRA adapters\\n\"\n \"are not enabled or a bias term (like in Qwen) is used.\"\n )\n\n # QKV attention patching\n q_proj = layer.self_attn.q_proj\n k_proj = layer.self_attn.k_proj\n v_proj = layer.self_attn.v_proj\n if (\n hasattr(q_proj, \"lora_A\")\n and hasattr(k_proj, \"lora_A\")\n and hasattr(v_proj, \"lora_A\")\n and (getattr(q_proj, \"base_layer\", q_proj).bias is None)\n and (getattr(k_proj, \"base_layer\", k_proj).bias is None)\n and (getattr(v_proj, \"base_layer\", v_proj).bias is None)\n and (len(getattr(q_proj, \"lora_magnitude_vector\", []) or []) == 0)\n and (len(getattr(k_proj, \"lora_magnitude_vector\", []) or []) == 0)\n and (len(getattr(v_proj, \"lora_magnitude_vector\", []) or []) == 0)\n ):\n layer.self_attn.apply_qkv = apply_lora_qkv\n n_qkv += 1\n else:\n if model_type == \"qwen2\":\n n_qkv += 1\n else:\n logger.warning_once(\n \"Not an error, but Unsloth cannot patch Attention layers with our manual autograd engine since either LoRA adapters\\n\"\n \"are not enabled or a bias term (like in Qwen) is used.\"\n )\n\n # O attention patching\n o_proj = layer.self_attn.o_proj\n if (\n hasattr(o_proj, \"lora_A\")\n and (getattr(o_proj, \"base_layer\", o_proj).bias is None)\n and (len(getattr(o_proj, \"lora_magnitude_vector\", []) or []) == 0)\n ):\n layer.self_attn.apply_o = apply_lora_o\n n_o += 1\n else:\n logger.warning_once(\n \"Not an error, but Unsloth cannot patch O projection layer with our manual autograd engine since either LoRA adapters\\n\"\n \"are not enabled or a bias term (like in Qwen) is used.\"\n )\n\n logger.warning_once(\n f\"Unsloth {__version__} patched {len(model.model.model.layers)} layers with \"\n f\"{n_qkv} QKV layers, {n_o} O layers and {n_mlp} MLP layers.\",\n )\n patch_saving_functions(model)\n\n # Patch cross entropy loss labels\n # Fixes https://github.com/unslothai/unsloth/issues/10\n max_seq_length = model.max_seq_length\n # extra_ignored_labels = torch.full((max_seq_length, 1), -100, device = \"cuda:0\")\n # model.model.extra_ignored_labels = extra_ignored_labels\n internal_model = model\n while hasattr(internal_model, \"model\"):\n internal_model.max_seq_length = max_seq_length\n internal_model = internal_model.model\n internal_model.max_seq_length = max_seq_length\n # Save to modules as well\n for module in model.modules():\n module.max_seq_length = max_seq_length\n\n # Patch tokenizer to pad to the right\n internal_model = model\n while hasattr(internal_model, \"model\"):\n if hasattr(internal_model, \"_saved_temp_tokenizer\"):\n internal_model._saved_temp_tokenizer.padding_side = \"right\"\n internal_model = internal_model.model\n if hasattr(internal_model, \"_saved_temp_tokenizer\"):\n internal_model._saved_temp_tokenizer.padding_side = \"right\"\n\n # Clear deleted GPU items\n for _ in range(3):\n gc.collect()\n clean_gpu_cache()\n\n patch_peft_fast_inference(model)\n\n # Add for_inference and for_training\n model.for_training = functools.partial(FastLlamaModel.for_training, model)\n model.for_inference = functools.partial(FastLlamaModel.for_inference, model)\n m = model\n while hasattr(m, \"model\"):\n m.for_training = functools.partial(FastBaseModel.for_training, m)\n m.for_inference = functools.partial(FastBaseModel.for_inference, m)\n m = m.model\n return model\n\n @staticmethod\n def for_inference(model):\n if not hasattr(model, \"parameters\"):\n raise TypeError(\n \"Unsloth: I think you're passing a tokenizer, not the model to for_inference!\"\n )\n\n def _for_inference(m):\n if hasattr(m, \"gradient_checkpointing\"):\n m.gradient_checkpointing = False\n if hasattr(m, \"training\"):\n m.training = False\n # Pad tokenizer to the left\n if hasattr(m, \"_saved_temp_tokenizer\"):\n m._saved_temp_tokenizer.padding_side = \"left\"\n # Set a flag for generation!\n m._flag_for_generation = True\n\n m = model\n while hasattr(m, \"model\"):\n _for_inference(m)\n m = m.model\n _for_inference(m)\n model.eval() # to turn off training on modules deeper in\n\n # Since transformers 4.53, must turn off explicitly\n for module in model.modules():\n if hasattr(module, \"gradient_checkpointing\"):\n module.gradient_checkpointing = False\n\n # Also disable training for embeddings for NEFTune\n if hasattr(model, \"get_input_embeddings\"):\n embeddings = model.get_input_embeddings()\n if hasattr(embeddings, \"training\"):\n embeddings.training = False\n if hasattr(model, \"get_output_embeddings\"):\n embeddings = model.get_output_embeddings()\n if hasattr(embeddings, \"training\"):\n embeddings.training = False\n return model\n\n @staticmethod\n def for_training(model, use_gradient_checkpointing = True):\n if not hasattr(model, \"parameters\"):\n raise TypeError(\n \"Unsloth: I think you're passing a tokenizer, not the model to for_training!\"\n )\n\n # Delete all fast inference loras\n for param in model.parameters():\n if hasattr(param, \"_fast_lora\"):\n del param._fast_lora\n\n def _for_training(m):\n if hasattr(m, \"gradient_checkpointing\"):\n m.gradient_checkpointing = use_gradient_checkpointing\n if hasattr(m, \"training\"):\n m.training = True\n # Pad tokenizer to the left\n if hasattr(m, \"_saved_temp_tokenizer\"):\n m._saved_temp_tokenizer.padding_side = \"right\"\n # Set a flag for generation!\n if hasattr(m, \"_flag_for_generation\"):\n del m._flag_for_generation\n\n m = model\n while hasattr(m, \"model\"):\n _for_training(m)\n m = m.model\n _for_training(m)\n model.train() # to turn on training on modules deeper in\n\n # Since transformers 4.53, must turn on explicitly\n for module in model.modules():\n if hasattr(module, \"gradient_checkpointing\"):\n module.gradient_checkpointing = use_gradient_checkpointing\n\n # Also re-enable training for embeddings for NEFTune\n if hasattr(model, \"get_input_embeddings\"):\n embeddings = model.get_input_embeddings()\n if hasattr(embeddings, \"training\"):\n embeddings.training = True\n if hasattr(model, \"get_output_embeddings\"):\n embeddings = model.get_output_embeddings()\n if hasattr(embeddings, \"training\"):\n embeddings.training = True\n return model\n\n\nfrom .rl import PatchFastRL\n\nPatchFastRL(FastLanguageModel = FastLlamaModel)\n"}}}} |