From 2afeb378392ac6f287c885c85113b31c7ac5d7cd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 5 Mar 2025 05:13:32 -0800 Subject: [PATCH] Many bug fixes (#1900) * Update rl.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * autocast * Update rl_replacements.py * Update llama.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update llama.py * Update rl_replacements.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update rl_replacements.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update pyproject.toml * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update llama.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update llama.py * Update _utils.py * Update llama.py * Update _utils.py * Update rl_replacements.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update rl_replacements.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * GRPO optimized * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Selective Log softmax * Fix GRPO bsz * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Fix TRL * Metrics GRPO * Update rl_replacements.py * Update rl_replacements.py * No compile * Update rl.py * Remove docs * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl_replacements.py * Update rl.py * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * llama-quantize on WINDOWS WSL error fix - edit save.py (gguf saving breaks) (#1649) * edit save.py to fix gguf saving breaks. * add check for .exe or not exe file extension for linux and windows * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update llama.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update llama.py * Update llama.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl.py * Update rl.py * Update rl_replacements.py * Update rl.py * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * unsloth_num_chunks * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py (#1754) Fix typo in comment: know -> now. This was printed when running the Llama3.1_(8B)-GRPO.ipynb example notebook, so I'd expect others to run into it as well. * Optional logits * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * Update rl_replacements.py * Update rl.py * Update rl.py * Update rl.py * Update rl.py * fix an import error (#1767) * fix an import error * Delete .gitignore * Update loader.py * Update save.py --------- Co-authored-by: Daniel Han * SamplingParams * Convert mask to float (#1762) * [Windows Support] Add latest `xformers` wheels to pyproject.toml (#1753) * Add latest xformers * Add a couple of lines to docs * vLLMSamplingParams * Update __init__.py * default num_chunks == -1 * Versioning * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update _utils.py * Update rl_replacements.py * Update rl_replacements.py * Update pyproject.toml * Update pyproject.toml * Export Model to ollama.com (#1648) * Ollama Export Model to ollama.com Signed-off-by: Jyotin Goel * Check for model_name Signed-off-by: Jyotin Goel * subprocess use instead of requests | added check for ollama server Signed-off-by: Jyotin Goel * create_ollama_model Signed-off-by: Jyotin Goel * create_ollama_model | fix Signed-off-by: Jyotin Goel * Push to Ollama Signed-off-by: Jyotin Goel --------- Signed-off-by: Jyotin Goel * Update cross_entropy_loss.py * torch_cuda_device * Update utils.py * Update utils.py * Update utils.py * device * device * Update loader.py * Update llama.py * Update README.md * Update llama.py * Update llama.py * Update _utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update utils.py * Update utils.py * Update utils.py * Update utils.py * __version__ * Update rl.py * Bug fixes * Bug fixes * Update llama.py * Update _utils.py * _wrap_fast_inference * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update _utils.py * SFT dataset prepare * Update pyproject.toml * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl.py * Update llama.py * Update llama.py * Update utils.py * bug fix * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update __init__.py --------- Signed-off-by: Jyotin Goel Co-authored-by: Gennadii Manzhos <105049664+everythingisc00l@users.noreply.github.com> Co-authored-by: Seth Weidman Co-authored-by: Nino Risteski <95188570+NinoRisteski@users.noreply.github.com> Co-authored-by: Edd <68678137+Erland366@users.noreply.github.com> Co-authored-by: Ben <6579034+versipellis@users.noreply.github.com> Co-authored-by: Jyotin Goel <120490013+gjyotin305@users.noreply.github.com> --- pyproject.toml | 34 ++-- unsloth/__init__.py | 2 +- unsloth/kernels/utils.py | 31 ++-- unsloth/models/_utils.py | 78 ++++----- unsloth/models/llama.py | 260 +++++++++++++----------------- unsloth/models/rl.py | 2 +- unsloth/models/rl_replacements.py | 19 +++ 7 files changed, 196 insertions(+), 230 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 73e69dcd4a..5a9d92202a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "unsloth" dynamic = ["version"] description = "2-5X faster LLM finetuning" readme = "README.md" -requires-python = ">=3.9" +requires-python = ">=3.9,<=3.12" license = {file = "LICENSE"} keywords = ["ai", "llm",] authors = [ @@ -39,8 +39,8 @@ triton = [ "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.2.0-windows.post10/triton-3.2.0-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'" ] -windows=[ - "unsloth_zoo>=2025.3.1", +huggingface = [ + "unsloth_zoo>=2025.3.2", "packaging", "tyro", "transformers>=4.46.1,!=4.47.0", @@ -51,34 +51,18 @@ windows=[ "wheel>=0.42.0", "numpy", "accelerate>=0.34.1", - "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,!=0.15.0", + "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,!=0.15.0,<=0.15.2", "peft>=0.7.1,!=0.11.0", "protobuf<4.0.0", "huggingface_hub", "hf_transfer", "unsloth[triton]", +] +windows=[ + "unsloth[huggingface]", "bitsandbytes>=0.41.1 ; platform_system == 'Windows'", "xformers>=0.0.22.post7 ; platform_system == 'Windows'", ] -huggingface = [ - "unsloth_zoo>=2025.3.1", - "packaging", - "tyro", - "transformers>=4.46.1,!=4.47.0", - "datasets>=2.16.0", - "sentencepiece>=0.2.0", - "tqdm", - "psutil", - "wheel>=0.42.0", - "numpy", - "accelerate>=0.34.1", - "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,!=0.15.0", - "peft>=0.7.1,!=0.11.0", - "protobuf<4.0.0", - "huggingface_hub", - "hf_transfer", - "unsloth[triton]", -] cu118only = [ "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.22.post7%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'", "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.22.post7%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10' and platform_system == 'Linux'", @@ -370,7 +354,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3", ] colab-new = [ - "unsloth_zoo>=2025.2.7", + "unsloth_zoo>=2025.3.1", "packaging", "tyro", "transformers>=4.46.1,!=4.47.0", @@ -388,7 +372,7 @@ colab-new = [ ] colab-no-deps = [ "accelerate>=0.34.1", - "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,!=0.15.0", + "trl>=0.7.9,!=0.9.0,!=0.9.1,!=0.9.2,!=0.9.3,!=0.15.0,<=0.15.2", "peft>=0.7.1", "xformers", "bitsandbytes>=0.46.1", diff --git a/unsloth/__init__.py b/unsloth/__init__.py index c8f2926985..8439ab8212 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -198,7 +198,7 @@ pass # Check for unsloth_zoo try: unsloth_zoo_version = importlib_version("unsloth_zoo") - if Version(unsloth_zoo_version) < Version("2025.3.1"): + if Version(unsloth_zoo_version) < Version("2025.3.2"): try: os.system("pip install --upgrade --no-cache-dir --no-deps unsloth_zoo") except: diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 8da152bcb3..db1d73c340 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -104,6 +104,11 @@ cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_ cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4 cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemm_4bit_inference_naive_fp16 cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemm_4bit_inference_naive_bf16 +torch_mm = torch.mm +torch_mv = torch.mv +torch_matmul = torch.matmul +torch_addmm = torch.addmm +torch_empty = torch.empty def QUANT_STATE(W): return getattr(W, "quant_state", None) @@ -194,8 +199,8 @@ if HAS_CUDA_STREAM: WEIGHT_BUFFER = WEIGHT_BUFFERS[device_index] ABSMAX_BUFFER = ABSMAX_BUFFERS[device_index] if WEIGHT_BUFFER is None: - WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = device, requires_grad = False) - ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch_empty(size, dtype = dtype, device = device, requires_grad = False) + ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size) if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax) @@ -204,11 +209,11 @@ if HAS_CUDA_STREAM: out_absmax = ABSMAX_BUFFER[:n_elements_absmax] else: if out is None: - out = torch.empty(shape, dtype = dtype, device = device, requires_grad = False) + out = torch_empty(shape, dtype = dtype, device = device, requires_grad = False) else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) pass # NF4 dequantization of statistics @@ -258,11 +263,11 @@ else: # Create weight matrix if out is None: - out = torch.empty(shape, dtype = dtype, device = device, requires_grad = False) + out = torch_empty(shape, dtype = dtype, device = device, requires_grad = False) else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) # Do dequantization ptr_out_absmax = get_ptr(out_absmax) @@ -286,7 +291,7 @@ pass if HAS_CUDA_STREAM: def fast_gemv(X, W, quant_state, out = None): - if quant_state is None: return torch.matmul(X, W, out = out) + if quant_state is None: return torch_matmul(X, W, out = out) # For fast X @ W where seq_len == 1 # From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469 _, q_len, hd = X.shape @@ -318,7 +323,7 @@ if HAS_CUDA_STREAM: bout = shape[0] if out is None: - out = torch.empty((1, 1, bout,), dtype = dtype, device = device) + out = torch_empty((1, 1, bout,), dtype = dtype, device = device) # else: # assert(out.shape == (1, 1, bout,)) # pass @@ -336,7 +341,7 @@ if HAS_CUDA_STREAM: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch.empty(absmax.shape, dtype = torch.float32, device = device) + df = torch_empty(absmax.shape, dtype = torch.float32, device = device) with torch_cuda_device(device): cdequantize_blockwise_fp32( get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), @@ -385,7 +390,7 @@ else: device = W.device if out is None: - out = torch.empty((1, 1, bout,), dtype = dtype, device = device) + out = torch_empty((1, 1, bout,), dtype = dtype, device = device) # else: # assert(out.shape == (1, 1, bout,)) # pass @@ -403,7 +408,7 @@ else: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch.empty(absmax.shape, dtype = torch.float32, device = device) + df = torch_empty(absmax.shape, dtype = torch.float32, device = device) cdequantize_blockwise_fp32( get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), @@ -423,10 +428,6 @@ else: pass -torch_mm = torch.mm -torch_mv = torch.mv -torch_matmul = torch.matmul -torch_addmm = torch.addmm def fast_linear_forward(proj, X, temp_lora = None, out = None): W, W_quant, lora_A, lora_B, lora_S, bias = get_lora_parameters_bias(proj) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 43828a358a..66926bca10 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2025.3.3" +__version__ = "2025.3.4" __all__ = [ "SUPPORTS_BFLOAT16", @@ -39,8 +39,8 @@ __all__ = [ "create_boolean_mask", "torch_amp_custom_fwd", "torch_amp_custom_bwd", - "accelerate_old_send_to_device", - "accelerate_new_send_to_device", + # "accelerate_old_send_to_device", + # "accelerate_new_send_to_device", "patch_gradient_accumulation_fix", "patch_compiling_bitsandbytes", "patch_regional_compilation", @@ -241,24 +241,24 @@ pass # ============================================= # Fix KeyError: 'Cache only has 0 layers, attempted to access layer with index 0' -import transformers.cache_utils -if hasattr(transformers.cache_utils, "DynamicCache") and \ - transformers.cache_utils.DynamicCache.__getitem__.__name__ != "__cache_utils_getitem__": +# import transformers.cache_utils +# if hasattr(transformers.cache_utils, "DynamicCache") and \ +# transformers.cache_utils.DynamicCache.__getitem__.__name__ != "__cache_utils_getitem__": - source = inspect.getsource(transformers.cache_utils.DynamicCache.__getitem__) - start = source.find("def") - spaces = start*" " - source = source.split("\n") - source = "\n".join(x[start:] for x in source) - where = source.find("raise KeyError") - source = source[:where] + \ - f"if len(self) == 0:\n{spaces}{spaces}"\ - " raise RuntimeError('Unsloth: You must call `FastLanguageModel.for_inference(model)` before doing inference for Unsloth models.')\n" + \ - f"{spaces}{spaces}else:\n{spaces}{spaces}{spaces}" + source[where:] - source = source.replace("__getitem__", "__cache_utils_getitem__", 1) - exec(source) - transformers.cache_utils.DynamicCache.__getitem__ = __cache_utils_getitem__ -pass +# source = inspect.getsource(transformers.cache_utils.DynamicCache.__getitem__) +# start = source.find("def") +# spaces = start*" " +# source = source.split("\n") +# source = "\n".join(x[start:] for x in source) +# where = source.find("raise KeyError") +# source = source[:where] + \ +# f"if len(self) == 0:\n{spaces}{spaces}"\ +# " raise RuntimeError('Unsloth: You must call `FastLanguageModel.for_inference(model)` before doing inference for Unsloth models.')\n" + \ +# f"{spaces}{spaces}else:\n{spaces}{spaces}{spaces}" + source[where:] +# source = source.replace("__getitem__", "__cache_utils_getitem__", 1) +# exec(source) +# transformers.cache_utils.DynamicCache.__getitem__ = __cache_utils_getitem__ +# pass # ============================================= # ============================================= @@ -411,25 +411,25 @@ pass # ============================================= # Fix new Xformers versions TypeError: Multiple dispatch failed for 'torch._ops.aten.to.dtype_layout' -accelerate_old_send_to_device = None -accelerate_new_send_to_device = None -if xformers_version is not None and Version(xformers_version) >= Version("0.0.27"): - import accelerate.utils.operations - if hasattr(accelerate.utils.operations, "send_to_device") and \ - accelerate.utils.operations.send_to_device.__name__ != "_fixed_send_to_device": - accelerate_old_send_to_device = accelerate.utils.operations.send_to_device - from accelerate.utils.operations import * - send_to_device = inspect.getsource(accelerate.utils.operations.send_to_device) - send_to_device = re.sub( - r"([ ]{4,})return tensor\.to\(device\)", - r"\1try: return tensor.to(device)\n\1except: return tensor", - send_to_device, - ).replace("def send_to_device", "def _fixed_send_to_device") - exec(send_to_device) - # accelerate.utils.operations.send_to_device = _fixed_send_to_device - accelerate_new_send_to_device = _fixed_send_to_device - pass -pass +# accelerate_old_send_to_device = None +# accelerate_new_send_to_device = None +# if xformers_version is not None and Version(xformers_version) >= Version("0.0.27"): +# import accelerate.utils.operations +# if hasattr(accelerate.utils.operations, "send_to_device") and \ +# accelerate.utils.operations.send_to_device.__name__ != "_fixed_send_to_device": +# accelerate_old_send_to_device = accelerate.utils.operations.send_to_device +# from accelerate.utils.operations import * +# send_to_device = inspect.getsource(accelerate.utils.operations.send_to_device) +# send_to_device = re.sub( +# r"([ ]{4,})return tensor\.to\(device\)", +# r"\1try: return tensor.to(device)\n\1except: return tensor", +# send_to_device, +# ).replace("def send_to_device", "def _fixed_send_to_device") +# exec(send_to_device) +# # accelerate.utils.operations.send_to_device = _fixed_send_to_device +# accelerate_new_send_to_device = _fixed_send_to_device +# pass +# pass # Transformers 4.46 breaks dynamic caching. This is a hack import transformers.generation.configuration_utils diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index bcabbd5125..3dacf5cdd5 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -15,7 +15,7 @@ import torch import gc import math -from functools import partial +import functools from typing import Optional, Tuple, List, Union from ._utils import * from ._utils import patch_unsloth_smart_gradient_checkpointing @@ -65,6 +65,7 @@ from peft import LoraConfig, TaskType, get_peft_model as _get_peft_model from peft import PeftModelForCausalLM from ..save import patch_saving_functions import re, os, inspect, math, sys +import types try: from huggingface_hub.utils import get_token except: @@ -217,14 +218,14 @@ def LlamaAttention_fast_forward_inference( RH_Q = self.RH_Q RH_Q[:,:,:,:h] = Qn[:,:,:,h:] RH_Q[:,:,:,h:] = Qn[:,:,:,:h] - torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h]) + RH_Q[:,:,:,:h].neg_() # torch.neg(RH_Q[:,:,:,:h], out = RH_Q[:,:,:,:h]) Qn *= cos Qn.addcmul_(RH_Q, sin) RH_K = RH_Q[:,:n_kv_heads,:,:] # torch.empty((n_kv_heads, 1, head_dim), dtype = dtype, device = "cuda:0") RH_K[:,:,:,:h] = Kn[:,:,:,h:] RH_K[:,:,:,h:] = Kn[:,:,:,:h] - torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h]) + RH_K[:,:,:,:h].neg_() #torch.neg(RH_K[:,:,:,:h], out = RH_K[:,:,:,:h]) Kn *= cos Kn.addcmul_(RH_K, sin) @@ -400,19 +401,20 @@ def LlamaAttention_fast_forward( else: # Extend RoPE dynamically to fit in VRA rotary_emb = self.rotary_emb - rotary_emb.extend_rope_embedding(V, seq_len=kv_seq_len) + rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len) if position_ids is None: # Useful for LongRoPE cos, sin = rotary_emb.get_cached(kv_seq_len) else: - cos, sin = rotary_emb(V, seq_len=kv_seq_len) + cos, sin = rotary_emb(V, seq_len = kv_seq_len) - Q, K = ( - fast_rope_embedding(Q, K, cos, sin) - if position_ids is None - else inplace_rope_embedding(Q, K, cos, sin, position_ids) - ) + # Q, K = ( + # fast_rope_embedding(Q, K, cos, sin) + # if position_ids is None + # else inplace_rope_embedding(Q, K, cos, sin, position_ids) + # ) + Q, K = fast_rope_embedding(Q, K, cos, sin) if past_key_value is not None: K = torch.cat([past_key_value[0], K], dim = 2) @@ -924,7 +926,6 @@ def LlamaModel_fast_forward_inference( X = X.to(self.config.torch_dtype) bsz, q_len, hd = X.shape assert(q_len == 1) - # Get saved buffers to reduce memory movement residual = torch.empty((bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") _XX = torch.empty((2, bsz, q_len, hd), dtype = torch.float32, device = "cuda:0") @@ -1020,7 +1021,6 @@ def CausalLM_fast_forward(fast_forward_inference): logits_to_keep: Optional[int] = 0, *args, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: - if past_key_values is not None: outputs = fast_forward_inference( self, @@ -1069,7 +1069,7 @@ def CausalLM_fast_forward(fast_forward_inference): if labels is not None: labels = labels.to(lm_head_device) # Output last hidden states without logits if asked - if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1": + if self.training and os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1": if num_logits_to_keep != 0: hidden_states = hidden_states[:, -num_logits_to_keep:, :] return CausalLMOutputWithPast( @@ -1534,78 +1534,58 @@ class LongRopeRotaryEmbedding(torch.nn.Module): pass -def _wrap_fast_inference(generate, device_type, dtype, model): - # Wraps inference with bfloat16 / float16 - @torch.inference_mode - def _fast_generate(*args, **kwargs): +def unsloth_fast_generate( + self, + *args, + **kwargs, +): + FastLlamaModel.for_inference(self) - if hasattr(model, "config") and hasattr(model.config, "max_position_embeddings"): - if "input_ids" in kwargs and kwargs["input_ids"] is not None and "max_new_tokens" in kwargs: - if kwargs["input_ids"].shape[-1] + kwargs["max_new_tokens"] > model.config.max_position_embeddings: - raise ValueError( - f'Unsloth: input length {kwargs["input_ids"].shape[-1]} + max_new_tokens {kwargs["max_new_tokens"]} exceeds the maximum sequence length of {model.config.max_position_embeddings}!\n'\ - 'You will need to do long context extension by increasing the `max_seq_length` in `FastLanguageModel.from_pretrained`.' - ) - pass + dtype = _get_dtype(self.config.torch_dtype) - # Set a flag for generation! - internal_model = model - while hasattr(internal_model, "model"): - internal_model._flag_for_generation = True - internal_model = internal_model.model - pass - internal_model._flag_for_generation = True - - # Must patch accelerate for Xformers - if accelerate_new_send_to_device is not None: - import accelerate.utils.operations - accelerate.utils.operations.send_to_device = accelerate_new_send_to_device - pass - - # For newer HF - kwargs["cache_implementation"] = "dynamic" - # For num_logits_to_keep - kwargs["num_logits_to_keep"] = 1 - - # Remove token_type_ids - kwargs.pop("token_type_ids", None) - - # Check pad_token - model_eos_token_id = getattr(model.config, "eos_token_id", None) - if model_eos_token_id is not None and hasattr(model_eos_token_id, "__iter__"): - model_eos_token_id = model_eos_token_id[0] - - kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id) - - # Set pad token - # old_pad_token_id = getattr(model.config, "pad_token_id", None) - # old_eos_token_id = getattr(model.config, "eos_token_id", None) - # model.config.pad_token_id = old_eos_token_id - - # Autocasted - with torch.autocast(device_type = device_type, dtype = dtype): - output = generate(*args, **kwargs) - pass - - # Revert - # model.config.pad_token_id = old_pad_token_id - - # Unset a flag for generation! - internal_model = model - while hasattr(internal_model, "model"): - if hasattr(internal_model, "_flag_for_generation"): del internal_model._flag_for_generation - internal_model = internal_model.model - pass - if hasattr(internal_model, "_flag_for_generation"): del internal_model._flag_for_generation - - # Return accelerate back - if accelerate_new_send_to_device is not None: - accelerate.utils.operations.send_to_device = accelerate_old_send_to_device - pass - - return output + if hasattr(self, "config") and hasattr(self.config, "max_position_embeddings"): + if "input_ids" in kwargs and kwargs["input_ids"] is not None and "max_new_tokens" in kwargs: + if kwargs["input_ids"].shape[-1] + kwargs["max_new_tokens"] > self.config.max_position_embeddings: + raise ValueError( + f'Unsloth: input length {kwargs["input_ids"].shape[-1]} + max_new_tokens {kwargs["max_new_tokens"]} exceeds the maximum sequence length of {model.config.max_position_embeddings}!\n'\ + 'You will need to do long context extension by increasing the `max_seq_length` in `FastLanguageModel.from_pretrained`.' + ) pass - return _fast_generate + + # Must patch accelerate for Xformers + # if accelerate_new_send_to_device is not None: + # import accelerate.utils.operations + # accelerate.utils.operations.send_to_device = accelerate_new_send_to_device + # pass + + # For newer HF + kwargs["cache_implementation"] = "dynamic" + # For num_logits_to_keep + kwargs["num_logits_to_keep"] = 1 + + # Remove token_type_ids + kwargs.pop("token_type_ids", None) + + # Check pad_token + model_eos_token_id = getattr(self.config, "eos_token_id", None) + if model_eos_token_id is not None and hasattr(model_eos_token_id, "__iter__"): + model_eos_token_id = model_eos_token_id[0] + + kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id) + + # Mixed precision autocast + with torch.inference_mode(), torch.autocast(device_type = "cuda", dtype = dtype): + output = self._old_generate(*args, **kwargs) + pass + + # Return accelerate back + # if accelerate_new_send_to_device is not None: + # accelerate.utils.operations.send_to_device = accelerate_old_send_to_device + # pass + + FastLlamaModel.for_training(self) + + return output pass @@ -1682,8 +1662,12 @@ class FastLlamaModel: gpu_stats = torch.cuda.get_device_properties(0) max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) + from importlib.metadata import version as importlib_version + try: vllm_version = f" vLLM: {importlib_version('vllm')}." + except: vllm_version = "" + statistics = \ - f"==((====))== Unsloth {__version__}: Fast {model_patcher.__name__[4:-5]} patching. Transformers: {transformers_version}.\n"\ + f"==((====))== Unsloth {__version__}: Fast {model_patcher.__name__[4:-5]} patching. Transformers: {transformers_version}.{vllm_version}\n"\ f" {chr(92)}{chr(92)} /| {gpu_stats.name}. Num GPUs = {torch.cuda.device_count()}. Max memory: {max_memory} GB. Platform: {platform_system}.\n"\ f"O^O/ {chr(92)}_/ {chr(92)} Torch: {torch.__version__}. CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {torch.version.cuda}. Triton: {triton_version}\n"\ f"{chr(92)} / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\n"\ @@ -1825,7 +1809,7 @@ class FastLlamaModel: model = convert_vllm_to_huggingface(quant_state_dict, model_config, dtype) model.vllm_engine = llm model.fast_generate = model.vllm_engine.generate - model.fast_generate_batches = partial(generate_batches, model.vllm_engine) + model.fast_generate_batches = functools.partial(generate_batches, model.vllm_engine) pass # Return old flag os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer @@ -1986,6 +1970,11 @@ class FastLlamaModel: layer.self_attn.rotary_emb = rotary_emb pass + # Patch generate + if model.generate.__name__ != "unsloth_fast_generate": + model._old_generate = model.generate + unsloth_fast_generate.__doc__ = model._old_generate.__doc__ + model.generate = types.MethodType(unsloth_fast_generate, model) return model, tokenizer pass @@ -2410,12 +2399,20 @@ class FastLlamaModel: model.fast_generate_batches = vllm_fast_generate_batches # Also saving and loading LoRA - from functools import partial from unsloth_zoo.vllm_utils import save_lora, load_lora - model.save_lora = partial(save_lora, model) - model.load_lora = partial(load_lora, model) + model.save_lora = functools.partial(save_lora, model) + model.load_lora = functools.partial(load_lora, model) pass + # Add for_inference and for_training + model.for_training = functools.partial(FastLlamaModel.for_training, model) + model.for_inference = functools.partial(FastLlamaModel.for_inference, model) + + # Patch generate + if model.generate.__name__ != "unsloth_fast_generate": + model._old_generate = model.generate + unsloth_fast_generate.__doc__ = model._old_generate.__doc__ + model.generate = types.MethodType(unsloth_fast_generate, model) return model pass @@ -2486,7 +2483,6 @@ class FastLlamaModel: n_mlp = 0 n_qkv = 0 n_o = 0 - import types active_adapter = model.active_adapters[0] if \ hasattr(model, "active_adapters") else model.active_adapter @@ -2496,9 +2492,8 @@ class FastLlamaModel: bias = model.peft_config[active_adapter].bias # We also do not inplace edit QKV for Cohere! - from functools import partial _apply_lora_mlp = \ - partial(apply_lora_mlp, inplace = False) \ + functools.partial(apply_lora_mlp, inplace = False) \ if model_type == "cohere" else \ apply_lora_mlp pass @@ -2611,52 +2606,30 @@ class FastLlamaModel: pass # Add for_inference and for_training - model.for_training = partial(FastLlamaModel.for_training, model) - model.for_inference = partial(FastLlamaModel.for_inference, model) + model.for_training = functools.partial(FastLlamaModel.for_training, model) + model.for_inference = functools.partial(FastLlamaModel.for_inference, model) return model pass @staticmethod def for_inference(model): - # if model.config.model_type == "qwen2": - # FastLlamaModel.for_training(model) - # return - # pass + if not hasattr(model, "parameters"): + raise TypeError("Unsloth: I think you're passing a tokenizer, not the model to for_inference!") + def _for_inference(m): + if hasattr(m, "gradient_checkpointing"): m.gradient_checkpointing = False + if hasattr(m, "training"): m.training = False + # Pad tokenizer to the left + if hasattr(m, "_saved_temp_tokenizer"): m._saved_temp_tokenizer.padding_side = "left" + # Set a flag for generation! + m._flag_for_generation = True + pass m = model while hasattr(m, "model"): - if hasattr(m, "gradient_checkpointing"): - m.gradient_checkpointing = False - if hasattr(m, "training"): - m.training = False - # Pad tokenizer to the left - if hasattr(m, "_saved_temp_tokenizer"): - m._saved_temp_tokenizer.padding_side = "left" + _for_inference(m) m = m.model - pass - if hasattr(m, "gradient_checkpointing"): - m.gradient_checkpointing = False - if hasattr(m, "training"): - m.training = False - # Pad tokenizer to the left - if hasattr(m, "_saved_temp_tokenizer"): - m._saved_temp_tokenizer.padding_side = "left" - - # Also check if lm_head / embeddings are trained - internal_model = model - while not hasattr(internal_model, "lm_head"): - internal_model = internal_model.model - pass - lm_head = internal_model.lm_head.weight - device_type = lm_head.device.type - dtype = _get_dtype(model.config.torch_dtype) - - # Wrap model.generate - if model.generate.__name__ != "_fast_generate": - model._unwrapped_old_generate = model.generate - model.generate = _wrap_fast_inference(model.generate, device_type, dtype, model) - pass + _for_inference(m) # Also disable training for embeddings for NEFTune if hasattr(model, "get_input_embeddings"): @@ -2667,13 +2640,14 @@ class FastLlamaModel: embeddings = model.get_output_embeddings() if hasattr(embeddings, "training"): embeddings.training = False pass - return model pass @staticmethod def for_training(model, use_gradient_checkpointing = True): + if not hasattr(model, "parameters"): + raise TypeError("Unsloth: I think you're passing a tokenizer, not the model to for_training!") # Delete all fast inference loras for param in model.parameters(): @@ -2681,30 +2655,19 @@ class FastLlamaModel: del param._fast_lora pass + def _for_training(m): + if hasattr(m, "gradient_checkpointing"): m.gradient_checkpointing = use_gradient_checkpointing + if hasattr(m, "training"): m.training = True + # Pad tokenizer to the left + if hasattr(m, "_saved_temp_tokenizer"): m._saved_temp_tokenizer.padding_side = "right" + # Set a flag for generation! + if hasattr(m, "_flag_for_generation"): del m._flag_for_generation + pass m = model while hasattr(m, "model"): - if hasattr(m, "gradient_checkpointing"): - m.gradient_checkpointing = use_gradient_checkpointing - if hasattr(m, "training"): - m.training = True - # Pad tokenizer to the right - if hasattr(m, "_saved_temp_tokenizer"): - m._saved_temp_tokenizer.padding_side = "right" + _for_training(m) m = m.model - pass - if hasattr(m, "gradient_checkpointing"): - m.gradient_checkpointing = use_gradient_checkpointing - if hasattr(m, "training"): - m.training = True - # Pad tokenizer to the right - if hasattr(m, "_saved_temp_tokenizer"): - m._saved_temp_tokenizer.padding_side = "right" - - # Also revert model.generate - if hasattr(model, "_unwrapped_old_generate"): - model.generate = model._unwrapped_old_generate - del model._unwrapped_old_generate - pass + _for_training(m) # Also re-enable training for embeddings for NEFTune if hasattr(model, "get_input_embeddings"): @@ -2715,7 +2678,6 @@ class FastLlamaModel: embeddings = model.get_output_embeddings() if hasattr(embeddings, "training"): embeddings.training = True pass - return model pass pass diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 3a9d651d11..c9ea922272 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -60,7 +60,7 @@ def PatchRL(FastLanguageModel): def unsloth_unwrap_model_for_generation(model, *args, **kwargs): with unwrap_model_for_generation(model, *args, **kwargs) as unwrapped_model: # Put the model in inference mode. - FastLanguageModel.for_inference(unwrapped_model) + FastLanguageModel.for_inference(model) # We must use .clone for Unsloth since we force inference_mode # Rather we should have used no_grad diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 5ea61cb9b3..7462d55944 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -78,6 +78,25 @@ def sft_trainer_prepare_dataset(function_name, function): if function_name != "_prepare_non_packed_dataloader" and \ function_name != "_prepare_dataset": return function + fast_sft_prepare_dataset = RL_REPLACEMENTS.get("sft_prepare_dataset", None) + if fast_sft_prepare_dataset is not None and "pack_examples" in function: + params = inspect.signature(fast_sft_prepare_dataset).parameters.keys() + params = ".*?".join(params) + matched = re.match( + r"[\s]{0,}def _prepare_dataset\(.*?" + params + r".*?\)", + function, + flags = re.MULTILINE | re.DOTALL, + ) + if matched: + # Use fast version! + function = inspect.getsource(fast_sft_prepare_dataset) + function = function.split("\n") + function = "\n".join(" "*4 + x for x in function) + function = function.replace("def sft_prepare_dataset", "def _prepare_dataset") + return function + pass + pass + check_text = \ "if 'tokenizer' not in locals(): tokenizer = processing_class\n"\ "if 'formatting_func' not in locals(): raise RuntimeError('Unsloth: Please file a bug report - `formatting_func` does not exist!')\n"\