diff --git a/README.md b/README.md index 4bdd7e2893..5b2dd6f129 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ All notebooks are **beginner friendly**! Add your dataset, click "Run All", and - Click [here](https://docs.unsloth.ai/) for detailed documentation for Unsloth. ## 🦥 Unsloth.ai News -- 📣 NEW! Introducing [Reasoning](https://unsloth.ai/blog/r1-reasoning) in Unsloth. You can now reproduce DeepSeek-R1's "aha" moment with just 7GB VRAM. Transform Llama, Phi, Mistral etc. into reasoning LLMs! +- 📣 NEW! Introducing Long-context [Reasoning (GRPO)](https://unsloth.ai/blog/grpo) in Unsloth. You can now reproduce DeepSeek-R1's "aha" moment with just 5GB VRAM. Transform Llama, Phi, Mistral etc. into reasoning LLMs! - 📣 NEW! [DeepSeek-R1](https://unsloth.ai/blog/deepseek-r1) - the most powerful open reasoning models with Llama & Qwen distillations. Run or fine-tune them now! More details: [unsloth.ai/blog/deepseek-r1](https://unsloth.ai/blog/deepseek-r1). All model uploads: [here](https://huggingface.co/collections/unsloth/deepseek-r1-all-versions-678e1c48f5d2fce87892ace5). - 📣 NEW! [Phi-4](https://unsloth.ai/blog/phi4) by Microsoft is now supported. We also [fixed bugs](https://unsloth.ai/blog/phi4) in Phi-4 and [uploaded GGUFs, 4-bit](https://huggingface.co/collections/unsloth/phi-4-all-versions-677eecf93784e61afe762afa). Try the [Phi-4 Colab notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Phi_4-Conversational.ipynb) - 📣 NEW! [Llama 3.3 (70B)](https://huggingface.co/collections/unsloth/llama-33-all-versions-67535d7d994794b9d7cf5e9f), Meta's latest model is supported. @@ -65,9 +65,8 @@ All notebooks are **beginner friendly**! Add your dataset, click "Run All", and | ------------------------------- | --------------------------------------- | | 📚 **Documentation & Wiki** | [Read Our Docs](https://docs.unsloth.ai) | |   **Twitter (aka X)** | [Follow us on X](https://twitter.com/unslothai)| -| 💾 **Installation** | [unsloth/README.md](https://github.com/unslothai/unsloth/tree/main#-installation-instructions)| -| 🥇 **Benchmarking** | [Performance Tables](https://github.com/unslothai/unsloth/tree/main#-performance-benchmarking) -| 🌐 **Released Models** | [Unsloth Releases](https://docs.unsloth.ai/get-started/all-our-models)| +| 💾 **Installation** | [Pip install](https://github.com/unslothai/unsloth/edit/main/README.md#-install-unsloth)| +| 🔮 **Our Models** | [Unsloth Releases](https://docs.unsloth.ai/get-started/all-our-models)| | ✍️ **Blog** | [Read our Blogs](https://unsloth.ai/blog)| |   **Reddit** | [Join our Reddit page](https://reddit.com/r/unsloth)| @@ -77,36 +76,15 @@ All notebooks are **beginner friendly**! Add your dataset, click "Run All", and - No change of hardware. Supports NVIDIA GPUs since 2018+. Minimum CUDA Capability 7.0 (V100, T4, Titan V, RTX 20, 30, 40x, A100, H100, L40 etc) [Check your GPU!](https://developer.nvidia.com/cuda-gpus) GTX 1070, 1080 works, but is slow. - Works on **Linux** and **Windows** via WSL. - Supports 4bit and 16bit QLoRA / LoRA finetuning via [bitsandbytes](https://github.com/TimDettmers/bitsandbytes). -- Open source trains 5x faster - see [Unsloth Pro](https://unsloth.ai/) for up to **30x faster training**! - If you trained a model with 🦥Unsloth, you can use this cool sticker!   +## 💾 Install Unsloth -## 🥇 Performance Benchmarking -- For our most detailed benchmarks, read our [Llama 3.3 Blog](https://unsloth.ai/blog/llama3-3). -- Benchmarking of Unsloth was also conducted by [🤗Hugging Face](https://huggingface.co/blog/unsloth-trl). - -We tested using the Alpaca Dataset, a batch size of 2, gradient accumulation steps of 4, rank = 32, and applied QLoRA on all linear layers (q, k, v, o, gate, up, down): - -| Model | VRAM | 🦥 Unsloth speed | 🦥 VRAM reduction | 🦥 Longer context | 😊 Hugging Face + FA2 | -|----------------|-------|-----------------|----------------|----------------|--------------------| -| Llama 3.3 (70B)| 80GB | 2x | >75% | 13x longer | 1x | -| Llama 3.1 (8B) | 80GB | 2x | >70% | 12x longer | 1x | - -
- -![](https://i.ibb.co/sJ7RhGG/image-41.png) - -## 💾 Installation Instructions - -Simply use pip install on Linux machines. Windows instructions are below. -
- -
- +- **Install with pip (recommended)** for Linux devices: +``` +pip install unsloth +``` +See below for Windows install instructions: ### Conda Installation (Optional) `⚠️Only use Conda if you have it. If not, use Pip`. Select either `pytorch-cuda=11.8,12.1` for CUDA 11.8 or CUDA 12.1. We support `python=3.10,3.11,3.12`. ```bash @@ -190,8 +168,50 @@ x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "") print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"') ``` -### Windows Installation +## Windows Installation +> [!warning] +> Python 3.13 does not support Unsloth. Use 3.12, 3.11 or 3.10 +### Step 1: NVIDIA Video Driver +You should install the latest version of your GPUs driver. You can download drivers here: + - [NVIDIA GPU Drive Download](https://www.nvidia.com/Download/index.aspx) + +### Step 2: Visual Studio C++ +You will need Visual Studio, with C++ installed. By default, C++ is not installed with Visual Studio, so make sure you select all of the C++ options. Also select options for Windows 10/11 SDK. + - [Visual Studio Community Edition](https://visualstudio.microsoft.com/vs/community/) + + + + + +
+ VSCode C++ Ref Image + +
+

Steps to configure VS C++

+
+
    +
  1. Launch the Installer downloaded from the link above.
    2. +
  2. In the installer, navigate to Individual components and select all the options mentioned in the image.
    3. +
  3. Click on install now.
    4. +
+
+ +### Step 3: CUDA Toolkit + + - [Download CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit-archive) + +### Step 4: Install PyTorch + +You will need the correct version of PyTorch that is compatibile with your CUDA drivers, so make sure to select them carefully + - [Install PyTorch](https://pytorch.org/get-started/locally/) + +### Step 5: Install Unsloth +```python +pip install "unsloth[windows] @ git+https://github.com/unslothai/unsloth.git" +``` + +### Side note To run Unsloth directly on Windows: - Install Triton from this Windows fork and follow the instructions: https://github.com/woct0rdho/triton-windows (be aware that the Windows fork requires PyTorch >= 2.4 and CUDA 12) - In the SFTTrainer, set `dataset_num_proc=1` to avoid a crashing issue: @@ -297,12 +317,18 @@ trainer.train() # (4) Customized chat templates ``` - -## DPO Support -DPO (Direct Preference Optimization), PPO, Reward Modelling all seem to work as per 3rd party independent testing from [Llama-Factory](https://github.com/hiyouga/LLaMA-Factory). We have a preliminary Google Colab notebook for reproducing Zephyr on Tesla T4 here: [notebook](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing). + +## 💡 Reinforcement Learning +RL including DPO, GRPO, PPO, Reward Modelling, Online DPO all work with Unsloth. We're in 🤗Hugging Face's official docs! We're on the [SFT docs](https://huggingface.co/docs/trl/main/en/sft_trainer#accelerate-fine-tuning-2x-using-unsloth) and the [DPO docs](https://huggingface.co/docs/trl/main/en/dpo_trainer#accelerate-dpo-fine-tuning-using-unsloth)! List of RL notebooks: -We're in 🤗Hugging Face's official docs! We're on the [SFT docs](https://huggingface.co/docs/trl/main/en/sft_trainer#accelerate-fine-tuning-2x-using-unsloth) and the [DPO docs](https://huggingface.co/docs/trl/main/en/dpo_trainer#accelerate-dpo-fine-tuning-using-unsloth)! +- ORPO notebook: [Link](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3_(8B)-ORPO.ipynb) +- DPO Zephyr notebook: [Link](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Zephyr_(7B)-DPO.ipynb) +- KTO notebook: [Link](https://colab.research.google.com/drive/1a2b3c4d5e6f7g8h9i0j) +- SimPO notebook: [Link](https://colab.research.google.com/drive/1a2b3c4d5e6f7g8h9i0j) +
+ Click for DPO code + ```python import os os.environ["CUDA_VISIBLE_DEVICES"] = "0" # Optional set GPU device ID @@ -360,9 +386,21 @@ dpo_trainer = DPOTrainer( ) dpo_trainer.train() ``` +
+ +## 🥇 Performance Benchmarking +- For our most detailed benchmarks, read our [Llama 3.3 Blog](https://unsloth.ai/blog/llama3-3). +- Benchmarking of Unsloth was also conducted by [🤗Hugging Face](https://huggingface.co/blog/unsloth-trl). + +We tested using the Alpaca Dataset, a batch size of 2, gradient accumulation steps of 4, rank = 32, and applied QLoRA on all linear layers (q, k, v, o, gate, up, down): + +| Model | VRAM | 🦥 Unsloth speed | 🦥 VRAM reduction | 🦥 Longer context | 😊 Hugging Face + FA2 | +|----------------|-------|-----------------|----------------|----------------|--------------------| +| Llama 3.3 (70B)| 80GB | 2x | >75% | 13x longer | 1x | +| Llama 3.1 (8B) | 80GB | 2x | >70% | 12x longer | 1x | -## 🥇 Detailed Benchmarking Tables ### Context length benchmarks + #### Llama 3.1 (8B) max. context length We tested Llama 3.1 (8B) Instruct and did 4bit QLoRA on all linear layers (Q, K, V, O, gate, up and down) with rank = 32 with a batch size of 1. We padded all sequences to a certain maximum sequence length to mimic long context finetuning workloads. | GPU VRAM | 🦥Unsloth context length | Hugging Face + FA2 | diff --git a/pyproject.toml b/pyproject.toml index 14797c8fa7..de1583e9e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,10 +33,32 @@ exclude = ["images*"] [project.optional-dependencies] triton = [ - "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp39-cp39-win_amd64.whl ; python_version=='3.9' and platform_system == 'Windows'", - "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp310-cp310-win_amd64.whl ; python_version=='3.10' and platform_system == 'Windows'", - "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp311-cp311-win_amd64.whl ; python_version=='3.11' and platform_system == 'Windows'", - "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'", + "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.2.0-windows.post10/triton-3.2.0-cp39-cp39-win_amd64.whl ; python_version=='3.9' and platform_system == 'Windows'", + "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.2.0-windows.post10/triton-3.2.0-cp310-cp310-win_amd64.whl ; python_version=='3.10' and platform_system == 'Windows'", + "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.2.0-windows.post10/triton-3.2.0-cp311-cp311-win_amd64.whl ; python_version=='3.11' and platform_system == 'Windows'", + "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.2.7", + "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]", + "bitsandbytes>=0.41.1 ; platform_system == 'Windows'", + "xformers>=0.0.22.post7 ; platform_system == 'Windows'", ] huggingface = [ "unsloth_zoo>=2025.2.7", diff --git a/unsloth/__init__.py b/unsloth/__init__.py index a3b3e68b2d..e33d16577a 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -17,6 +17,27 @@ from packaging.version import Version import os, re, subprocess, inspect import numpy as np +# Check if modules that need patching are already imported +critical_modules = ['trl', 'transformers', 'peft'] +already_imported = [mod for mod in critical_modules if mod in sys.modules] + +# This check is critical because Unsloth optimizes these libraries by modifying +# their code at import time. If they're imported first, the original (slower, +# more memory-intensive) implementations will be used instead of Unsloth's +# optimized versions, potentially causing OOM errors or slower training. + +if already_imported: + # stacklevel=2 makes warning point to user's import line rather than this library code, + # showing them exactly where to fix the import order in their script + warnings.warn( + f"WARNING: Unsloth should be imported before {', '.join(already_imported)} " + f"to ensure all optimizations are applied. Your code may run slower or encounter " + f"memory issues without these optimizations.\n\n" + f"Please restructure your imports with 'import unsloth' at the top of your file.", + stacklevel = 2, + ) +pass + # Unsloth currently does not work on multi GPU setups - sadly we are a 2 brother team so # enabling it will require much more work, so we have to prioritize. Please understand! # We do have a beta version, which you can contact us about! @@ -25,25 +46,6 @@ import numpy as np # Fixes https://github.com/unslothai/unsloth/issues/1266 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" -if "CUDA_VISIBLE_DEVICES" in os.environ: - os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" - devices = os.environ["CUDA_VISIBLE_DEVICES"] - # Check if there are multiple cuda devices set in env - if not devices.isdigit(): - first_id = devices.split(",")[0] - warnings.warn( - f"Unsloth: 'CUDA_VISIBLE_DEVICES' is currently {devices} \n"\ - "Unsloth currently does not support multi GPU setups - but we are working on it!\n"\ - "Multiple CUDA devices detected but we require a single device.\n"\ - f"We will override CUDA_VISIBLE_DEVICES to first device: {first_id}." - ) - os.environ["CUDA_VISIBLE_DEVICES"] = str(first_id) -else: - # warnings.warn("Unsloth: 'CUDA_VISIBLE_DEVICES' is not set. We shall set it ourselves.") - os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" - os.environ["CUDA_VISIBLE_DEVICES"] = "0" -pass - # Reduce VRAM usage by reducing fragmentation # And optimize pinning of memory os.environ["PYTORCH_CUDA_ALLOC_CONF"] = \ diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index c401393234..5785894a23 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -1684,7 +1684,7 @@ extra_eos_tokens = None, for j in range(1, len(response_part)): try_find = re.escape(response_part[:j]) - try: found = next(re.finditer("(" + try_find + ").+?\{INPUT\}", chat_template, flags = re.DOTALL | re.MULTILINE)) + try: found = next(re.finditer("(" + try_find + ").+?\\{INPUT\\}", chat_template, flags = re.DOTALL | re.MULTILINE)) except: break pass separator = found.group(1) @@ -2125,7 +2125,7 @@ def test_hf_gguf_equivalence(tokenizer, gguf_model = "./model-unsloth.F16.gguf") gguf_tokens = "".join(datas) # Now extract GGUF tokenization attempt - gguf_tokenized = re.findall("([\d]{1,}) \-\> \'([^\']{1,})\'", gguf_tokens, flags = re.MULTILINE) + gguf_tokenized = re.findall(r"([\d]{1,}) \-\> \'([^\']{1,})\'", gguf_tokens, flags = re.MULTILINE) gguf_tokenized = [(int(x[0]), x[1],) for x in gguf_tokenized] input_ids = tokenizer(prompt).input_ids diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index c2b7929a29..a4fb2a89b6 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -98,9 +98,6 @@ class LoRA_MLP(torch.autograd.Function): gateA, gateB, upA, upB, downA, downB, \ X, e, g = ctx.saved_tensors - gateA, gateB, upA, upB, downA, downB = \ - gateA.t(), gateB.t(), upA.t(), upB.t(), downA.t(), downB.t() - batch, seq_len, hd = X.shape dY = dY.view(-1, dY.shape[-1]) X = X .view(-1, X .shape[-1]) @@ -108,39 +105,61 @@ class LoRA_MLP(torch.autograd.Function): g = g .view(-1, g .shape[-1]) dtype = X.dtype + gateA, gateB, upA, upB, downA, downB = \ + gateA.to(dtype), gateB.to(dtype), upA.to(dtype), upB.to(dtype), downA.to(dtype), downB.to(dtype) + + gateA, gateB, upA, upB, downA, downB = \ + gateA.t(), gateB.t(), upA.t(), upB.t(), downA.t(), downB.t() + DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) DW, e, g = _backward_function(DW, e, g) h, df, de = DW, e, g + d_downA = torch.empty_like(downA) + d_downB = torch.empty_like(downB) + d_gateA = torch.empty_like(gateA) + d_gateB = torch.empty_like(gateB) + d_upA = torch.empty_like(upA) + d_upB = torch.empty_like(upB) + # Down projection LoRA weights - d_downA = h.t() @ (dY @ downB.t()) - d_downB = (downA.t() @ h.t()) @ dY - d_downA *= downS - d_downB *= downS + # d_downA = h.t() @ (dY @ downB.t()) + # d_downB = (downA.t() @ h.t()) @ dY + # d_downA *= downS + # d_downB *= downS + d_downA.addmm_(h.t(), dY @ downB.t(), alpha = downS, beta = 0) + d_downB.addmm_(downA.t() @ h.t(), dY, alpha = downS, beta = 0) # Up projection LoRA weights - d_upA = X.t() @ (df @ upB.t()) - d_upB = (upA.t() @ X.t()) @ df - d_upA *= upS - d_upB *= upS + # d_upA = X.t() @ (df @ upB.t()) + # d_upB = (upA.t() @ X.t()) @ df + # d_upA *= upS + # d_upB *= upS + d_upA.addmm_(X.t(), df @ upB.t(), alpha = upS, beta = 0) + d_upB.addmm_(upA.t() @ X.t(), df, alpha = upS, beta = 0) # Gate projection LoRA weights - d_gateA = X.t() @ (de @ gateB.t()) - d_gateB = (gateA.t() @ X.t()) @ de - d_gateA *= gateS - d_gateB *= gateS + # d_gateA = X.t() @ (de @ gateB.t()) + # d_gateB = (gateA.t() @ X.t()) @ de + # d_gateA *= gateS + # d_gateB *= gateS + d_gateA.addmm_(X.t(), de @ gateB.t(), alpha = gateS, beta = 0) + d_gateB.addmm_(gateA.t() @ X.t(), de, alpha = gateS, beta = 0) # dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS) # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS) upW = fast_dequantize(upW.t(), upW_quant) dX = torch.matmul(df, upW.t(), out = X if ctx.inplace else None) del upW - dX += df @ upB.to(dtype).t() @ (upS * upA.to(dtype).t()) + # dX += df @ upB.to(dtype).t() @ (upS * upA.to(dtype).t()) + dX.addmm_(df @ upB.t(), upA.t(), alpha = upS) gateW = fast_dequantize(gateW.t(), gateW_quant) - dX += de @ gateW.t() + # dX += de @ gateW.t() + dX.addmm_(de, gateW.t()) del gateW - dX += de @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t()) + # dX += de @ gateB.to(dtype).t() @ (gateS * gateA.to(dtype).t()) + dX.addmm_(de @ gateB.t(), gateA.t(), alpha = gateS) # gateW, gateW_quant, gateA, gateB, gateS, # upW, upW_quant, upA, upB, upS, @@ -258,9 +277,6 @@ class LoRA_QKV(torch.autograd.Function): ctx.custom_saved_tensors X, QA, QB, KA, KB, VA, VB, = ctx.saved_tensors - QA, QB, KA, KB, VA, VB = \ - QA.t(), QB.t(), KA.t(), KB.t(), VA.t(), VB.t() - batch, seq_len, hd = X.shape dQ = dQ.view(-1, dQ.shape[-1]) dK = dK.reshape(-1, dK.shape[-1]) # view doesn't work on K.T @@ -268,45 +284,68 @@ class LoRA_QKV(torch.autograd.Function): X = X .view(-1, X .shape[-1]) dtype = X.dtype + QA, QB, KA, KB, VA, VB = \ + QA.to(dtype), QB.to(dtype), KA.to(dtype), KB.to(dtype), VA.to(dtype), VB.to(dtype) + + QA, QB, KA, KB, VA, VB = \ + QA.t(), QB.t(), KA.t(), KB.t(), VA.t(), VB.t() + ### Weight projection LoRA weights # See our blogpost for more details. + d_QA = torch.empty_like(QA) + d_QB = torch.empty_like(QB) + d_KA = torch.empty_like(KA) + d_KB = torch.empty_like(KB) + d_VA = torch.empty_like(VA) + d_VB = torch.empty_like(VB) # Q Projection - d_QA = X.t() @ (dQ @ QB.t()) - d_QB = (QA.t() @ X.t()) @ dQ - d_QA *= QS - d_QB *= QS + # d_QA = X.t() @ (dQ @ QB.t()) + # d_QB = (QA.t() @ X.t()) @ dQ + # d_QA *= QS + # d_QB *= QS + d_QA.addmm_(X.t(), dQ @ QB.t(), alpha = QS, beta = 0) + d_QB.addmm_(QA.t() @ X.t(), dQ, alpha = QS, beta = 0) # K Projection - d_KA = X.t() @ (dK @ KB.t()) - d_KB = (KA.t() @ X.t()) @ dK - d_KA *= KS - d_KB *= KS + # d_KA = X.t() @ (dK @ KB.t()) + # d_KB = (KA.t() @ X.t()) @ dK + # d_KA *= KS + # d_KB *= KS + d_KA.addmm_(X.t(), dK @ KB.t(), alpha = KS, beta = 0) + d_KB.addmm_(KA.t() @ X.t(), dK, alpha = KS, beta = 0) # V Projection - d_VA = X.t() @ (dV @ VB.t()) - d_VB = (VA.t() @ X.t()) @ dV - d_VA *= VS - d_VB *= VS + # d_VA = X.t() @ (dV @ VB.t()) + # d_VB = (VA.t() @ X.t()) @ dV + # d_VA *= VS + # d_VB *= VS + d_VA.addmm_(X.t(), dV @ VB.t(), alpha = VS, beta = 0) + d_VB.addmm_(VA.t() @ X.t(), dV, alpha = VS, beta = 0) # Combine derivatives to find dX # dQ QW = fast_dequantize(QW.t(), QW_quant) dX = torch.matmul(dQ, QW.t(), out = X if ctx.inplace else None) del QW - dX += (dQ @ QB.to(dtype).t() @ (QS * QA.to(dtype).t())) + # dX += (dQ @ QB.to(dtype).t() @ (QS * QA.to(dtype).t())) + dX.addmm_(dQ @ QB.t(), QA.t(), alpha = QS) # dK KW = fast_dequantize(KW.t(), KW_quant) - dX += dK @ KW.t() + # dX += dK @ KW.t() + dX.addmm_(dK, KW.t()) del KW - dX += dK @ KB.to(dtype).t() @ (KS * KA.to(dtype).t()) + # dX += dK @ KB.to(dtype).t() @ (KS * KA.to(dtype).t()) + dX.addmm_(dK @ KB.t(), KA.t(), alpha = KS) # dV VW = fast_dequantize(VW.t(), VW_quant) - dX += dV @ VW.t() + # dX += dV @ VW.t() + dX.addmm_(dV, VW.t()) del VW - dX += dV @ VB.to(dtype).t() @ (VS * VA.to(dtype).t()) + # dX += dV @ VB.to(dtype).t() @ (VS * VA.to(dtype).t()) + dX.addmm_(dV @ VB.t(), VA.t(), alpha = VS) # QW, QW_quant, QA, QB, QS, # KW, KW_quant, KA, KB, KS, @@ -378,25 +417,33 @@ class LoRA_W(torch.autograd.Function): W, W_quant, S = ctx.custom_saved_tensors A, B, X = ctx.saved_tensors - A, B = A.t(), B.t() - batch, seq_len, hd = X.shape dY = dY.reshape(-1, dY.shape[-1]) # Must be reshape X = X .reshape(-1, X .shape[-1]) # Must be reshape dtype = X.dtype + A, B = A.to(dtype), B.to(dtype) + + A, B = A.t(), B.t() + + d_A = torch.empty_like(A) + d_B = torch.empty_like(B) + ### Weight projection LoRA weights # Weight projection - d_A = X.t() @ (dY @ B.t()) - d_B = (A.t() @ X.t()) @ dY - d_A *= S - d_B *= S + # d_A = X.t() @ (dY @ B.t()) + # d_B = (A.t() @ X.t()) @ dY + # d_A *= S + # d_B *= S + d_A.addmm_(X.t(), dY @ B.t(), alpha = S, beta = 0) + d_B.addmm_(A.t() @ X.t(), dY, alpha = S, beta = 0) # Get derivative for dX W = fast_dequantize(W.t(), W_quant) dX = dY @ W.t() del W - dX += dY @ B.to(dtype).t() @ (S * A.to(dtype).t()) + # dX += dY @ B.to(dtype).t() @ (S * A.to(dtype).t()) + dX.addmm_(dY @ B.t(), A.t(), alpha = S) # W, W_quant, A, B, S return dX.view(batch, seq_len, hd), \ diff --git a/unsloth/kernels/layernorm.py b/unsloth/kernels/layernorm.py index a5f7926e2e..ffcc5cc13c 100644 --- a/unsloth/kernels/layernorm.py +++ b/unsloth/kernels/layernorm.py @@ -49,7 +49,8 @@ def layernorm_forward( b_row = tl.load(b + col_offsets, mask = mask, other = 0).to(tl.float32) mean_X = tl.sum(X_row, axis = 0) / n_cols - XX = X_row - mean_X + # (X[0] - mean) == -mean so we need to mask it out + XX = tl.where(mask, X_row - mean_X, 0) row_var = tl.sum(XX * XX, axis = 0) / n_cols inv_var = tl.math.rsqrt(row_var + eps) tl.store (r, inv_var) @@ -105,10 +106,10 @@ class Fast_Layernorm(torch.autograd.Function): X = X.view(-1, dim) n_rows, n_cols = X.shape BLOCK_SIZE, num_warps = calculate_settings(n_cols) - - Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0") - r = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") - mu = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") + device = X.device + Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = device) + r = torch.empty(n_rows, dtype = torch.float32, device = device) + mu = torch.empty(n_rows, dtype = torch.float32, device = device) layernorm_forward[(n_rows,)]( Y, Y.stride(0), diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index 6310f7f392..7487c10eeb 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -148,9 +148,10 @@ class Fast_RMS_Layernorm(torch.autograd.Function): BLOCK_SIZE : int num_warps : int BLOCK_SIZE, num_warps = calculate_settings(n_cols) + device = X.device - Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0") - r = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") + Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = device) + r = torch.empty(n_rows, dtype = torch.float32, device = device) fx = _gemma_rms_layernorm_forward if gemma else _rms_layernorm_forward fx[(n_rows,)]( @@ -180,7 +181,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function): n_cols : int n_rows, n_cols = dY.shape # dW = X - dX = torch.empty_like(dY, device = "cuda:0") if ctx.GEMMA else dY + dX = torch.empty_like(dY) if ctx.GEMMA else dY _rms_layernorm_backward[(n_rows,)]( dY, dY.stride(0), diff --git a/unsloth/kernels/swiglu.py b/unsloth/kernels/swiglu.py index f81b7aae9b..688e9f9a48 100644 --- a/unsloth/kernels/swiglu.py +++ b/unsloth/kernels/swiglu.py @@ -41,7 +41,7 @@ pass def swiglu_fg_kernel(e, g): batch, seq_len, hd = e.shape n_elements = e.numel() - h = torch.empty((batch, seq_len, hd), dtype = e.dtype, device = "cuda:0") + h = torch.empty((batch, seq_len, hd), dtype = e.dtype, device = e.device) grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) _fg_kernel[grid](e, g, h, n_elements, BLOCK_SIZE = 1024,) return h diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index f052914f98..985adaaa44 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -61,12 +61,29 @@ pass import bitsandbytes as bnb +import ctypes + # https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1330/files HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3") -global CUDA_STREAM -CUDA_STREAM = None get_ptr = bnb.functional.get_ptr -import ctypes + +# Get array of CUDA streams and other buffers +global CUDA_STREAMS +global WEIGHT_BUFFERS +global ABSMAX_BUFFERS + +_CUDA_STREAMS = { + (index := torch.cuda.device(i).idx) : ctypes.c_void_p(torch._C._cuda_getCurrentRawStream(index)) + for i in range(torch.cuda.device_count()) +} +CUDA_STREAMS = [None] * (max(_CUDA_STREAMS.keys()) + 1) +WEIGHT_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1) +ABSMAX_BUFFERS = [None] * (max(_CUDA_STREAMS.keys()) + 1) +for k, v in _CUDA_STREAMS.items(): CUDA_STREAMS[k] = v +CUDA_STREAMS = tuple(CUDA_STREAMS) +del _CUDA_STREAMS + +# Bitsandbytes operations ctypes_c_int = ctypes.c_int ctypes_c_int32 = ctypes.c_int32 cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32 @@ -118,11 +135,6 @@ def get_lora_parameters_bias(proj): return W, QUANT_STATE(W), A, B, s, bias pass -global WEIGHT_BUFFER -WEIGHT_BUFFER = None -global ABSMAX_BUFFER -ABSMAX_BUFFER = None - if HAS_CUDA_STREAM: @torch.inference_mode def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False): @@ -145,8 +157,10 @@ if HAS_CUDA_STREAM: offset, state2 = compressed_stats absmax2, code2, blocksize2, _, _, _, _ = state2 pass - global CUDA_STREAM - if CUDA_STREAM is None: CUDA_STREAM = torch.cuda.current_stream("cuda:0") + global CUDA_STREAMS + device = W.device + device_index = device.index + CUDA_STREAM = CUDA_STREAMS[device_index] n_elements_absmax = absmax.numel() @@ -155,11 +169,13 @@ if HAS_CUDA_STREAM: # Use same buffers for faster inference size = shape[0]*shape[1] - global WEIGHT_BUFFER - global ABSMAX_BUFFER + global WEIGHT_BUFFERS + global ABSMAX_BUFFERS + WEIGHT_BUFFER = WEIGHT_BUFFERS[device_index] + ABSMAX_BUFFER = ABSMAX_BUFFERS[device_index] if WEIGHT_BUFFER is None: - WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = "cuda:0", requires_grad = False) - ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0", 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) @@ -168,11 +184,11 @@ if HAS_CUDA_STREAM: out_absmax = ABSMAX_BUFFER[:n_elements_absmax] else: if out is None: - out = torch.empty(shape, dtype = dtype, device = "cuda:0", 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 = "cuda:0", requires_grad = False) + out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) pass # NF4 dequantization of statistics @@ -217,31 +233,15 @@ else: pass n_elements_absmax = absmax.numel() + device = W.device # Create weight matrix - if use_global_buffer: - - # Use same buffers for faster inference - size = shape[0]*shape[1] - global WEIGHT_BUFFER - global ABSMAX_BUFFER - if WEIGHT_BUFFER is None: - WEIGHT_BUFFER = torch.empty(size, dtype = dtype, device = "cuda:0", requires_grad = False) - ABSMAX_BUFFER = torch.empty(n_elements_absmax, dtype = dtype, device = "cuda:0", 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) - - out = WEIGHT_BUFFER[:size].view(shape) - out_absmax = ABSMAX_BUFFER[:n_elements_absmax] + if out is None: + out = torch.empty(shape, dtype = dtype, device = device, requires_grad = False) else: - if out is None: - out = torch.empty(shape, dtype = dtype, device = "cuda:0", requires_grad = False) - else: - assert(out.shape == shape) - assert(out.dtype == dtype) - out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = "cuda:0", requires_grad = False) - pass + assert(out.shape == shape) + assert(out.dtype == dtype) + out_absmax = torch.empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) # Do dequantization ptr_out_absmax = get_ptr(out_absmax) @@ -288,14 +288,16 @@ if HAS_CUDA_STREAM: offset, state2 = compressed_stats absmax2, code2, blocksize2, _, _, _, _ = state2 pass - global CUDA_STREAM - if CUDA_STREAM is None: CUDA_STREAM = torch.cuda.current_stream("cuda:0") + global CUDA_STREAMS + device = W.device + device_index = device.index + CUDA_STREAM = CUDA_STREAMS[device_index] # assert(dtype == X.dtype) bout = shape[0] if out is None: - out = torch.empty((1, 1, bout,), dtype = dtype, device = "cuda:0") + out = torch.empty((1, 1, bout,), dtype = dtype, device = device) # else: # assert(out.shape == (1, 1, bout,)) # pass @@ -313,7 +315,7 @@ if HAS_CUDA_STREAM: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch.empty(absmax.shape, dtype = torch.float32, device = "cuda:0") + 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()), CUDA_STREAM, @@ -357,9 +359,10 @@ else: pass # assert(dtype == X.dtype) bout = shape[0] + device = W.device if out is None: - out = torch.empty((1, 1, bout,), dtype = dtype, device = "cuda:0") + out = torch.empty((1, 1, bout,), dtype = dtype, device = device) # else: # assert(out.shape == (1, 1, bout,)) # pass @@ -377,7 +380,7 @@ else: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch.empty(absmax.shape, dtype = torch.float32, device = "cuda:0") + 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()), @@ -400,6 +403,7 @@ 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) @@ -461,7 +465,9 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None): if A is not None: # LoRA is enabled A, B = A.t(), B.t() - out += (X @ A.to(dtype)) @ (s * B.to(dtype)) + XA = torch_matmul(X, A.to(dtype)) + out.addmm_(XA, B.to(dtype), alpha = s) + # out += (X @ A.to(dtype)) @ (s * B.to(dtype)) pass return out.view(batch, seq_len, -1) if reshape else out diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index e1259af3ae..cca77bb60b 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.2.14" +__version__ = "2025.3.1" __all__ = [ "SUPPORTS_BFLOAT16", @@ -25,7 +25,6 @@ __all__ = [ "__version__", "HAS_FLASH_ATTENTION", "HAS_FLASH_ATTENTION_SOFTCAPPING", - "PRE_CHECK", "platform_system", "patch_tokenizer", "get_statistics", @@ -37,7 +36,6 @@ __all__ = [ "torch_compile_options", "patch_linear_scaling", "patch_llama_rope_scaling", - "check_nvidia", "create_boolean_mask", "torch_amp_custom_fwd", "torch_amp_custom_bwd", @@ -589,7 +587,7 @@ if Version(peft_version) < Version("0.12.0"): spaces = len(re.match(r"[\s]{1,}", source).group(0)) lines = source.split("\n") source = "\n".join(x[spaces:] for x in lines) - source = re.sub("([^\.])nn\.", r"\1torch.nn.", source) + source = re.sub(r"([^\.])nn\.", r"\1torch.nn.", source) source = source.replace("def update_layer", "def LoraLayer_update_layer") exec(source, globals()) @@ -703,9 +701,7 @@ pass # ============================================= # Fixes Bitsandbytes to remove missing warnings from transformers.utils.quantization_config import BitsAndBytesConfig, QuantizationMethod -from inspect import getsource -from accelerate.utils.dataclasses import DistributedType -BitsAndBytesConfig__init__ = getsource(BitsAndBytesConfig.__init__) +BitsAndBytesConfig__init__ = inspect.getsource(BitsAndBytesConfig.__init__) BitsAndBytesConfig__init__ = re.sub( r"if[\s]{1,}kwargs\:[\s]{1,}.+?\n", "", @@ -719,28 +715,30 @@ BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.replace( "__init__", "_BitsAndBytesConfig__init__", ) - -def _prepare_backend( - self, cpu = False, sagemaker_dp = False, backend: str = None, -) -> tuple[str, DistributedType]: - return None, DistributedType.NO -pass -import accelerate.state -accelerate.state.PartialState._prepare_backend = _prepare_backend - -import accelerate.accelerator -prepare = inspect.getsource(accelerate.accelerator.Accelerator.prepare) -prepare = prepare.split("\n") -spaces = prepare[0].find("def") -prepare = "\n".join(x[spaces:] for x in prepare) -x = "for obj in args:" -s = " "*spaces -prepare = prepare.replace(x, f'self.state.distributed_type = DistributedType.NO\n{s}{x}', 1) -exec(prepare, globals()) -accelerate.accelerator.Accelerator.prepare = prepare - exec(BitsAndBytesConfig__init__, globals()) +if torch.cuda.device_count() == 1: + from accelerate.utils.dataclasses import DistributedType + def _prepare_backend( + self, cpu = False, sagemaker_dp = False, backend: str = None, + ) -> tuple[str, DistributedType]: + return None, DistributedType.NO + pass + import accelerate.state + accelerate.state.PartialState._prepare_backend = _prepare_backend + + import accelerate.accelerator + prepare = inspect.getsource(accelerate.accelerator.Accelerator.prepare) + prepare = prepare.split("\n") + spaces = prepare[0].find("def") + prepare = "\n".join(x[spaces:] for x in prepare) + x = "for obj in args:" + s = " "*spaces + prepare = prepare.replace(x, f'self.state.distributed_type = DistributedType.NO\n{s}{x}', 1) + exec(prepare, globals()) + accelerate.accelerator.Accelerator.prepare = prepare +pass + import transformers.utils.quantization_config transformers.utils.quantization_config.BitsAndBytesConfig.__init__ = _BitsAndBytesConfig__init__ # ============================================= @@ -852,7 +850,7 @@ def patch_linear_scaling( scaled_rope_function = scaled_rope_module.__name__, ) rotary_emb = re.findall( - "self.rotary_emb = .+?\)", function, + r"self\.rotary\_emb \= .+?\)", function, flags = re.DOTALL | re.MULTILINE, ) if len(rotary_emb) == 0: @@ -952,7 +950,7 @@ def patch_llama_rope_scaling( (longrope_module if longrope_module is not None else rope_module).__name__ ) rotary_emb = re.findall( - "self.rotary_emb = .+?\)", function, + r"self\.rotary\_emb \= .+?\)", function, flags = re.DOTALL | re.MULTILINE, ) if len(rotary_emb) == 0: return None, function @@ -963,21 +961,6 @@ def patch_llama_rope_scaling( pass -def check_nvidia(): - # Unsloth doesn't work yet on AMD devices - we're working on it! - output = np.array([0,]) - try: - output = subprocess.check_output("nvidia-smi --query-gpu=memory.used --format=csv", shell = True) - output = re.findall(rb'([\d]{1,})[\s]{1,}M', output) - output = np.array([int(x.decode('utf-8'))/1024 for x in output]) - except: - if not torch.cuda.is_available(): - raise RuntimeError("Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!") - return output -pass -PRE_CHECK = check_nvidia() - - def create_boolean_mask(n = 4096, sliding_window = 2048): # Creates a boolean mask for attention mask = torch.ones(n, n, dtype = torch.bool) @@ -1122,8 +1105,6 @@ def patch_gradient_accumulation_fix(Trainer): items_in_trainer = dir(transformers.trainer) good_items = [] for item in items_in_trainer: - # TODO: Support Deepspeed - if item.startswith(("deepspeed", "xm", "met", "smp")): continue if item in function: good_items.append(item) pass exec("from transformers.trainer import (" + ", ".join(x for x in good_items) + ")", globals()) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index bc29c46abc..873bdcf2eb 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -245,8 +245,8 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module): emb = torch.cat((radians_new, radians_new), dim = -1) # We must do RoPE in float32! - cos = emb.cos().to(device = "cuda:0", non_blocking = True)#, dtype = dtype) - sin = emb.sin().to(device = "cuda:0", non_blocking = True)#, dtype = dtype) + cos = emb.cos().to(device = "cuda", non_blocking = True)#, dtype = dtype) + sin = emb.sin().to(device = "cuda", non_blocking = True)#, dtype = dtype) self.register_buffer("cos_cached", cos, persistent = False) self.register_buffer("sin_cached", sin, persistent = False) pass @@ -270,7 +270,7 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module): if seq_len <= self.current_rope_size: return # Iteratively grow by increments of 8192 self.current_rope_size = math.ceil(seq_len / 8192) * 8192 - self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype) + self._set_cos_sin_cache(self.current_rope_size, device = "cuda", dtype = x.dtype) pass pass @@ -304,8 +304,8 @@ class GemmaFixedLinearScalingRotaryEmbedding(GemmaFixedRotaryEmbedding): emb = torch.cat((radians_new, radians_new), dim = -1) # We must do RoPE in float32! - cos = emb.cos().to(device = "cuda:0", non_blocking = True)#, dtype = dtype) - sin = emb.sin().to(device = "cuda:0", non_blocking = True)#, dtype = dtype) + cos = emb.cos().to(device = "cuda", non_blocking = True)#, dtype = dtype) + sin = emb.sin().to(device = "cuda", non_blocking = True)#, dtype = dtype) self.register_buffer("cos_cached", cos, persistent = False) self.register_buffer("sin_cached", sin, persistent = False) pass diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index be6b0469d9..316b4e8f0e 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -265,21 +265,22 @@ def Gemma2Attention_fast_forward_inference( attention_size = n_heads*head_dim seq_len = K1.shape[-2] kv_seq_len = seq_len + 1 + device = hidden_states.device # Prefill phase # if not hasattr(self, "paged_attention"): if do_prefill: - self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda:0") + self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = device) self.paged_attention_K = self.paged_attention[:,0] self.paged_attention_V = self.paged_attention[:,1] self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3) - self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = "cuda:0") - self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda:0") - self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda:0") + self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = device) + self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = device) + self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device) # Only for Gemma2 - self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = "cuda:0") - self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda:0") + self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device) + self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = device) # See https://github.com/google/gemma_pytorch/commit/03e657582d17cb5a8617ebf333c1c16f3694670e # Gemma 9b should use 256 and not 224 (hs / nah). 27b uses the below diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index fb7e96d8d2..df498d18ba 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -274,21 +274,22 @@ def GraniteAttention_fast_forward_inference( attention_size = n_heads*head_dim seq_len = K1.shape[-2] kv_seq_len = seq_len + 1 + device = hidden_states.device # Prefill phase # if not hasattr(self, "paged_attention"): if do_prefill: - self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda:0") + self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = device) self.paged_attention_K = self.paged_attention[:,0] self.paged_attention_V = self.paged_attention[:,1] self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3) - self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = "cuda:0") - self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda:0") - self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda:0") + self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = device) + self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = device) + self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device) # Only for Gemma2 - self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = "cuda:0") - self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda:0") + self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device) + self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = device) self.half_head_dim = head_dim // 2 diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 3e0717a872..fe0627f8d7 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -167,24 +167,25 @@ def LlamaAttention_fast_forward_inference( # Prefill phase # if not hasattr(self, "paged_attention"): + device = hidden_states.device if do_prefill: - self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = "cuda:0") + self.paged_attention = torch.empty((KV_CACHE_INCREMENT+seq_len+1, 2, bsz, n_kv_heads, head_dim), dtype = dtype, device = device) self.paged_attention_K = self.paged_attention[:,0] self.paged_attention_V = self.paged_attention[:,1] self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3) - self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = "cuda:0") - self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda:0") - self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda:0") + self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = device) + self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = device) + self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device) # Mistral Nemo 12b has weird dimensions if attention_size != hidden_size: - self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = "cuda:0") + self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device) else: self.temp_O = self.temp_QA[1][:,:,:hidden_size] pass - self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda:0") + self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = device) self.scalar = 1.0 / math_sqrt(self.head_dim) self.half_head_dim = head_dim // 2 elif kv_seq_len >= self.paged_attention.shape[0]: @@ -813,13 +814,13 @@ def LlamaModel_fast_forward( is_causal = True, sliding_window = self.config.sliding_window, )\ - .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda:0",)\ + .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda",)\ .squeeze(0).squeeze(0) self.GA_mask = AttentionMaskConverter( is_causal = True, )\ - .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda:0",)\ + .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda",)\ .squeeze(0).squeeze(0) pass pass @@ -1075,10 +1076,16 @@ def CausalLM_fast_forward(fast_forward_inference): bsz, q_len, hd = hidden_states.shape lm_head = self.lm_head.weight + lm_head_device = lm_head.device + logit_softcapping = getattr(self.config, "final_logit_softcapping", 0) logit_scaling = getattr(self.config, "logit_scale", 0) dtype = lm_head.dtype num_logits_to_keep = max(num_logits_to_keep, logits_to_keep) + + # Move items to same device as lm_head + hidden_states = hidden_states.to(lm_head_device) + 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": @@ -1148,11 +1155,14 @@ def CausalLM_fast_forward(fast_forward_inference): if labels is not None: shift_logits = logits - if not hasattr(self, "extra_ignored_labels"): - # Fixes https://github.com/unslothai/unsloth/issues/10 - self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0") - pass - shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]])) + # if not hasattr(self, "extra_ignored_labels"): + # # Fixes https://github.com/unslothai/unsloth/issues/10 + # self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0") + # pass + shift_labels = torch.empty_like(labels) + shift_labels[..., :-1] = labels[..., 1:] + shift_labels[..., -1] = -100 + # shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]])) loss = fast_cross_entropy_loss( logits = shift_logits, labels = shift_labels, @@ -1297,7 +1307,7 @@ class LlamaRotaryEmbedding(torch.nn.Module): if seq_len <= self.current_rope_size: return # Iteratively grow by increments of 8192 self.current_rope_size = ((seq_len // 8192) + ((seq_len % 8192) != 0)) * 8192 - self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype) + self._set_cos_sin_cache(self.current_rope_size, device = "cuda", dtype = x.dtype) pass pass @@ -1423,7 +1433,7 @@ class LlamaExtendedRotaryEmbedding(torch.nn.Module): if seq_len <= self.current_rope_size: return # Iteratively grow by increments of 8192 self.current_rope_size = ((seq_len // 8192) + ((seq_len % 8192) != 0)) * 8192 - self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype) + self._set_cos_sin_cache(self.current_rope_size, device = "cuda", dtype = x.dtype) pass pass @@ -1538,7 +1548,7 @@ class LongRopeRotaryEmbedding(torch.nn.Module): if seq_len <= self.current_rope_size: return # Iteratively grow by increments of 8192 self.current_rope_size = ((seq_len // 8192) + ((seq_len % 8192) != 0)) * 8192 - self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype) + self._set_cos_sin_cache(self.current_rope_size, device = "cuda", dtype = x.dtype) pass pass @@ -1771,8 +1781,6 @@ class FastLlamaModel: # Add to kwargs kwargs["rope_scaling"] = rope_scaling pass - # We currently only support NVIDIA GPUs - AMD / Intel is a work in progress! - pre_check = check_nvidia() bnb_config = None if load_in_4bit: @@ -1840,8 +1848,6 @@ class FastLlamaModel: pass # Return old flag os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer - # We currently only support NVIDIA GPUs - AMD / Intel is a work in progress! - post_check = check_nvidia() # Counteract saved tokenizers tokenizer_name = model_name if tokenizer_name is None else tokenizer_name @@ -1874,25 +1880,20 @@ class FastLlamaModel: except: raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!') pass - - if ((post_check - pre_check) >= 1).sum() > 1: - raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!') - + import transformers.trainer items_in_trainer = dir(transformers.trainer) good_items = [] for item in items_in_trainer: - # TODO: Support Deepspeed - if item.startswith(("deepspeed", "xm", "met", "smp")): continue if item in inner_training_loop: good_items.append(item) pass exec("from transformers.trainer import (" + ", ".join(x for x in good_items) + ")", globals()) - start = re.search('logger\.info\([\"\'].+?Running training', inner_training_loop).span(0)[0] + start = re.search(r'logger\.info\([\"\'].+?Running training', inner_training_loop).span(0)[0] end = inner_training_loop.find("\n\n", start) original_debug = inner_training_loop[start:end] - spaces = re.search('\n([\s\t]{1,})', original_debug).group(0)[1:] - front_spaces = re.match('([\s\t]{1,})', inner_training_loop).group(0) + spaces = re.search(r'\n([\s\t]{1,})', original_debug).group(0)[1:] + front_spaces = re.match(r'([\s\t]{1,})', inner_training_loop).group(0) # Cannot use \\ since it will cause a SyntaxWarning in Python 3.12 # Instead use chr(92) == \\ @@ -1903,17 +1904,7 @@ class FastLlamaModel: f"{chr(92)} / Total batch size = {total_train_batch_size:,} | Total steps = {max_steps:,}\\n"\\ f' "-____-" Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}' logger.warning(debug_info) - import subprocess, re, gc, numpy as np - a = np.array([0,]) - try: - a = subprocess.check_output('nvidia-smi --query-gpu=memory.used --format=csv', shell = True) - a = re.findall(rb'([\\d]{1,})[\\s]{1,}M', a) - a = np.array([int(x.decode('utf-8'))/1024 for x in a]) - except: - if not torch.cuda.is_available(): - raise RuntimeError('Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!') - if ((a - PRE_CHECK) >= 1).sum() > 1: - raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!') + import subprocess, re, gc for _ in range(3): gc.collect() torch.cuda.empty_cache()""" @@ -1925,7 +1916,7 @@ class FastLlamaModel: debug_info = """n_total_devices = total_train_batch_size // \\ args.gradient_accumulation_steps // self._train_batch_size if n_total_devices > 1: - logger.warning_once('Unsloth currently does not support multi GPU setups - but we are working on it!') + logger.warning_once('Unsloth is running with multi GPUs - the effective batch size is multiplied by ' + str(n_total_devices)) debug_info =""" debug_info = debug_info.split('\n') debug_info = "\n".join([debug_info[0]] + [spaces + x[8:] for x in debug_info[1:]]) @@ -1937,31 +1928,6 @@ class FastLlamaModel: "train_dataloader = tpu_spmd_dataloader(train_dataloader)", "raise RuntimeError('Unsloth: TPUs are not yet supported!')" ) - inner_training_loop = inner_training_loop.replace( - "self.accelerator.free_memory()", - "self.accelerator.free_memory()\n" + \ - front_spaces + "if self.is_deepspeed_enabled:"\ - "raise RuntimeError('Unsloth: Deepspeed is not yet supported!')\n", 1, - ) - - check_batches = """train_dataloader = self.get_train_dataloader() - ga = args.gradient_accumulation_steps - bsz = self._train_batch_size - total_batches = bsz * ga * args.world_size - n_total_devices = total_batches // ga // bsz - if n_total_devices > 1: - logger.warning_once('Unsloth currently does not support multi GPU setups - but we are working on it!') - divisor = n_total_devices / 1 - bsz = self._train_batch_size = max(int(bsz / divisor), 1) - if total_batches // ga // bsz > 1: - divisor = n_total_devices / 1 - ga = args.gradient_accumulation_steps = max(int(ga / divisor), 1)""" - check_batches = check_batches.split('\n') - check_batches = "\n".join([check_batches[0]] + [front_spaces + x[8:] for x in check_batches[1:]]) - inner_training_loop = inner_training_loop.replace( - "train_dataloader = self.get_train_dataloader()", - check_batches, 1, - ) inner_training_loop = inner_training_loop.replace( "_inner_training_loop", "_fast_inner_training_loop", 1, @@ -1973,13 +1939,6 @@ class FastLlamaModel: "is_torch_tpu_available()", "False", ) - if "n_total_devices >" not in inner_training_loop: - raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!') - pass - inner_training_loop = inner_training_loop.replace( - "is_sagemaker_mp_enabled()", - "False", - ) exec(inner_training_loop, globals()) Trainer._inner_training_loop = _fast_inner_training_loop @@ -2136,7 +2095,7 @@ class FastLlamaModel: pass model.get_input_embeddings().modules_to_save.default\ - .to(device = "cuda:0", dtype = new_dtype, non_blocking = True) + .to(device = "cuda", dtype = new_dtype, non_blocking = True) model.get_input_embeddings().modules_to_save.default.requires_grad_(True) # [TODO] Move old embed_tokens to CPU - should be disk! @@ -2156,7 +2115,7 @@ class FastLlamaModel: pass model.get_output_embeddings().modules_to_save.default\ - .to(device = "cuda:0", dtype = new_dtype, non_blocking = True) + .to(device = "cuda", dtype = new_dtype, non_blocking = True) model.get_output_embeddings().modules_to_save.default.requires_grad_(True) # [TODO] Move old lm_head to CPU - should be disk! @@ -2413,7 +2372,7 @@ class FastLlamaModel: pass model.get_input_embeddings().modules_to_save.default\ - .to(device = "cuda:0", dtype = new_dtype, non_blocking = True) + .to(device = "cuda", dtype = new_dtype, non_blocking = True) model.get_input_embeddings().modules_to_save.default.requires_grad_(True) pass @@ -2429,7 +2388,7 @@ class FastLlamaModel: pass model.get_output_embeddings().modules_to_save.default\ - .to(device = "cuda:0", dtype = new_dtype, non_blocking = True) + .to(device = "cuda", dtype = new_dtype, non_blocking = True) model.get_output_embeddings().modules_to_save.default.requires_grad_(True) pass @@ -2515,12 +2474,7 @@ class FastLlamaModel: from transformers.trainer import Trainer if Trainer._inner_training_loop.__name__ != "_fast_inner_training_loop": - raise RuntimeError( - 'Unsloth currently does not work on multi GPU setups - sadly we are a 2 brother team so '\ - 'enabling it will require much more work, so we have to prioritize. Please understand!\n'\ - 'We do have a separate beta version, which you can contact us about!\n'\ - 'Thank you for your understanding and we appreciate it immensely!' - ) + raise RuntimeError("Unsloth: Unsuccessfully patched Trainer! Please file a bug report!") pass # Fix loftq issues @@ -2636,8 +2590,8 @@ class FastLlamaModel: # Patch cross entropy loss labels # Fixes https://github.com/unslothai/unsloth/issues/10 max_seq_length = model.max_seq_length - extra_ignored_labels = torch.full((max_seq_length, 1), -100, device = "cuda:0") - model.model.extra_ignored_labels = extra_ignored_labels + # extra_ignored_labels = torch.full((max_seq_length, 1), -100, device = "cuda:0") + # model.model.extra_ignored_labels = extra_ignored_labels internal_model = model while hasattr(internal_model, "model"): internal_model.max_seq_length = max_seq_length diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 784ca9cb41..303c3d9589 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -35,6 +35,7 @@ except: MistralSdpaAttention = MistralAttention MistralFlashAttention2 = MistralAttention pass +from unsloth_zoo.utils import Version, _get_dtype def MistralAttention_fast_forward( @@ -183,6 +184,7 @@ def MistralForCausalLM_fast_forward( output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, num_logits_to_keep: Optional[int] = 0, + logits_to_keep: Optional[int] = 0, *args, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: @@ -194,7 +196,6 @@ def MistralForCausalLM_fast_forward( elif q_len <= sliding_window: causal_mask = xformers.attn_bias.LowerTriangularMask() else: - # Fix from https://github.com/Rypo causal_mask = xformers.attn_bias.BlockDiagonalCausalMask\ .from_seqlens([q_len]*bsz)\ .make_local_attention(window_size = sliding_window) @@ -219,41 +220,92 @@ def MistralForCausalLM_fast_forward( ) else: outputs = self.model( - input_ids=input_ids, - causal_mask=causal_mask, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, + input_ids = input_ids, + causal_mask = causal_mask, + attention_mask = attention_mask, + position_ids = position_ids, + past_key_values = past_key_values, + inputs_embeds = inputs_embeds, + use_cache = use_cache, + output_attentions = output_attentions, + output_hidden_states = output_hidden_states, + return_dict = return_dict, ) pass hidden_states = outputs[0] + bsz, q_len, hd = hidden_states.shape lm_head = self.lm_head.weight + lm_head_device = lm_head.device + + # Move items to same device as lm_head + hidden_states = hidden_states.to(lm_head_device) + if labels is not None: labels = labels.to(lm_head_device) + + # If we are in GRPO mode, return raw hidden states + if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1": + num_logits_to_keep = max(num_logits_to_keep, logits_to_keep) + if num_logits_to_keep != 0: + hidden_states = hidden_states[:, -num_logits_to_keep:, :] + return CausalLMOutputWithPast( + loss = None, + logits = hidden_states, + past_key_values = outputs.past_key_values, + hidden_states = outputs.hidden_states, + attentions = outputs.attentions, + ) + pass + if bsz == 1 and q_len == 1: logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype)) logits = logits.unsqueeze(0).unsqueeze(0) elif num_logits_to_keep != 0: logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :].to(lm_head.dtype)) else: + RETURN_LOGITS = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1" + # < 1024 Normal Unsloth uses less VRAM! + if bsz * q_len <= 1024: RETURN_LOGITS = True + + if not RETURN_LOGITS and HAS_CUT_CROSS_ENTROPY and labels is not None: + n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None) + logit_softcapping = getattr(self.config, "final_logit_softcapping", 0) + loss = fused_linear_cross_entropy( + hidden_states = hidden_states, + lm_weight = lm_head, + labels = labels, + num_items_in_batch = n_items, + logit_softcapping = logit_softcapping, + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + output = CausalLMOutputWithPast( + loss = loss, + logits = EMPTY_LOGITS, + past_key_values = outputs.past_key_values, + hidden_states = outputs.hidden_states, + attentions = outputs.attentions, + ) + return output + pass logits = self.lm_head(hidden_states.to(lm_head.dtype)) pass - logits = logits.to(self.config.torch_dtype) + logits = logits.to(_get_dtype(self.config.torch_dtype)) loss = None if labels is not None: shift_logits = logits - if not hasattr(self, "extra_ignored_labels"): - # Fixes https://github.com/unslothai/unsloth/issues/10 - self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0") - pass - - shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]])) + # if not hasattr(self, "extra_ignored_labels"): + # # Fixes https://github.com/unslothai/unsloth/issues/10 + # self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0") + # pass + # shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]])) + shift_labels = torch.empty_like(labels) + shift_labels[..., :-1] = labels[..., 1:] + shift_labels[..., -1] = -100 loss = fast_cross_entropy_loss( logits = shift_logits, labels = shift_labels, @@ -266,11 +318,11 @@ def MistralForCausalLM_fast_forward( return (loss,) + output if loss is not None else output return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, + loss = loss, + logits = logits, + past_key_values = outputs.past_key_values, + hidden_states = outputs.hidden_states, + attentions = outputs.attentions, ) pass diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 06ae82140b..fe7f4accee 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -93,7 +93,7 @@ def sft_trainer_prepare_dataset(function_name, function): " tokenizer = partial(tokenizer, add_special_tokens = False)\n"\ " processing_class = tokenizer\n"\ "else:\n"\ - " add_special_tokens = False if has_bos_token_already else add_special_tokens\n" + " add_special_tokens = False if has_bos_token_already else locals().get('add_special_tokens', False)\n" check_text = check_text.split("\n") check_text = "\n".join(" "*8 + x for x in check_text) @@ -101,7 +101,7 @@ def sft_trainer_prepare_dataset(function_name, function): # .*? matches first match. .+? matches final match. replacer = re.findall( - r"def {function_name}\(.*?\).*?\:\n", + r"def " + function_name + r"\(.*?\).*?\:\n", function, flags = re.MULTILINE | re.DOTALL, ) @@ -164,7 +164,7 @@ RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__prepare_inputs) # Remove _move_model_to_vllm def grpo_trainer__move_model_to_vllm(function_name, function): if function_name != "_move_model_to_vllm": return function - + def _move_model_to_vllm(self, *args, **kwargs): return None function = inspect.getsource(_move_model_to_vllm) @@ -246,14 +246,20 @@ def grpo_trainer_compute_loss(function_name, function): self, _input_ids, logits_to_keep, completion_mask, advantages, n_chunks = self.args.unsloth_num_chunks, ) - + # Log the metrics # completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item() - self._metrics["completion_length"].append(completion_length.item()) # mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() # self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) - self._metrics["kl"].append(mean_kl.item()) + + if "train" in self._metrics: + mode = "eval" if self.control.should_evaluate else "train" + self._metrics[mode]["completion_length"].append(completion_length.item()) + self._metrics[mode]["kl"].append(mean_kl.item()) + else: + self._metrics["completion_length"].append(completion_length.item()) + self._metrics["kl"].append(mean_kl.item()) return loss pass diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 51450aa0d9..d13d394669 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1,18 +1,16 @@ -# Unsloth Zoo - Utilities for Unsloth # Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. # -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Lesser General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at # -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. +# http://www.apache.org/licenses/LICENSE-2.0 # -# You should have received a copy of the GNU Lesser General Public License -# along with this program. If not, see . +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. import torch from transformers import ( @@ -98,9 +96,9 @@ class FastBaseVisionModel: statistics = \ f"==((====))== Unsloth {__version__}: Fast {model_types[0].title()} vision patching. Transformers: {transformers_version}.\n"\ - f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB. Platform: {platform_system}.\n"\ - f"O^O/ \_/ \\ Torch: {torch.__version__}. CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {torch.version.cuda}. Triton: {triton_version}\n"\ - f"\ / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\n"\ + f" {chr(92)}{chr(92)} /| GPU: {gpu_stats.name}. 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"\ f' "-____-" Free Apache license: http://github.com/unslothai/unsloth' print(statistics) @@ -123,9 +121,6 @@ class FastBaseVisionModel: assert(dtype == torch.float16 or dtype == torch.bfloat16 or dtype == torch.float32) - # We currently only support NVIDIA GPUs - AMD / Intel is a work in progress! - pre_check = check_nvidia() - bnb_config = None if load_in_4bit: bnb_config = BitsAndBytesConfig( @@ -154,8 +149,6 @@ class FastBaseVisionModel: ) # Return old flag os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer - # We currently only support NVIDIA GPUs - AMD / Intel is a work in progress! - post_check = check_nvidia() # Counteract saved tokenizers tokenizer_name = model_name if tokenizer_name is None else tokenizer_name diff --git a/unsloth/save.py b/unsloth/save.py index 6770d658c8..d03f47e874 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -484,8 +484,8 @@ def unsloth_save_model( max_ram = psutil.virtual_memory().available sharded_ram_usage = 5 * 1024 * 1024 * 1024 if type(max_shard_size) is str: - gb_found = re.match("([0-9]{1,})[\s]{0,}GB", max_shard_size, flags = re.IGNORECASE) - mb_found = re.match("([0-9]{1,})[\s]{0,}MB", max_shard_size, flags = re.IGNORECASE) + gb_found = re.match(r"([0-9]{1,})[\s]{0,}GB", max_shard_size, flags = re.IGNORECASE) + mb_found = re.match(r"([0-9]{1,})[\s]{0,}MB", max_shard_size, flags = re.IGNORECASE) if gb_found: sharded_ram_usage = int(gb_found.group(1)) * 1024 * 1024 * 1024 elif mb_found: sharded_ram_usage = int(mb_found.group(1)) * 1024 * 1024 elif type(max_shard_size) is int: @@ -1019,9 +1019,9 @@ def save_to_gguf( print_info = \ f"==((====))== Unsloth: Conversion from QLoRA to GGUF information\n"\ - f" \\\ /| [0] Installing llama.cpp might take 3 minutes.\n"\ - f"O^O/ \_/ \\ [1] Converting HF to GGUF 16bits might take 3 minutes.\n"\ - f"\ / [2] Converting GGUF 16bits to {quantization_method} might take 10 minutes each.\n"\ + f" {chr(92)}{chr(92)} /| [0] Installing llama.cpp might take 3 minutes.\n"\ + f"O^O/ {chr(92)}_/ {chr(92)} [1] Converting HF to GGUF 16bits might take 3 minutes.\n"\ + f"{chr(92)} / [2] Converting GGUF 16bits to {quantization_method} might take 10 minutes each.\n"\ f' "-____-" In total, you will have to wait at least 16 minutes.\n' print(print_info) diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 048bee7797..9c5f825a0c 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -857,21 +857,6 @@ def check_tokenizer( pass -def check_nvidia(): - # Unsloth doesn't work yet on AMD devices - we're working on it! - output = np.array([0,]) - try: - output = subprocess.check_output("nvidia-smi --query-gpu=memory.used --format=csv", shell = True) - output = re.findall(rb'([\d]{1,})[\s]{1,}M', output) - output = np.array([int(x.decode('utf-8'))/1024 for x in output]) - except: - if not torch.cuda.is_available(): - raise RuntimeError("Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!") - return output -pass -PRE_CHECK = check_nvidia() - - import inspect from inspect import getsource import trl.trainer.sft_trainer