diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml index 4ebb6df3d0..ae5dade42d 100644 --- a/.github/FUNDING.yml +++ b/.github/FUNDING.yml @@ -3,7 +3,7 @@ github: unslothai patreon: # Replace with a single Patreon username open_collective: # Replace with a single Open Collective username -ko_fi: unsloth +ko_fi: # unsloth tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry liberapay: # Replace with a single Liberapay username diff --git a/.github/ISSUE_TEMPLATE/bug---issue.md b/.github/ISSUE_TEMPLATE/bug---issue.md index 397d725f95..83e0fd73a9 100644 --- a/.github/ISSUE_TEMPLATE/bug---issue.md +++ b/.github/ISSUE_TEMPLATE/bug---issue.md @@ -18,4 +18,4 @@ assignees: '' Put Minimal code to reproduce error here ###Remove Hugging Face token### ``` -🦥 You can also ask via our Reddit page: https://www.reddit.com/r/unsloth/ +🦥 You can also ask via our Reddit page: https://reddit.com/r/unsloth/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 545c7899aa..bd37ece943 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.10 + rev: v0.14.13 hooks: - id: ruff args: diff --git a/README.md b/README.md index 43c09381fc..8ddfa80c35 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@
- + unsloth logo @@ -8,7 +8,7 @@ - + ### Train gpt-oss, DeepSeek, Gemma, Qwen & Llama 2x faster with 70% less VRAM! @@ -18,25 +18,25 @@ ## ✨ Train for Free -Notebooks are beginner friendly. Read our [guide](https://docs.unsloth.ai/get-started/fine-tuning-guide). Add dataset, run, then export your trained model to GGUF, llama.cpp, Ollama, vLLM, SGLang or Hugging Face. +Notebooks are beginner friendly. Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Add dataset, run, then deploy your trained model. | Model | Free Notebooks | Performance | Memory use | |-----------|---------|--------|----------| | **gpt-oss (20B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-Fine-tuning.ipynb) | 1.5x faster | 70% less | -| **Mistral Ministral 3 (3B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Ministral_3_VL_(3B)_Vision.ipynb) | 1.5x faster | 60% less | | **gpt-oss (20B): GRPO** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-GRPO.ipynb) | 2x faster | 80% less | | **Qwen3: Advanced GRPO** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_(4B)-GRPO.ipynb) | 2x faster | 50% less | | **Qwen3-VL (8B): GSPO** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_VL_(8B)-Vision-GRPO.ipynb) | 1.5x faster | 80% less | -| **Gemma 3 (270M)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma3_(270M).ipynb) | 1.7x faster | 60% less | -| **Gemma 3n (4B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma3N_(4B)-Conversational.ipynb) | 1.5x faster | 50% less | -| **DeepSeek-OCR (3B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Deepseek_OCR_(3B).ipynb) | 1.5x faster | 30% less | +| **Gemma 3 (4B) Vision** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma3_(4B)-Vision.ipynb) | 1.7x faster | 60% less | +| **Gemma 3n (e4B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma3N_(4B)-Conversational.ipynb) | 1.5x faster | 50% less | +| **embeddinggemma (300M)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/EmbeddingGemma_(300M).ipynb) | 2x faster | 20% less | +| **Mistral Ministral 3 (3B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Ministral_3_VL_(3B)_Vision.ipynb) | 1.5x faster | 60% less | | **Llama 3.1 (8B) Alpaca** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-Alpaca.ipynb) | 2x faster | 70% less | | **Llama 3.2 Conversational** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.2_(1B_and_3B)-Conversational.ipynb) | 2x faster | 70% less | | **Orpheus-TTS (3B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Orpheus_(3B)-TTS.ipynb) | 1.5x faster | 50% less | -- See all our notebooks for: [Kaggle](https://github.com/unslothai/notebooks?tab=readme-ov-file#-kaggle-notebooks), [GRPO](https://docs.unsloth.ai/get-started/unsloth-notebooks#grpo-reasoning-rl-notebooks), **[TTS](https://docs.unsloth.ai/get-started/unsloth-notebooks#text-to-speech-tts-notebooks)** & [Vision](https://docs.unsloth.ai/get-started/unsloth-notebooks#vision-multimodal-notebooks) -- See [all our models](https://docs.unsloth.ai/get-started/all-our-models) and [all our notebooks](https://docs.unsloth.ai/get-started/unsloth-notebooks) -- See detailed documentation for Unsloth [here](https://docs.unsloth.ai/) +- See all our notebooks for: [Kaggle](https://github.com/unslothai/notebooks?tab=readme-ov-file#-kaggle-notebooks), [GRPO](https://unsloth.ai/docs/get-started/unsloth-notebooks#grpo-reasoning-rl-notebooks), [TTS](https://unsloth.ai/docs/get-started/unsloth-notebooks#text-to-speech-tts-notebooks), [embedding](https://unsloth.ai/docs/new/embedding-finetuning) & [Vision](https://unsloth.ai/docs/get-started/unsloth-notebooks#vision-multimodal-notebooks) +- See [all our models](https://unsloth.ai/docs/get-started/unsloth-model-catalog) and [all our notebooks](https://unsloth.ai/docs/get-started/unsloth-notebooks) +- See detailed documentation for Unsloth [here](https://unsloth.ai/docs) ## ⚡ Quickstart ### Linux or WSL @@ -44,33 +44,36 @@ Notebooks are beginner friendly. Read our [guide](https://docs.unsloth.ai/get-st pip install unsloth ``` ### Windows -For Windows, `pip install unsloth` works only if you have Pytorch installed. Read our [Windows Guide](https://docs.unsloth.ai/get-started/installing-+-updating/windows-installation). +For Windows, `pip install unsloth` works only if you have Pytorch installed. Read our [Windows Guide](https://unsloth.ai/docs/get-started/install-and-update/windows-installation). + ### Docker -Use our official [Unsloth Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Read our [Docker Guide](https://docs.unsloth.ai/get-started/install-and-update/docker). +Use our official [Unsloth Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Read our [Docker Guide](https://unsloth.ai/docs/get-started/install-and-update/docker). + ### Blackwell & DGX Spark -For RTX 50x, B200, 6000 GPUs: `pip install unsloth`. Read our [Blackwell Guide](https://docs.unsloth.ai/basics/training-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark Guide](https://docs.unsloth.ai/new/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth) for more details. +For RTX 50x, B200, 6000 GPUs: `pip install unsloth`. Read our [Blackwell Guide](https://unsloth.ai/docs/basics/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark Guide](https://unsloth.ai/docs/basics/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth) for more details. ## 🦥 Unsloth News -- New RoPE & MLP **Triton Kernels** & **Padding Free + Packing**: 3x faster training & 30% less VRAM. [Blog](https://docs.unsloth.ai/new/3x-faster-training-packing) -- **Ministral 3** by Mistral: Run Ministral 3 or fine-tune with vision/RL sodoku notebooks. [Guide](https://docs.unsloth.ai/new/ministral-3) • [Notebooks](https://docs.unsloth.ai/new/ministral-3#fine-tuningb) -- **500K Context**: Training a 20B model with >500K context is now possible on an 80GB GPU. [Blog](https://docs.unsloth.ai/new/500k-context-length-fine-tuning) -- **FP8 Reinforcement Learning**: You can now do FP8 GRPO on consumer GPUs. [Blog](https://docs.unsloth.ai/new/fp8-reinforcement-learning) • [Notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_8B_FP8_GRPO.ipynb) -- **DeepSeek-OCR**: Fine-tune to improve language understanding by 89%. [Guide](https://docs.unsloth.ai/new/deepseek-ocr-run-and-fine-tune) • [Notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Deepseek_OCR_(3B).ipynb) -- **Docker**: Use Unsloth with no setup & environment issues with our new image. [Guide](https://docs.unsloth.ai/new/how-to-train-llms-with-unsloth-and-docker) • [Docker image](https://hub.docker.com/r/unsloth/unsloth) -- **gpt-oss RL**: Introducing the fastest possible inference for gpt-oss RL! [Read blog](https://docs.unsloth.ai/new/gpt-oss-reinforcement-learning) -- **Vision RL**: You can now train VLMs with GRPO or GSPO in Unsloth! [Read guide](https://docs.unsloth.ai/new/vision-reinforcement-learning-vlm-rl) -- **gpt-oss** by OpenAI: Read our [Unsloth Flex Attention](https://docs.unsloth.ai/new/long-context-gpt-oss-training) blog and [gpt-oss Guide](https://docs.unsloth.ai/basics/gpt-oss). 20B works on 14GB VRAM. 120B on 65GB. +- **Embedding models**: Unsloth now supports ~1.8-3.3x faster embedding fine-tuning. [Blog](https://unsloth.ai/docs/new/embedding-finetuning) • [Notebooks](https://unsloth.ai/docs/get-started/unsloth-notebooks#embedding-models) +- New **7x longer context RL** vs. all other setups, via our new batching algorithms. [Blog](https://unsloth.ai/docs/new/grpo-long-context) +- New RoPE & MLP **Triton Kernels** & **Padding Free + Packing**: 3x faster training & 30% less VRAM. [Blog](https://unsloth.ai/docs/new/3x-faster-training-packing) +- **500K Context**: Training a 20B model with >500K context is now possible on an 80GB GPU. [Blog](https://unsloth.ai/docs/new/500k-context-length-fine-tuning) +- **FP8 Reinforcement Learning**: You can now do FP8 GRPO on consumer GPUs. [Blog](https://unsloth.ai/docs/new/fp8-reinforcement-learning) • [Notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_8B_FP8_GRPO.ipynb) +- **DeepSeek-OCR**: Fine-tune to improve language understanding by 89%. [Guide](https://unsloth.ai/docs/models/deepseek-ocr-how-to-run-and-fine-tune) • [Notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Deepseek_OCR_(3B).ipynb) +- **Docker**: Use Unsloth with no setup & environment issues with our new image. [Guide](https://unsloth.ai/docs/new/how-to-fine-tune-llms-with-unsloth-and-docker) • [Docker image](https://hub.docker.com/r/unsloth/unsloth) +- **Vision RL**: You can now train VLMs with GRPO or GSPO in Unsloth! [Read guide](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/vision-reinforcement-learning-vlm-rl) +- **gpt-oss** by OpenAI: Read our [RL blog](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/gpt-oss-reinforcement-learning), [Flex Attention](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/long-context-gpt-oss-training) blog and [gpt-oss Guide](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune). 20B works on 14GB VRAM. 120B on 65GB.
Click for more news -- **Quantization-Aware Training**: We collabed with Pytorch, recovering ~70% accuracy. [Read blog](https://docs.unsloth.ai/new/quantization-aware-training-qat) -- **Memory-efficient RL**: We're introducing even better RL. Our new kernels & algos allows faster RL with 50% less VRAM & 10× more context. [Read blog](https://docs.unsloth.ai/new/memory-efficient-rl) -- **Gemma 3n** by Google: [Read Blog](https://docs.unsloth.ai/basics/gemma-3n-how-to-run-and-fine-tune). We [uploaded GGUFs, 4-bit models](https://huggingface.co/collections/unsloth/gemma-3n-685d3874830e49e1c93f9339). -- **[Text-to-Speech (TTS)](https://docs.unsloth.ai/basics/text-to-speech-tts-fine-tuning)** is now supported, including `sesame/csm-1b` and STT `openai/whisper-large-v3`. -- **[Qwen3](https://docs.unsloth.ai/basics/qwen3-how-to-run-and-fine-tune)** is now supported. Qwen3-30B-A3B fits on 17.5GB VRAM. -- Introducing **[Dynamic 2.0](https://docs.unsloth.ai/basics/unsloth-dynamic-2.0-ggufs)** quants that set new benchmarks on 5-shot MMLU & Aider Polyglot. -- [**EVERYTHING** is now supported](https://unsloth.ai/blog/gemma3#everything) - all models (TTS, BERT, Mamba), FFT, etc. [MultiGPU](https://docs.unsloth.ai/basics/multi-gpu-training-with-unsloth) coming soon. Enable FFT with `full_finetuning = True`, 8-bit with `load_in_8bit = True`. +- **Quantization-Aware Training**: We collabed with Pytorch, recovering ~70% accuracy. [Read blog](https://unsloth.ai/docs/basics/quantization-aware-training-qat) +- **Memory-efficient RL**: We're introducing even better RL. Our new kernels & algos allows faster RL with 50% less VRAM & 10× more context. [Read blog](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/memory-efficient-rl) +- **Mistral 3**: Run Ministral 3 or Devstral 2 and fine-tune with vision/RL sodoku notebooks. [Guide](https://unsloth.ai/docs/models/ministral-3) • [Notebooks](https://unsloth.ai/docs/models/ministral-3#fine-tuning-ministral-3) +- **Gemma 3n** by Google: [Read Blog](https://unsloth.ai/docs/models/gemma-3-how-to-run-and-fine-tune/gemma-3n-how-to-run-and-fine-tune). We [uploaded GGUFs, 4-bit models](https://huggingface.co/collections/unsloth/gemma-3n-685d3874830e49e1c93f9339). +- **[Text-to-Speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning)** is now supported, including `sesame/csm-1b` and STT `openai/whisper-large-v3`. +- **[Qwen3](https://unsloth.ai/docs/models/qwen3-how-to-run-and-fine-tune)** is now supported. Qwen3-30B-A3B fits on 17.5GB VRAM. +- Introducing **[Dynamic 2.0](https://unsloth.ai/docs/basics/unsloth-dynamic-2.0-ggufs)** quants that set new benchmarks on 5-shot MMLU & Aider Polyglot. +- [**EVERYTHING** is now supported](https://unsloth.ai/blog/gemma3#everything) - all models (TTS, BERT, Mamba), FFT, etc. [MultiGPU](https://unsloth.ai/docs/basics/multi-gpu-training-with-unsloth) coming soon. Enable FFT with `full_finetuning = True`, 8-bit with `load_in_8bit = True`. - 📣 [DeepSeek-R1](https://unsloth.ai/blog/deepseek-r1) - run or fine-tune them [with our guide](https://unsloth.ai/blog/deepseek-r1). All model uploads: [here](https://huggingface.co/collections/unsloth/deepseek-r1-all-versions-678e1c48f5d2fce87892ace5). - 📣 Introducing Long-context [Reasoning (GRPO)](https://unsloth.ai/blog/grpo) in Unsloth. Train your own reasoning model with just 5GB VRAM. Transform Llama, Phi, Mistral etc. into reasoning LLMs! - 📣 Introducing Unsloth [Dynamic 4-bit Quantization](https://unsloth.ai/blog/dynamic-4bit)! We dynamically opt not to quantize certain parameters and this greatly increases accuracy while only using <10% more VRAM than BnB 4-bit. See our collection on [Hugging Face here.](https://huggingface.co/collections/unsloth/unsloth-4-bit-dynamic-quants-67503bb873f89e15276c44e7) @@ -84,27 +87,29 @@ For RTX 50x, B200, 6000 GPUs: `pip install unsloth`. Read our [Blackwell Guide](
## 🔗 Links and Resources -| Type | Links | -| ------------------------------- | --------------------------------------- | -|   **r/unsloth Reddit** | [Join Reddit community](https://reddit.com/r/unsloth)| -| 📚 **Documentation & Wiki** | [Read Our Docs](https://docs.unsloth.ai) | -|   **Twitter (aka X)** | [Follow us on X](https://twitter.com/unslothai)| -| 💾 **Installation** | [Pip & Docker Install](https://docs.unsloth.ai/get-started/installing-+-updating)| -| 🔮 **Our Models** | [Unsloth Catalog](https://docs.unsloth.ai/get-started/all-our-models)| -| ✍️ **Blog** | [Read our Blogs](https://unsloth.ai/blog)| +| Type | Links | +| ----------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ | +|   **r/unsloth Reddit** | [Join Reddit community](https://reddit.com/r/unsloth) | +| 📚 **Documentation & Wiki** | [Read Our Docs](https://unsloth.ai/docs) | +|   **Twitter (aka X)** | [Follow us on X](https://twitter.com/unslothai) | +| 💾 **Installation** | [Pip & Docker Install](https://unsloth.ai/docs/get-started/install-and-update) | +| 🔮 **Our Models** | [Unsloth Catalog](https://unsloth.ai/docs/get-started/unsloth-model-catalog) | +| ✍️ **Blog** | [Read our Blogs](https://unsloth.ai/blog) | ## ⭐ Key Features -- Supports **full-finetuning**, pretraining, 4b-bit, 16-bit and **FP8** training -- Supports **all models** including [TTS](https://docs.unsloth.ai/basics/text-to-speech-tts-fine-tuning), multimodal, [BERT](https://docs.unsloth.ai/get-started/unsloth-notebooks#other-important-notebooks) and more! Any model that works in transformers, works in Unsloth. -- The most efficient library for [Reinforcement Learning (RL)](https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide), using 80% less VRAM. Supports GRPO, GSPO, DrGRPO, DAPO etc. -- **0% loss in accuracy** - no approximation methods - all exact. -- Supports NVIDIA (since 2018), [AMD](https://docs.unsloth.ai/get-started/install-and-update/amd) and Intel GPUs. Minimum CUDA Capability 7.0 (V100, T4, Titan V, RTX 20, 30, 40x, A100, H100, L40 etc) -- Works on **Linux**, WSL and **Windows** -- All kernels written in [OpenAI's Triton](https://openai.com/index/triton/) language. Manual backprop engine. -- If you trained a model with 🦥Unsloth, you can use this cool sticker!   + +* Supports **full-finetuning**, pretraining, 4b-bit, 16-bit and **FP8** training +* Supports **all models** including [TTS](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning), multimodal, [embedding](https://unsloth.ai/docs/new/embedding-finetuning) and more! Any model that works in transformers, works in Unsloth. +* The most efficient library for [Reinforcement Learning (RL)](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide), using 80% less VRAM. Supports GRPO, GSPO, DrGRPO, DAPO etc. +* **0% loss in accuracy** - no approximation methods - all exact. +* Export and [deploy your model](https://unsloth.ai/docs/basics/inference-and-deployment) to GGUF, llama.cpp, vLLM, SGLang and Hugging Face. +* Supports NVIDIA (since 2018), [AMD](https://unsloth.ai/docs/get-started/install-and-update/amd) and Intel GPUs. Minimum CUDA Capability 7.0 (V100, T4, Titan V, RTX 20, 30, 40x, A100, H100, L40 etc) +* Works on **Linux**, WSL and **Windows** +* All kernels written in OpenAI's Triton language. Manual backprop engine. +* If you trained a model with 🦥Unsloth, you can use this cool sticker!   ## 💾 Install Unsloth -You can also see our docs for more detailed installation and updating instructions [here](https://docs.unsloth.ai/get-started/installing-+-updating). +You can also see our docs for more detailed installation and updating instructions [here](https://unsloth.ai/docs/get-started/install-and-update). Unsloth supports Python 3.13 or lower. @@ -124,7 +129,7 @@ See [here](#advanced-pip-installation) for advanced pip install instructions. You should install the latest driver for your GPU. Download drivers here: [NVIDIA GPU Driver](https://www.nvidia.com/Download/index.aspx). 3. **Install Visual Studio C++:** - You will need Visual Studio, with C++ installed. By default, C++ is not installed with [Visual Studio](https://visualstudio.microsoft.com/vs/community/), so make sure you select all of the C++ options. Also select options for Windows 10/11 SDK. For detailed instructions with options, see [here](https://docs.unsloth.ai/get-started/installing-+-updating). + You will need Visual Studio, with C++ installed. By default, C++ is not installed with [Visual Studio](https://visualstudio.microsoft.com/vs/community/), so make sure you select all of the C++ options. Also select options for Windows 10/11 SDK. For detailed instructions with options, see [here](https://unsloth.ai/docs/get-started/install-and-update/windows-installation#method-3-windows-directly). 5. **Install CUDA Toolkit:** Follow the instructions to install [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit-archive). @@ -139,19 +144,7 @@ See [here](#advanced-pip-installation) for advanced pip install instructions. pip install unsloth ``` -#### Notes -To run Unsloth directly on Windows: -- Install Triton from this Windows fork and follow the instructions [here](https://github.com/woct0rdho/triton-windows) (be aware that the Windows fork requires PyTorch >= 2.4 and CUDA 12) -- In the `SFTConfig`, set `dataset_num_proc=1` to avoid a crashing issue: -```python -SFTConfig( - dataset_num_proc=1, - ... -) -``` - #### Advanced/Troubleshooting - For **advanced installation instructions** or if you see weird errors during installations: First try using an isolated environment via then `pip install unsloth` @@ -268,7 +261,7 @@ print(f'pip install --upgrade pip && pip install --no-deps git+https://github.co ``` ### Docker Installation You can use our pre-built Docker container with all dependencies to use Unsloth instantly with no setup required. -[Read our guide](https://docs.unsloth.ai/get-started/install-and-update/docker). +[Read our guide](https://unsloth.ai/docs/get-started/install-and-update/docker). This container requires installing [NVIDIA's Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html). @@ -283,9 +276,9 @@ docker run -d -e JUPYTER_PASSWORD="mypassword" \ Access Jupyter Lab at `http://localhost:8888` and start fine-tuning! ## 📜 Documentation -- Go to our official [Documentation](https://docs.unsloth.ai) for [running models](https://docs.unsloth.ai/basics/running-and-saving-models), [saving to GGUF](https://docs.unsloth.ai/basics/running-and-saving-models/saving-to-gguf), [checkpointing](https://docs.unsloth.ai/basics/finetuning-from-last-checkpoint), [evaluation](https://docs.unsloth.ai/get-started/fine-tuning-llms-guide#evaluation) and more! -- Read our Guides for: [Fine-tuning](https://docs.unsloth.ai/get-started/fine-tuning-llms-guide), [Reinforcement Learning](https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide), [Text-to-Speech (TTS)](https://docs.unsloth.ai/basics/text-to-speech-tts-fine-tuning), [Vision](https://docs.unsloth.ai/basics/vision-fine-tuning) and [any model](https://docs.unsloth.ai/models/tutorials-how-to-fine-tune-and-run-llms). -- We support Huggingface's transformers, TRL, Trainer, Seq2SeqTrainer and Pytorch code. +* Go to our official [Documentation](https://unsloth.ai/docs) for [running models](https://unsloth.ai/docs/basics/inference-and-deployment), [saving to GGUF](https://unsloth.ai/docs/basics/inference-and-deployment/saving-to-gguf), [checkpointing](https://unsloth.ai/docs/basics/finetuning-from-last-checkpoint), [evaluation](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide#evaluation) and more! +* Read our Guides for: [Fine-tuning](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide), [Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide), [Text-to-Speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning), [Vision](https://unsloth.ai/docs/basics/vision-fine-tuning) and [any model](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms). +* We support Huggingface's transformers, TRL, Trainer, Seq2SeqTrainer and Pytorch code. Unsloth example code to fine-tune gpt-oss-20b: @@ -310,8 +303,9 @@ model, tokenizer = FastModel.from_pretrained( max_seq_length = 2048, # Choose any for long context! load_in_4bit = True, # 4-bit quantization. False = 16-bit LoRA. load_in_8bit = False, # 8-bit quantization - load_in_16bit = False, # [NEW!] 16-bit LoRA + load_in_16bit = False, # 16-bit LoRA full_finetuning = False, # Use for full fine-tuning. + trust_remote_code = False, # Enable to support new models # token = "hf_...", # use one if using gated models ) @@ -350,7 +344,7 @@ trainer = SFTTrainer( ) trainer.train() -# Go to https://docs.unsloth.ai for advanced tips like +# Go to https://unsloth.ai/docs for advanced tips like # (1) Saving to GGUF / merging to 16bit for vLLM or SGLang # (2) Continued training from a saved LoRA adapter # (3) Adding an evaluation loop / OOMs @@ -359,14 +353,15 @@ trainer.train() ## 💡 Reinforcement Learning -[RL](https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide) including [GRPO](https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide#training-with-grpo), [GSPO](https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide/gspo-reinforcement-learning), **FP8** traning, DrGRPO, DAPO, PPO, Reward Modelling, Online DPO all work with Unsloth. -Read our [Reinforcement Learning Guide](https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide) or our [advanced RL docs](https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide/advanced-rl-documentation) for batching, generation & training parameters. +[RL](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide) including [GRPO](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide#training-with-grpo), [GSPO](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/gspo-reinforcement-learning), [**FP8** training](https://unsloth.ai/docs/new/fp8-reinforcement-learning), DrGRPO, DAPO, PPO, Reward Modelling, Online DPO all work with Unsloth. + +Read our [Reinforcement Learning Guide](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide) or our [advanced RL docs](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/advanced-rl-documentation) for batching, generation & training parameters. List of RL notebooks: - gpt-oss GSPO notebook: [Link](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-GRPO.ipynb) -- Qwen2.5-VL GSPO notebook: [Link](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen2_5_7B_VL_GRPO.ipynb) +- - ***FP8*** Qwen3-8B GRPO notebook (L4): [Link](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_8B_FP8_GRPO.ipynb) +- Qwen2.3-VL GSPO notebook: [Link](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_VL_(8B)-Vision-GRPO.ipynb) - Advanced Qwen3 GRPO notebook: [Link](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_(4B)-GRPO.ipynb) -- ***FP8*** Qwen3-8B GRPO notebook (L4): [Link](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_8B_FP8_GRPO.ipynb) - 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/1MRgGtLWuZX4ypSfGguFgC-IblTvO2ivM?usp=sharing) @@ -426,6 +421,4 @@ You can cite the Unsloth repo as follows: - The [llama.cpp library](https://github.com/ggml-org/llama.cpp) that lets users save models with Unsloth - The Hugging Face team and their libraries: [transformers](https://github.com/huggingface/transformers) and [TRL](https://github.com/huggingface/trl) - The Pytorch and [Torch AO](https://github.com/unslothai/unsloth/pull/3391) team for their contributions -- [Erik](https://github.com/erikwijmans) for his help adding [Apple's ML Cross Entropy](https://github.com/apple/ml-cross-entropy) in Unsloth -- [Etherl](https://github.com/Etherll) for adding support for [TTS, diffusion and BERT models](https://github.com/unslothai/notebooks/pull/34) - And of course for every single person who has contributed or has used Unsloth! diff --git a/pyproject.toml b/pyproject.toml index 5f51c097fc..18519eb750 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,16 +55,16 @@ huggingfacenotorch = [ "sentencepiece>=0.2.0", "datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0", "accelerate>=0.34.1", - "peft>=0.7.1,!=0.11.0", + "peft>=0.18.0,!=0.11.0", "huggingface_hub>=0.34.0", "hf_transfer", "diffusers", - "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,<=4.57.3", + "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,<=4.57.6", "trl>=0.18.2,!=0.19.0,<=0.24.0", ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2025.12.7", + "unsloth_zoo>=2026.1.4", "torchvision", "unsloth[triton]", ] @@ -517,10 +517,10 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2025.12.7", + "unsloth_zoo>=2026.1.4", "packaging", "tyro", - "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,<=4.57.3", + "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,<=4.57.6", "datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0", "sentencepiece>=0.2.0", "tqdm", @@ -536,7 +536,7 @@ colab-new = [ colab-no-deps = [ "accelerate>=0.34.1", "trl>=0.18.2,!=0.19.0,<=0.24.0", - "peft>=0.7.1", + "peft>=0.18.0", "xformers ; ('linux' in sys_platform or sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "protobuf", diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py new file mode 100644 index 0000000000..9f2e8cda4e --- /dev/null +++ b/tests/test_raw_text.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +""" +Minimal test for raw text training implementation. +Tests basic functionality without heavy dependencies. +""" + +import sys +import os +import tempfile +from pathlib import Path +import importlib.util + + +# Mock the datasets module since it's not installed +class MockDataset: + def __init__(self, data_dict): + self.data = data_dict + self.column_names = list(data_dict.keys()) + + def __len__(self): + return len(next(iter(self.data.values()))) + + def __getitem__(self, idx): + if isinstance(idx, str): + # Allow accessing columns by name like dataset['text'] + return self.data[idx] + elif isinstance(idx, int): + # Allow accessing individual rows by index + return {key: values[idx] for key, values in self.data.items()} + else: + raise TypeError(f"Invalid index type: {type(idx)}") + + @classmethod + def from_dict(cls, data_dict): + return cls(data_dict) + + +# Mock datasets module +datasets_mock = type(sys)("datasets") +datasets_mock.Dataset = MockDataset +sys.modules["datasets"] = datasets_mock + +# Import the raw_text module directly to avoid unsloth/__init__.py dependencies +current_dir = os.path.dirname(__file__) +raw_text_path = os.path.join( + os.path.dirname(current_dir), "unsloth", "dataprep", "raw_text.py" +) + +spec = importlib.util.spec_from_file_location("raw_text", raw_text_path) +raw_text_module = importlib.util.module_from_spec(spec) +spec.loader.exec_module(raw_text_module) + +RawTextDataLoader = raw_text_module.RawTextDataLoader +TextPreprocessor = raw_text_module.TextPreprocessor + + +def test_raw_text_loader(): + """Test basic RawTextDataLoader functionality.""" + + # Mock tokenizer for testing + class MockTokenizer: + def __init__(self): + self.eos_token = "" + self.eos_token_id = 2 # Mock EOS token ID + + def __call__(self, text, return_tensors = None, add_special_tokens = False): + words = text.split() + token_ids = list(range(len(words))) + + if return_tensors == "pt": + # Mock tensor-like object + class MockTensor: + def __init__(self, data): + self.data = data + + def __getitem__(self, idx): + return self.data + + def __len__(self): + return len(self.data) + + def tolist(self): + return self.data + + return {"input_ids": [MockTensor(token_ids)]} + return {"input_ids": token_ids} + + def decode(self, token_ids, skip_special_tokens = False): + return " ".join([f"word_{i}" for i in token_ids]) + + # Create test file + test_content = "This is a test file for raw text training. " * 10 + with tempfile.NamedTemporaryFile(mode = "w", suffix = ".txt", delete = False) as f: + f.write(test_content) + test_file = f.name + + try: + # Test loader + tokenizer = MockTokenizer() + loader = RawTextDataLoader(tokenizer, chunk_size = 5, stride = 2) + + # Test loading with text output (legacy mode) + text_dataset = loader.load_from_file(test_file, return_tokenized = False) + assert len(text_dataset) > 0, "Should create at least one chunk" + assert "text" in text_dataset.column_names, "Dataset should have 'text' column" + + # Test loading with tokenized output (new efficient mode) + tokenized_dataset = loader.load_from_file(test_file, return_tokenized = True) + assert len(tokenized_dataset) > 0, "Should create at least one tokenized chunk" + assert ( + "input_ids" in tokenized_dataset.column_names + ), "Dataset should have 'input_ids' column" + assert ( + "attention_mask" in tokenized_dataset.column_names + ), "Dataset should have 'attention_mask' column" + + # Verify tokenized data structure + first_sample = tokenized_dataset[0] + assert isinstance(first_sample["input_ids"], list), "input_ids should be a list" + assert isinstance( + first_sample["attention_mask"], list + ), "attention_mask should be a list" + assert len(first_sample["input_ids"]) == len( + first_sample["attention_mask"] + ), "input_ids and attention_mask should have same length" + + # Verify labels field exists (for causal LM training) + assert ( + "labels" in tokenized_dataset.column_names + ), "Dataset should have 'labels' column" + assert ( + first_sample["labels"] == first_sample["input_ids"] + ), "labels should match input_ids" + + # Test constructor validation + try: + bad_loader = RawTextDataLoader(tokenizer, chunk_size = 0, stride = 2) + assert False, "Should raise ValueError for chunk_size=0" + except ValueError as e: + assert "chunk_size must be positive" in str(e) + + try: + bad_loader = RawTextDataLoader(tokenizer, chunk_size = 5, stride = 10) + assert False, "Should raise ValueError for stride >= chunk_size" + except ValueError as e: + assert "stride" in str(e) and "chunk_size" in str(e) + + # Test preprocessor + preprocessor = TextPreprocessor() + clean_text = preprocessor.clean_text(" messy text \n\n\n ") + assert "messy text" in clean_text, "Should clean text properly" + + # Test validation + stats = preprocessor.validate_dataset(text_dataset) + assert stats["total_samples"] > 0, "Should count samples" + assert "warnings" in stats, "Should include warnings" + + print("✅ All tests passed!") + return True + + except Exception as e: + print(f"❌ Test failed: {e}") + return False + + finally: + # Cleanup + os.unlink(test_file) + + +if __name__ == "__main__": + success = test_raw_text_loader() + sys.exit(0 if success else 1) diff --git a/tests/utils/test_qat.py b/tests/utils/test_qat.py index 79251cf2ff..1083712d78 100644 --- a/tests/utils/test_qat.py +++ b/tests/utils/test_qat.py @@ -4,12 +4,19 @@ from typing import Dict import pytest import torch -from torchao.quantization.qat import FakeQuantizedLinear -from torchao.quantization.qat.fake_quantizer import ( - FakeQuantizerBase, - Float8FakeQuantizer, - Int4WeightPreshuffledFakeQuantizer, -) + +try: + from torchao.quantization.qat import FakeQuantizedLinear + from torchao.quantization.qat.fake_quantizer import ( + FakeQuantizerBase, + Float8FakeQuantizer, + Int4WeightFakeQuantizer, + IntxFakeQuantizer, + ) +except ImportError: + print( + "Missing torchao import, please install or upgrade torchao with: pip install 'torchao>=0.15.0'" + ) class _CountingFakeQuantizer(torch.nn.Module): @@ -49,14 +56,20 @@ def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str): """ Verify that the given linear contains fake quantizers according to the `qat_scheme`. """ + weight_only = False if qat_scheme == "fp8-int4": act_fq_class = Float8FakeQuantizer - weight_fq_class = Int4WeightPreshuffledFakeQuantizer + weight_fq_class = Int4WeightFakeQuantizer min_in_features = 128 elif qat_scheme == "fp8-fp8": act_fq_class = Float8FakeQuantizer weight_fq_class = Float8FakeQuantizer min_in_features = -1 + elif qat_scheme == "int8": + act_fq_class = None + weight_fq_class = IntxFakeQuantizer + min_in_features = 128 + weight_only = True else: raise ValueError(f"Unknown qat_scheme: {qat_scheme}") @@ -64,7 +77,8 @@ def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str): base_layer = getattr(linear, "base_layer", linear) if base_layer.in_features >= min_in_features: assert isinstance(base_layer, FakeQuantizedLinear) - assert isinstance(base_layer.activation_fake_quantizer, act_fq_class) + if not weight_only: + assert isinstance(base_layer.activation_fake_quantizer, act_fq_class) assert isinstance(base_layer.weight_fake_quantizer, weight_fq_class) # Check lora A and B (only for full_finetuning=False) @@ -73,11 +87,13 @@ def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str): lora_B = linear.lora_B.default if lora_A.in_features >= min_in_features: assert isinstance(lora_A, FakeQuantizedLinear) - assert isinstance(lora_A.activation_fake_quantizer, act_fq_class) + if not weight_only: + assert isinstance(lora_A.activation_fake_quantizer, act_fq_class) assert isinstance(lora_A.weight_fake_quantizer, weight_fq_class) if lora_B.in_features >= min_in_features: assert isinstance(lora_B, FakeQuantizedLinear) - assert isinstance(lora_B.activation_fake_quantizer, act_fq_class) + if not weight_only: + assert isinstance(lora_B.activation_fake_quantizer, act_fq_class) assert isinstance(lora_B.weight_fake_quantizer, weight_fq_class) @@ -85,10 +101,12 @@ def _test_fake_quantizers_are_called( model: torch.nn.Module, example_inputs: Dict, full_finetuning: bool, + qat_scheme: str, ): """ Verify that the fake quantizers are actually called when the model is called. """ + weight_only = qat_scheme == "int8" def _swap_fake_quantizers(model: torch.nn.Module): for name, child in model.named_children(): @@ -99,7 +117,8 @@ def _test_fake_quantizers_are_called( for name, child in model.named_children(): if full_finetuning: if isinstance(child, FakeQuantizedLinear): - assert child.activation_fake_quantizer.count == 1 + if not weight_only: + assert child.activation_fake_quantizer.count == 1 assert child.weight_fake_quantizer.count == 1 else: # For LoRA, we only fake quantize the input activations once per block: @@ -107,12 +126,14 @@ def _test_fake_quantizers_are_called( # For mlp, we only fake quantize the gate_proj's input activations if name == "self_attn": base_layer = child.q_proj.base_layer - assert hasattr(base_layer, "activation_fake_quantizer") - assert base_layer.activation_fake_quantizer.count == 1 + if not weight_only: + assert hasattr(base_layer, "activation_fake_quantizer") + assert base_layer.activation_fake_quantizer.count == 1 elif name == "mlp": base_layer = child.gate_proj.base_layer - assert hasattr(base_layer, "activation_fake_quantizer") - assert base_layer.activation_fake_quantizer.count == 1 + if not weight_only: + assert hasattr(base_layer, "activation_fake_quantizer") + assert base_layer.activation_fake_quantizer.count == 1 elif isinstance(child, FakeQuantizedLinear): # Weight fake quantizers should always be called assert child.weight_fake_quantizer.count == 1 @@ -124,7 +145,7 @@ def _test_fake_quantizers_are_called( model.apply(_assert_fake_quantizers_are_called) -def _test_model_fake_quantize(qat_scheme: bool, full_finetuning: bool): +def _test_model_fake_quantize(qat_scheme: str, full_finetuning: bool): """ Test that all linear layers in the model are fake quantized according to the `qat_scheme`. """ @@ -141,16 +162,16 @@ def _test_model_fake_quantize(qat_scheme: bool, full_finetuning: bool): _test_linear_is_fake_quantized(layer.mlp.up_proj, qat_scheme) _test_linear_is_fake_quantized(layer.mlp.down_proj, qat_scheme) inputs = tokenizer("How are you?", return_tensors = "pt") - _test_fake_quantizers_are_called(model, inputs, full_finetuning) + _test_fake_quantizers_are_called(model, inputs, full_finetuning, qat_scheme) # TODO: there are bad interactions across tests right now, need to figure out # how to disable model caching before re-enabling this test -@pytest.mark.parametrize("qat_scheme", ["fp8-int4", "fp8-fp8"]) -def _test_full_model_fake_quantize(qat_scheme: bool): +@pytest.mark.parametrize("qat_scheme", ["fp8-int4", "fp8-fp8", "int8"]) +def _test_full_model_fake_quantize(qat_scheme: str): _test_model_fake_quantize(qat_scheme, full_finetuning = True) -@pytest.mark.parametrize("qat_scheme", ["fp8-int4", "fp8-fp8"]) -def test_lora_model_fake_quantize(qat_scheme: bool): +@pytest.mark.parametrize("qat_scheme", ["fp8-int4", "fp8-fp8", "int8"]) +def test_lora_model_fake_quantize(qat_scheme: str): _test_model_fake_quantize(qat_scheme, full_finetuning = False) diff --git a/unsloth-cli.py b/unsloth-cli.py index 0222afe0c7..612da11eb2 100644 --- a/unsloth-cli.py +++ b/unsloth-cli.py @@ -41,6 +41,7 @@ def run(args): from unsloth import is_bfloat16_supported from unsloth.models.loader_utils import prepare_device_map import logging + from unsloth import RawTextDataLoader logging.getLogger("hf-to-gguf").setLevel(logging.WARNING) @@ -99,15 +100,36 @@ def run(args): texts.append(text) return {"text": texts} - use_modelscope = strtobool(os.environ.get("UNSLOTH_USE_MODELSCOPE", "False")) - if use_modelscope: - from modelscope import MsDataset + def load_dataset_smart(args): + from transformers.utils import strtobool - dataset = MsDataset.load(args.dataset, split = "train") - else: - # Load and format dataset - dataset = load_dataset(args.dataset, split = "train") - dataset = dataset.map(formatting_prompts_func, batched = True) + if args.raw_text_file: + # Use raw text loader + loader = RawTextDataLoader(tokenizer, args.chunk_size, args.stride) + dataset = loader.load_from_file(args.raw_text_file) + elif args.dataset.endswith((".txt", ".md", ".json", ".jsonl")): + # Auto-detect local raw text files + loader = RawTextDataLoader(tokenizer) + dataset = loader.load_from_file(args.dataset) + else: + # Check for modelscope usage + use_modelscope = strtobool( + os.environ.get("UNSLOTH_USE_MODELSCOPE", "False") + ) + if use_modelscope: + from modelscope import MsDataset + + dataset = MsDataset.load(args.dataset, split = "train") + else: + # Existing HuggingFace dataset logic + dataset = load_dataset(args.dataset, split = "train") + + # Apply formatting for structured datasets + dataset = dataset.map(formatting_prompts_func, batched = True) + return dataset + + # Load dataset using smart loader + dataset = load_dataset_smart(args) print("Data is formatted and ready!") # Configure training arguments @@ -437,5 +459,15 @@ if __name__ == "__main__": help = "Token for pushing the model to Hugging Face hub", ) + parser.add_argument( + "--raw_text_file", type = str, help = "Path to raw text file for training" + ) + parser.add_argument( + "--chunk_size", type = int, default = 2048, help = "Size of text chunks for training" + ) + parser.add_argument( + "--stride", type = int, default = 512, help = "Overlap between chunks" + ) + args = parser.parse_args() run(args) diff --git a/unsloth/__init__.py b/unsloth/__init__.py index d10a0f8030..0b819a546e 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -30,16 +30,19 @@ from .import_fixes import ( check_fbgemm_gpu_version, torchvision_compatibility_check, fix_diffusers_warnings, + fix_huggingface_hub, ) fix_message_factory_issue() check_fbgemm_gpu_version() torchvision_compatibility_check() fix_diffusers_warnings() +fix_huggingface_hub() del fix_message_factory_issue del check_fbgemm_gpu_version del torchvision_compatibility_check del fix_diffusers_warnings +del fix_huggingface_hub # This check is critical because Unsloth optimizes these libraries by modifying # their code at import time. If they're imported first, the original (slower, @@ -76,7 +79,7 @@ from importlib.metadata import PackageNotFoundError # Check for unsloth_zoo try: unsloth_zoo_version = importlib_version("unsloth_zoo") - if Version(unsloth_zoo_version) < Version("2025.12.4"): + if Version(unsloth_zoo_version) < Version("2026.1.2"): print( "Unsloth: Please update Unsloth and Unsloth-Zoo to the latest version!\n" "Do this via `pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo`" @@ -123,6 +126,8 @@ from .import_fixes import ( fix_xformers_performance_issue, fix_vllm_aimv2_issue, fix_vllm_guided_decoding_params, + fix_vllm_pdl_blackwell, + fix_rocm_triton_key_error, ignore_logger_messages, patch_ipykernel_hf_xet, patch_trackio, @@ -130,11 +135,14 @@ from .import_fixes import ( patch_enable_input_require_grads, fix_openenv_no_vllm, fix_executorch, + patch_vllm_for_notebooks, ) fix_xformers_performance_issue() fix_vllm_aimv2_issue() fix_vllm_guided_decoding_params() +fix_vllm_pdl_blackwell() +fix_rocm_triton_key_error() ignore_logger_messages() patch_ipykernel_hf_xet() patch_trackio() @@ -142,10 +150,13 @@ patch_datasets() patch_enable_input_require_grads() fix_openenv_no_vllm() fix_executorch() +patch_vllm_for_notebooks() del fix_xformers_performance_issue del fix_vllm_aimv2_issue del fix_vllm_guided_decoding_params +del fix_vllm_pdl_blackwell +del fix_rocm_triton_key_error del ignore_logger_messages del patch_ipykernel_hf_xet del patch_trackio @@ -153,6 +164,7 @@ del patch_datasets del patch_enable_input_require_grads del fix_openenv_no_vllm del fix_executorch +del patch_vllm_for_notebooks # Torch 2.4 has including_emulation if DEVICE_TYPE == "cuda": @@ -273,6 +285,9 @@ from .save import * from .chat_templates import * from .tokenizer_utils import * from .trainer import * + +# Export dataprep utilities for CLI and downstream users +from .dataprep.raw_text import RawTextDataLoader, TextPreprocessor from unsloth_zoo.rl_environments import ( check_python_modules, create_locked_down_function, diff --git a/unsloth/dataprep/__init__.py b/unsloth/dataprep/__init__.py index b36122eb74..048f9b8010 100644 --- a/unsloth/dataprep/__init__.py +++ b/unsloth/dataprep/__init__.py @@ -13,3 +13,4 @@ # limitations under the License. from .synthetic import * +from .raw_text import * diff --git a/unsloth/dataprep/raw_text.py b/unsloth/dataprep/raw_text.py new file mode 100644 index 0000000000..ba010edabb --- /dev/null +++ b/unsloth/dataprep/raw_text.py @@ -0,0 +1,348 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import re +import json +import csv +from typing import List, Dict, Any, Union, Optional +from datasets import Dataset +from pathlib import Path + +__all__ = [ + "RawTextDataLoader", + "TextPreprocessor", +] + +SUPPORTED_FORMATS = { + ".txt": "plain_text", + ".md": "markdown", + ".json": "json_lines", + ".jsonl": "json_lines", + ".csv": "csv_text_column", +} + + +class RawTextDataLoader: + def __init__(self, tokenizer, chunk_size = 2048, stride = 512, return_tokenized = True): + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + if stride >= chunk_size: + raise ValueError( + f"stride ({stride}) must be smaller than chunk_size ({chunk_size})" + ) + self.tokenizer = tokenizer + self.chunk_size = chunk_size + self.stride = stride + self.return_tokenized = return_tokenized + + def detect_format(self, file_path): + """Auto-detect file format and parse accordingly""" + extension = Path(file_path).suffix.lower() + return SUPPORTED_FORMATS.get(extension, "plain_text") + + def load_from_file(self, file_path, return_tokenized = None): + """Load raw text and convert to dataset""" + if return_tokenized is None: + return_tokenized = self.return_tokenized + file_format = self.detect_format(file_path) + text_content = self._read_file_by_format(file_path, file_format) + if not text_content or not text_content.strip(): + raise ValueError(f"File '{file_path}' is empty or contains only whitespace") + chunks = self.smart_chunk_text( + text_content, self.chunk_size, self.stride, return_tokenized + ) + return self.create_causal_dataset(chunks) + + def load_from_files(self, file_paths, return_tokenized = None): + """Load multiple text files""" + if return_tokenized is None: + return_tokenized = self.return_tokenized + all_chunks = [] + for file_path in file_paths: + file_format = self.detect_format(file_path) + text_content = self._read_file_by_format(file_path, file_format) + chunks = self.smart_chunk_text( + text_content, self.chunk_size, self.stride, return_tokenized + ) + all_chunks.extend(chunks) + return self.create_causal_dataset(all_chunks) + + def chunk_text(self, text, return_tokenized = None): + """Split text into overlapping chunks""" + if return_tokenized is None: + return_tokenized = self.return_tokenized + return self.smart_chunk_text( + text, self.chunk_size, self.stride, return_tokenized + ) + + def create_causal_dataset(self, chunks): + """Create dataset for causal language modeling""" + if chunks and isinstance(chunks[0], dict): + # If chunks are already tokenized (dict with input_ids, attention_mask) + # Reorganize the data structure for Dataset.from_dict + input_ids = [chunk["input_ids"] for chunk in chunks] + attention_mask = [chunk["attention_mask"] for chunk in chunks] + # Labels are same as input_ids for causal LM training + labels = [list(ids) for ids in input_ids] + return Dataset.from_dict( + { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + } + ) + else: + # If chunks are text strings (backward compatibility) + return Dataset.from_dict({"text": chunks}) + + def smart_chunk_text(self, text, chunk_size, stride, return_tokenized = True): + """ + Intelligent chunking that: + 1. Respects sentence/paragraph boundaries + 2. Handles various text formats (.txt, .md, .json, etc.) + 3. Maintains context with stride overlap + 4. Returns tokenized chunks directly (more efficient) or text chunks + """ + # First pass: tokenize the entire text to get accurate token counts + tokenized = self.tokenizer(text, return_tensors = "pt", add_special_tokens = False) + tokens = tokenized["input_ids"] + + # Handle different tokenizer return formats + if hasattr(tokens, "__len__") and len(tokens) > 0: + # If it's a nested structure, get the first element + if hasattr(tokens[0], "__len__"): + tokens = tokens[0] + elif isinstance(tokens, int): + # If tokenizer returns just a count, create a simple range + tokens = list(range(tokens)) + + if len(tokens) <= chunk_size: + # Text is small enough to fit in one chunk + if return_tokenized: + # Add EOS token to the tokens if available + eos_token_id = getattr(self.tokenizer, "eos_token_id", None) + if eos_token_id is not None: + tokens = ( + tokens.tolist() if hasattr(tokens, "tolist") else list(tokens) + ) + tokens.append(eos_token_id) + + # Create attention mask + attention_mask = [1] * len(tokens) + return [{"input_ids": tokens, "attention_mask": attention_mask}] + else: + eos_token = self.tokenizer.eos_token if self.tokenizer.eos_token else "" + return [text + eos_token] + + chunks = [] + start_idx = 0 + + while start_idx < len(tokens): + # Calculate end index for this chunk + end_idx = min(start_idx + chunk_size, len(tokens)) + + # Extract tokens for this chunk + chunk_tokens = tokens[start_idx:end_idx] + + if return_tokenized: + # Convert to list if it's a tensor + chunk_tokens_list = ( + chunk_tokens.tolist() + if hasattr(chunk_tokens, "tolist") + else list(chunk_tokens) + ) + + # Add EOS token if it's the last chunk or chunk is complete + if end_idx == len(tokens) or len(chunk_tokens_list) == chunk_size: + eos_token_id = getattr(self.tokenizer, "eos_token_id", None) + if eos_token_id is not None: + chunk_tokens_list.append(eos_token_id) + + # Create attention mask (all tokens are attended to) + attention_mask = [1] * len(chunk_tokens_list) + + chunks.append( + {"input_ids": chunk_tokens_list, "attention_mask": attention_mask} + ) + else: + # Decode back to text (backward compatibility) + chunk_text = self.tokenizer.decode( + chunk_tokens, skip_special_tokens = True + ) + + # Add EOS token if it's the last chunk or chunk is complete + if end_idx == len(tokens) or len(chunk_tokens) == chunk_size: + eos_token = ( + self.tokenizer.eos_token if self.tokenizer.eos_token else "" + ) + chunk_text += eos_token + + chunks.append(chunk_text) + + # Move to next chunk with stride overlap + if end_idx == len(tokens): + break + start_idx += chunk_size - stride + + return chunks + + def _read_file_by_format(self, file_path, file_format): + """Read file content based on detected format.""" + with open(file_path, "r", encoding = "utf-8") as f: + if file_format == "plain_text" or file_format == "markdown": + return f.read() + elif file_format == "json_lines": + lines = [] + for line in f: + try: + data = json.loads(line.strip()) + text = self._extract_text_from_json(data) + if text: + lines.append(text) + except json.JSONDecodeError: + continue + return "\n\n".join(lines) + elif file_format == "csv_text_column": + reader = csv.DictReader(f) + texts = [] + for row in reader: + text = self._extract_text_from_csv_row(row) + if text: + texts.append(text) + return "\n\n".join(texts) + return "" + + def _extract_text_from_json(self, data): + """Extract text from JSON object using common field names.""" + text_fields = ["text", "content", "message", "body", "description", "prompt"] + for field in text_fields: + if field in data and isinstance(data[field], str): + return data[field] + return "" + + def _extract_text_from_csv_row(self, row): + """Extract text from CSV row using common column names.""" + text_columns = ["text", "content", "message", "body", "description", "prompt"] + for column in text_columns: + if column in row and row[column]: + return row[column] + return "" + + +class TextPreprocessor: + def clean_text(self, text): + """Remove unwanted characters, normalize whitespace""" + text = re.sub(r"\s+", " ", text) + text = re.sub(r"[^\x20-\x7E\n\t]", "", text) + text = text.replace("\r\n", "\n").replace("\r", "\n") + text = re.sub(r"\n{3,}", "\n\n", text) + return text.strip() + + def extract_sections(self, text, patterns): + """Extract specific sections (e.g., code blocks, quotes)""" + sections = [] + for pattern in patterns: + matches = re.findall(pattern, text, re.MULTILINE | re.DOTALL) + sections.extend(matches) + return sections + + def add_structure_tokens(self, text): + """Add special tokens for structure (chapters, sections)""" + text = re.sub( + r"^# (.+)$", r"<|chapter|>\1<|/chapter|>", text, flags = re.MULTILINE + ) + text = re.sub( + r"^## (.+)$", r"<|section|>\1<|/section|>", text, flags = re.MULTILINE + ) + text = re.sub( + r"^### (.+)$", r"<|subsection|>\1<|/subsection|>", text, flags = re.MULTILINE + ) + text = re.sub( + r"```(\w*)\n(.*?)\n```", r"<|code|\1|>\2<|/code|>", text, flags = re.DOTALL + ) + return text + + def validate_dataset(self, dataset): + """ + Check for: + - Minimum/maximum sequence lengths + - Character encoding issues + - Repeated content + - Empty chunks + """ + stats = { + "total_samples": len(dataset), + "empty_samples": 0, + "min_length": float("inf"), + "max_length": 0, + "avg_length": 0, + "repeated_content": 0, + "encoding_issues": 0, + "warnings": [], + } + + texts = dataset["text"] + text_lengths = [] + seen_texts = set() + + for i, text in enumerate(texts): + if not text or len(text.strip()) == 0: + stats["empty_samples"] += 1 + continue + + # Check for encoding issues + try: + text.encode("utf-8") + except UnicodeEncodeError: + stats["encoding_issues"] += 1 + + # Calculate lengths + length = len(text) + text_lengths.append(length) + stats["min_length"] = min(stats["min_length"], length) + stats["max_length"] = max(stats["max_length"], length) + + # Check for repeated content + text_hash = hash(text.strip()) + if text_hash in seen_texts: + stats["repeated_content"] += 1 + else: + seen_texts.add(text_hash) + + # Calculate average length + if text_lengths: + stats["avg_length"] = sum(text_lengths) / len(text_lengths) + stats["min_length"] = ( + stats["min_length"] if stats["min_length"] != float("inf") else 0 + ) + + # Generate warnings + if stats["empty_samples"] > 0: + stats["warnings"].append(f"Found {stats['empty_samples']} empty samples") + + if stats["repeated_content"] > 0: + stats["warnings"].append( + f"Found {stats['repeated_content']} repeated samples" + ) + + if stats["encoding_issues"] > 0: + stats["warnings"].append( + f"Found {stats['encoding_issues']} encoding issues" + ) + + if stats["min_length"] < 10: + stats["warnings"].append("Some samples are very short (< 10 characters)") + + return stats diff --git a/unsloth/device_type.py b/unsloth/device_type.py index 6a2d859d15..c7026ccf5f 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -24,6 +24,7 @@ __all__ = [ import torch import functools +import inspect from unsloth_zoo.utils import Version import inspect diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index efc7a7f4cd..89fd152857 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -20,6 +20,7 @@ from packaging.version import Version as TrueVersion import re import logging import textwrap +import warnings # We cannot do from unsloth_zoo.log import logger since FBGEMM might cause seg faults. UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") in ( @@ -93,10 +94,34 @@ class HidePrintMessage: if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") != "1": import sys - # Apply to stderr for FBGEMM + # Apply to stderr for FBGEMM and CUTLASS errors sys.stderr = HidePrintMessage(sys.stderr) # https://github.com/pytorch/FBGEMM/blob/d99cd96490ec4aabac2ee95b1e76ea4dcfcfa628/fbgemm_gpu/experimental/gemm/triton_gemm/utils.py#L43-L52 sys.stderr.add_filter("TMA benchmarks will be running") + # CUTLASS/FBGEMM MMA instruction error on SM90 vs SM100 (Blackwell) GPUs + # https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp + sys.stderr.add_filter("Arch conditional MMA instruction used without targeting") + # CUTLASS arch conditional errors for various architectures + sys.stderr.add_filter("CUTE_INVALID_CONTROL_PATH") + # CUTLASS TMA-related errors when not targeting correct architecture + sys.stderr.add_filter("Trying to use tma without CUTE_ARCH_TMA") + # Skipping import of cpp extensions due to incompatible torch version 2.9.0+cu128 for torchao version 0.15.0 + logging.getLogger("torchao").setLevel(logging.ERROR) + # Also filter torchao print to stderr about cpp extensions + sys.stderr.add_filter("Skipping import of cpp extensions") + # SyntaxWarning: invalid escape sequence '\.' + warnings.filterwarnings( + "ignore", message = "invalid escape sequence", category = SyntaxWarning + ) + # PYTORCH_CUDA_ALLOC_CONF is deprecated warning from torch + warnings.filterwarnings("ignore", message = "PYTORCH_CUDA_ALLOC_CONF is deprecated") + # TF32 precision deprecation warning from torch + warnings.filterwarnings( + "ignore", message = "Please use the new API settings to control TF32" + ) + # Deprecation warnings from torchao + warnings.filterwarnings("ignore", message = "`int4_weight_only` is deprecated") + warnings.filterwarnings("ignore", message = "`int8_weight_only` is deprecated") # Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype' @@ -179,6 +204,65 @@ def fix_xformers_performance_issue(): logger.info(f"Unsloth: Failed patching Xformers with error = {str(e)}") +def patch_vllm_for_notebooks(): + import sys + + ipython = None + try: + from IPython import get_ipython as _get_ipython + except Exception: + _get_ipython = None + + if _get_ipython is not None: + try: + ipython = _get_ipython() + except Exception: + ipython = None + + if ipython is None: + try: + import builtins + + _get_ipython = getattr(builtins, "get_ipython", None) + if callable(_get_ipython): + ipython = _get_ipython() + except Exception: + ipython = None + + if ipython is None: + return + + try: + shell = ipython.__class__.__name__ + is_notebook = shell == "ZMQInteractiveShell" or "google.colab" in str( + type(ipython) + ) + except Exception: + return + + if not is_notebook: + return + + if not hasattr(sys.stdout, "fileno"): + return + + needs_patch = False + try: + fd = sys.stdout.fileno() + if not isinstance(fd, int) or fd < 0: + needs_patch = True + except Exception: + needs_patch = True + + if not needs_patch: + return + + logger.info( + "Unsloth: Notebook detected - Patching sys.stdout.fileno for newer `vllm>=0.12.0` versions" + ) + sys.stdout.fileno = lambda: 1 + + # ValueError: 'aimv2' is already used by a Transformers config, pick another name. def fix_vllm_aimv2_issue(): spec = importlib.util.find_spec("vllm") @@ -223,16 +307,43 @@ def fix_vllm_aimv2_issue(): def fix_vllm_guided_decoding_params(): + def _maybe_raise_vllm_transformers_mismatch(error): + error_text = str(error) + if ( + "ALLOWED_LAYER_TYPES" in error_text + or "transformers.configuration_utils" in error_text + ): + try: + vllm_version = importlib_version("vllm") + except Exception: + vllm_version = "unknown" + raise RuntimeError( + "Unsloth: vLLM with version " + f"{vllm_version} does not yet support transformers>=5.0.0. " + "Please downgrade to transformers==4.57.3 via " + 'pip install --force-reinstall "transformers==4.57.3". ' + f"Original error: {error}" + ) from error + if importlib.util.find_spec("vllm") is None: return # GuidedDecodingParmas is renamed to StructuredOutputsParams in vLLM # https://github.com/vllm-project/vllm/pull/22772/files # trl still wants to use GuidedDecodingParams. This is a temporary patch till trl updates - import vllm + try: + import vllm + except ImportError as e: + _maybe_raise_vllm_transformers_mismatch(e) + raise try: from vllm.sampling_params import GuidedDecodingParams - except ImportError: + except ImportError as e: + _maybe_raise_vllm_transformers_mismatch(e) + if not hasattr(vllm, "sampling_params") or not hasattr( + vllm.sampling_params, "StructuredOutputsParams" + ): + raise vllm.sampling_params.GuidedDecodingParams = ( vllm.sampling_params.StructuredOutputsParams ) @@ -316,10 +427,14 @@ def check_fbgemm_gpu_version(): except: return # We noticed some SegFault or bad alloc errors on lower versions of fbgemm_gpu. + # Instead of raising an error, disable FBGEMM and fall back to Triton kernels. if Version(fbgemm_gpu_version) < Version("1.4.0"): - raise ImportError( - f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected. It might cause unexpected issues like segmentation faults. Please uninstall the current one by doing `pip uninstall fbgemm-gpu` && `pip install fbgemm-gpu` to install fbgemm-gpu 1.4.0 or newer!" + os.environ["UNSLOTH_HAS_FBGEMM"] = "0" + logger.info( + f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} is old and may cause issues. " + f"Disabling FBGEMM - using Triton kernels instead." ) + return logger.info(f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected.") @@ -539,3 +654,161 @@ def fix_executorch(): def fix_diffusers_warnings(): # Silence Flax classes are deprecated and will be removed in Diffusers v1.0.0. os.environ["DIFFUSERS_VERBOSITY"] = "error" + + +def fix_huggingface_hub(): + # huggingface_hub.is_offline_mode got removed, so add it back + import huggingface_hub + + if not hasattr(huggingface_hub, "is_offline_mode"): + huggingface_hub.is_offline_mode = ( + lambda: huggingface_hub.constants.HF_HUB_OFFLINE + ) + + +def fix_rocm_triton_key_error(): + """ + ROCm + torch.compile can fail if Triton lacks `triton_key`. + Disable Inductor/compile only on ROCm when that symbol is missing. + """ + try: + import torch + except (ImportError, ModuleNotFoundError): + return + + if not getattr(torch.version, "hip", None): + return + + try: + import triton + except (ImportError, ModuleNotFoundError): + return + + try: + from triton.runtime import triton_key # noqa: F401 + + return + except ImportError: + pass + + os.environ.setdefault("TORCHINDUCTOR_DISABLE", "1") + os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + logger.info( + "Unsloth: ROCm detected and Triton lacks triton_key; " + "disabling torch.compile/Inductor to avoid backend crash." + ) + + +def fix_vllm_pdl_blackwell(): + """ + Fix vLLM PDL (Programmatic Dependent Launch) bug on Blackwell GPUs (SM100). + + The issue: vLLM's LoRA Triton kernels use tl.extra.cuda.gdc_wait() for PDL + optimization on SM90+ GPUs. This fails on SM100 (B200/B100) during CUDA graph + capture because Triton's pipeliner can't handle gdc_wait in complex kernels. + + See: https://github.com/vllm-project/vllm/issues/30872 + """ + if importlib.util.find_spec("vllm") is None: + return + + # Check if any CUDA GPU is SM100 (Blackwell) + try: + import torch + + if not torch.cuda.is_available(): + return + + # Scan all GPUs for SM100 - fix applies globally via env var and monkey-patch + has_sm100 = False + sm100_gpu_name = None + for i in range(torch.cuda.device_count()): + major, minor = torch.cuda.get_device_capability(i) + if major == 10: + has_sm100 = True + sm100_gpu_name = torch.cuda.get_device_name(i) + break + + if not has_sm100: + return + except Exception: + return + + # Helper to check if module spec exists + def _spec_exists(name): + try: + return importlib.util.find_spec(name) is not None + except (ModuleNotFoundError, ValueError): + return False + + # Check if vLLM has the PDL-related modules before doing internet check + has_utils = _spec_exists("vllm.lora.ops.triton_ops.utils") + has_expand_op = _spec_exists("vllm.lora.ops.triton_ops.lora_expand_op") + has_shrink_op = _spec_exists("vllm.lora.ops.triton_ops.lora_shrink_op") + + if not has_utils and not has_expand_op and not has_shrink_op: + # Old vLLM version without PDL support - nothing to patch + return + + # Check if vLLM version includes the fix + VLLM_PDL_FIX_VERSION = "0.13.2" + try: + vllm_version = Version(importlib_version("vllm")) + if vllm_version > Version(VLLM_PDL_FIX_VERSION): + logger.info( + f"Unsloth: SM100 ({sm100_gpu_name}) detected but vLLM {vllm_version} " + f"should include PDL fix - skipping workaround" + ) + return + except Exception as e: + logger.debug( + f"Unsloth: vLLM version check failed ({e}), applying PDL workaround." + ) + + # Apply the PDL fix + os.environ["TRITON_DISABLE_PDL"] = "1" + + def fake_supports_pdl(*args, **kwargs): + return False + + patched = [] + + # First, patch the source module (utils.py) where supports_pdl is defined. + # This is critical because supports_pdl uses @lru_cache - we must clear the + # cache to prevent stale cached results from the original function. + try: + utils_module = importlib.import_module("vllm.lora.ops.triton_ops.utils") + if hasattr(utils_module, "supports_pdl"): + original_fn = utils_module.supports_pdl + if hasattr(original_fn, "cache_clear"): + original_fn.cache_clear() + utils_module.supports_pdl = fake_supports_pdl + patched.append("utils") + except (ImportError, ModuleNotFoundError, AttributeError): + pass + + # Also patch the consumer modules that import supports_pdl from utils. + # This ensures the patched function is used even if the module was already + # imported before this fix runs. + consumer_modules = { + "lora_expand_op": "vllm.lora.ops.triton_ops.lora_expand_op", + "lora_shrink_op": "vllm.lora.ops.triton_ops.lora_shrink_op", + "fused_moe_lora_op": "vllm.lora.ops.triton_ops.fused_moe_lora_op", + } + for name, path in consumer_modules.items(): + try: + module = importlib.import_module(path) + if hasattr(module, "supports_pdl"): + module.supports_pdl = fake_supports_pdl + patched.append(name) + except (ImportError, ModuleNotFoundError, AttributeError): + pass + + if patched: + logger.info( + f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - " + f"patched: {', '.join(patched)}" + ) + else: + # Just set the env var - vLLM might be an older version without supports_pdl + logger.info(f"Unsloth: Set TRITON_DISABLE_PDL=1 for SM100 ({sm100_gpu_name})") diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 912e6f7e3f..fbb14013ff 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -24,7 +24,7 @@ from .utils import ( is_cdna, ) from transformers.models.llama.modeling_llama import logger -from packaging.version import Version +from unsloth_zoo.utils import Version from unsloth_zoo.loss_utils import ( patch_loss_functions as _patch_loss_functions, diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 60d0c318c3..f1c0e298d9 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -379,9 +379,22 @@ class LoRA_QKV(torch.autograd.Function): ): dtype = X.dtype - Q = matmul_lora(X, QW, QW_quant, QA, QB, QS) - K = matmul_lora(X, KW, KW_quant, KA, KB, KS) - V = matmul_lora(X, VW, VW_quant, VA, VB, VS) + # bitsandbytes 8-bit matmul expects 2D inputs. + # TorchInductor/AOTAutograd fails on 3D tensors during backward, + # so we explicitly flatten the sequence dimension. + orig_shape = X.shape + X_for_matmul = X + if X.dim() == 3: + X_for_matmul = X.view(-1, X.shape[-1]) + Q = matmul_lora(X_for_matmul, QW, QW_quant, QA, QB, QS) + K = matmul_lora(X_for_matmul, KW, KW_quant, KA, KB, KS) + V = matmul_lora(X_for_matmul, VW, VW_quant, VA, VB, VS) + + # Restore original shape after matmul + if len(orig_shape) == 3: + Q = Q.view(orig_shape[0], orig_shape[1], -1) + K = K.view(orig_shape[0], orig_shape[1], -1) + V = V.view(orig_shape[0], orig_shape[1], -1) ctx.custom_saved_tensors = ( QW, diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index 3093bf61b1..e9f9161709 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -523,6 +523,7 @@ def fp8_fbgemm_block_linear(X, weight, weight_scale, bias = None): def test_has_fbgemm(): # We must manually check if the faster FBGEMM works on the specific GPU # For example RTX 5090 and RTX 4090 does not work + # Also SM100 (Blackwell B200/B100) GPUs fail with CUTLASS SM90 kernels # [TODO] Investigate with TorchAO why FBGEMM fails on consumer GPUs M, N, K = 128, 128, 128 xq = torch.ones(M, K, dtype = torch.float8_e4m3fn, device = "cuda") @@ -537,10 +538,25 @@ def test_has_fbgemm(): has_fbgemm = True del out except Exception as e: - e = str(e) - if "cutlass cannot initialize" in e.lower(): + error_str = str(e).lower() + # Catch any CUTLASS/CUDA errors and disable FBGEMM + # This includes MMA instruction errors, architecture mismatches, kernel launch failures, etc. + cutlass_cuda_errors = ( + "cutlass", + "cuda error", + "cuda runtime error", + "no kernel image", + "arch conditional", + "mma instruction", + "compute capability", + "cute_invalid_control_path", + "tma", + ) + is_cutlass_cuda_error = any(err in error_str for err in cutlass_cuda_errors) + + if is_cutlass_cuda_error: print( - f"Unsloth: FBGEMM on the current GPU cannot load - will switch to Triton kernels" + "Unsloth: FBGEMM on the current GPU cannot load - will switch to Triton kernels" ) else: print( diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index a032e0f7fc..fcc9cb923b 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/unsloth/kernels/rope_embedding.py @@ -312,8 +312,8 @@ class Fast_RoPE_Embedding_QK(torch.autograd.Function): _, n_heads_K, _, _ = K.shape # Inplace rotary embedding is generally fine - Q_out = Q.clone() if not Q.is_contiguous else Q - K_out = K.clone() if not K.is_contiguous else K + Q_out = Q.clone() if not Q.is_contiguous() else Q + K_out = K.clone() if not K.is_contiguous() else K if has_indices: # TRL's rotary indices are always in int32, so casting is just for safety @@ -383,21 +383,21 @@ class Fast_RoPE_Embedding_QK(torch.autograd.Function): else ctx.cos.new_empty(1, dtype = torch.int32) ) + # Inplace rotary embedding is generally fine + dQ_out = dQ.clone() if not dQ.is_contiguous() else dQ + dK_out = dK.clone() if not dK.is_contiguous() else dK + Q_batch_stride, Q_head_stride, Q_seq_stride = ( - dQ.stride(0), - dQ.stride(1), - dQ.stride(2), + dQ_out.stride(0), + dQ_out.stride(1), + dQ_out.stride(2), ) K_batch_stride, K_head_stride, K_seq_stride = ( - dK.stride(0), - dK.stride(1), - dK.stride(2), + dK_out.stride(0), + dK_out.stride(1), + dK_out.stride(2), ) - # Inplace rotary embedding is generally fine - dQ_out = dQ.clone() if not dQ.is_contiguous else dQ - dK_out = dK.clone() if not dK.is_contiguous else dK - with torch_gpu_device(dQ.device): _rope_embedding_QK[(batch * ctx.seq_len, ctx.n_heads_Q)]( dQ_out, diff --git a/unsloth/kernels/swiglu.py b/unsloth/kernels/swiglu.py index b321f5179e..b3ae9d40e6 100644 --- a/unsloth/kernels/swiglu.py +++ b/unsloth/kernels/swiglu.py @@ -128,7 +128,7 @@ def _DWf_DW_dfg_kernel( def swiglu_DWf_DW_dfg_kernel(DW, e, g): - batch_seq_len, hd = e.shape + batch_seq_len, hd = e.shape # Flattened to 2D, so 1st dim is bsz * seq_len n_elements = e.numel() grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) with torch_gpu_device(e.device): diff --git a/unsloth/models/__init__.py b/unsloth/models/__init__.py index d7b2393c89..138f309032 100644 --- a/unsloth/models/__init__.py +++ b/unsloth/models/__init__.py @@ -19,6 +19,7 @@ from .qwen2 import FastQwen2Model from .qwen3 import FastQwen3Model from .qwen3_moe import FastQwen3MoeModel from .granite import FastGraniteModel +from .sentence_transformer import FastSentenceTransformer try: from .falcon_h1 import FastFalconH1Model diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index abc8380562..76952b00a5 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.12.9" +__version__ = "2026.1.4" __all__ = [ "SUPPORTS_BFLOAT16", @@ -73,6 +73,7 @@ __all__ = [ "verify_fp8_support_if_applicable", "_get_inference_mode_context_manager", "hf_login", + "make_fast_generate_wrapper", ] import torch @@ -174,6 +175,8 @@ warnings.filterwarnings(action = "ignore", category = UserWarning, module = "bit # Stop "Special tokens have been added in the vocabulary, ..." logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITICAL + 1) +TORCHAO_MSG = "Error: torchao not found, please install with `pip install torchao`" + # Ignore logging messages class HideLoggingMessage(logging.Filter): @@ -1105,53 +1108,66 @@ def _get_statistics(statistics = None, force_download = True): global USE_MODELSCOPE USE_MODELSCOPE = os.environ.get("UNSLOTH_USE_MODELSCOPE", "0") == "1" - if statistics is not None: - pass - elif "\nCOLAB_" in keynames and n_cpus == 1: - statistics = "colab" - elif "\nCOLAB_" in keynames: - statistics = "colabpro" - elif "\nKAGGLE_" in keynames: - statistics = "kaggle" - elif "\nRUNPOD_" in keynames: - statistics = "runpod" - elif "\nAWS_" in keynames: - statistics = "aws" - elif "\nAZURE_" in keynames: - statistics = "azure" - # elif "\nK_" in keynames or "\nFUNCTION_" in keynames: statistics = "gcp" - elif "\nINVOCATION_ID" in keynames: - statistics = "lambda" - # else: statistics = "other" - else: - - def try_vllm_check(): - vendor_files = ( - "/sys/class/dmi/id/product_version", - "/sys/class/dmi/id/bios_vendor", - "/sys/class/dmi/id/product_name", - "/sys/class/dmi/id/chassis_asset_tag", - "/sys/class/dmi/id/sys_vendor", - ) + if statistics is None: + # Prefer filesystem markers (harder to misidentify) before env-key matching + try: from pathlib import Path - for vendor_file in vendor_files: - path = Path(vendor_file) - if path.is_file(): - file_content = path.read_text().lower() - if "amazon" in file_content: - return "aws" - elif "microsoft corporation" in file_content: - return "azure" - elif "google" in file_content: - return "gcp" - return "other" + if Path("/kaggle/working").exists(): + statistics = "kaggle" + elif Path("/content").exists() and Path("/opt/colab").exists(): + statistics = "colab" if n_cpus == 1 else "colabpro" + elif Path("/runpod-volume").exists(): + statistics = "runpod" + except Exception: + pass + + # Fallback to env-key detection + if statistics is None: + if "\nKAGGLE_" in keynames: + statistics = "kaggle" + elif "\nCOLAB_" in keynames and n_cpus == 1: + statistics = "colab" + elif "\nCOLAB_" in keynames: + statistics = "colabpro" + elif "\nRUNPOD_" in keynames: + statistics = "runpod" + elif "\nAWS_" in keynames: + statistics = "aws" + elif "\nAZURE_" in keynames: + statistics = "azure" + # elif "\nK_" in keynames or "\nFUNCTION_" in keynames: statistics = "gcp" + elif "\nINVOCATION_ID" in keynames: + statistics = "lambda" + # else: statistics = "other" + else: + + def try_vllm_check(): + vendor_files = ( + "/sys/class/dmi/id/product_version", + "/sys/class/dmi/id/bios_vendor", + "/sys/class/dmi/id/product_name", + "/sys/class/dmi/id/chassis_asset_tag", + "/sys/class/dmi/id/sys_vendor", + ) + + for vendor_file in vendor_files: + path = Path(vendor_file) + if path.is_file(): + file_content = path.read_text().lower() + if "amazon" in file_content: + return "aws" + elif "microsoft corporation" in file_content: + return "azure" + elif "google" in file_content: + return "gcp" + return "other" + + try: + statistics = try_vllm_check() + except Exception: + statistics = "other" - pass - try: - statistics = try_vllm_check() - except: - statistics = "other" if statistics is not None: import tempfile from huggingface_hub import snapshot_download @@ -1183,7 +1199,7 @@ def _get_statistics(statistics = None, force_download = True): "model = FastLanguageModel.from_pretrained('unsloth/gpt-oss-20b')\n" "```" ) - except: + except Exception: # Try no time limit check stats_check() @@ -1196,7 +1212,10 @@ def get_statistics(local_files_only = False): # You can disable this by setting UNSLOTH_DISABLE_STATISTICS import os - if "UNSLOTH_DISABLE_STATISTICS" in os.environ: + if ( + "UNSLOTH_DISABLE_STATISTICS" in os.environ + or os.environ.get("UNSLOTH_USE_MODELSCOPE", "0") == "1" + ): return if local_files_only: return @@ -1981,9 +2000,10 @@ def validate_loftq_config(loftq_config, lora_dropout, bias, init_lora_weights, m type(init_lora_weights) is bool or init_lora_weights == "gaussian" or init_lora_weights == "loftq" + or init_lora_weights == "corda" ): raise ValueError( - 'Unsloth: `init_lora_weights` must be either [True, False, "gaussian", "loftq"].' + 'Unsloth: `init_lora_weights` must be either [True, False, "gaussian", "loftq", "corda"].' ) if init_lora_weights == "loftq": @@ -2193,15 +2213,32 @@ def _prepare_model_for_qat( QAT can be optionally combined with LoRA fine-tuning to for additional throughput improvement. For more details: https://dev-discuss.pytorch.org/t/speeding-up-qat-by-1-89x-with-lora/2700 """ - from torchao.quantization import PerRow, quantize_ - from torchao.quantization.granularity import PerGroup, PerAxis - from torchao.quantization.qat import QATConfig + try: + from torchao.quantization import PerRow, quantize_ + from torchao.quantization.granularity import PerGroup, PerAxis + from torchao.quantization.qat import QATConfig + except ImportError: + raise ImportError(TORCHAO_MSG) + + # Gemma3 models have issues with int8 embedding quantization due to their + # large vocabulary size (262144). Auto-switch to int4 weight-only instead. + if qat_scheme == "int8-int4": + model_types = get_transformers_model_type(model.config) + is_gemma3 = any("gemma3" in mt or "gemma_3" in mt for mt in model_types) + if is_gemma3: + print( + "Unsloth: Gemma3 has a large vocabulary causing int8 embedding issues. " + "Switching to int4 weight-only QAT for training stability." + ) + qat_scheme = "int4" if not isinstance(qat_scheme, TorchAOConfig): torchao_config: Optional[TorchAOConfig] = None if qat_scheme == "fp8-int4": - from torchao.quantization import Float8DynamicActivationInt4WeightConfig - + try: + from torchao.quantization import Float8DynamicActivationInt4WeightConfig + except ImportError: + raise ImportError(TORCHAO_MSG) group_size = 128 base_config = Float8DynamicActivationInt4WeightConfig() filter_fn = ( @@ -2213,8 +2250,12 @@ def _prepare_model_for_qat( base_config_and_filter_fns = [(base_config, filter_fn)], ) elif qat_scheme == "fp8-fp8": - from torchao.quantization import Float8DynamicActivationFloat8WeightConfig - + try: + from torchao.quantization import ( + Float8DynamicActivationFloat8WeightConfig, + ) + except ImportError: + raise ImportError(TORCHAO_MSG) base_config = Float8DynamicActivationFloat8WeightConfig( granularity = PerRow() ) @@ -2222,11 +2263,13 @@ def _prepare_model_for_qat( qat_scheme = qat_scheme, base_config_and_filter_fns = [(base_config, None)] ) elif qat_scheme == "int8-int4": - from torchao.quantization import ( - Int8DynamicActivationIntxWeightConfig, - IntxWeightOnlyConfig, - ) - + try: + from torchao.quantization import ( + Int8DynamicActivationIntxWeightConfig, + IntxWeightOnlyConfig, + ) + except ImportError: + raise ImportError(TORCHAO_MSG) torchao_config = TorchAOConfig( qat_scheme = qat_scheme, base_config_and_filter_fns = [ @@ -2246,8 +2289,10 @@ def _prepare_model_for_qat( prequantization_transform = _untie_input_output_embeddings, ) elif qat_scheme == "int4": - from torchao.quantization import Int4WeightOnlyConfig - + try: + from torchao.quantization import Int4WeightOnlyConfig + except ImportError: + raise ImportError(TORCHAO_MSG) group_size = 128 base_config = Int4WeightOnlyConfig(group_size = group_size) filter_fn = ( @@ -2258,6 +2303,22 @@ def _prepare_model_for_qat( qat_scheme = qat_scheme, base_config_and_filter_fns = [(base_config, filter_fn)], ) + elif qat_scheme == "int8": + try: + from torchao.quantization import IntxWeightOnlyConfig + from torchao.quantization.granularity import PerAxis + except ImportError: + raise ImportError(TORCHAO_MSG) + + base_config = IntxWeightOnlyConfig( + weight_dtype = torch.int8, + granularity = PerAxis(0), + ) + filter_fn = lambda m, _: isinstance(m, torch.nn.Linear) + torchao_config = TorchAOConfig( + qat_scheme = qat_scheme, + base_config_and_filter_fns = [(base_config, filter_fn)], + ) else: raise ValueError(f"Unexpected QAT scheme {qat_scheme}") assert torchao_config is not None, f"TorchAOConfig was not set for {qat_scheme}" @@ -2365,3 +2426,59 @@ def hf_login(token: Optional[str] = None) -> Optional[str]: except Exception as e: logger.info(f"Failed to login to huggingface using token with error: {e}") return token + + +def make_fast_generate_wrapper(original_generate): + """ + Creates a wrapper around model.generate that checks for incorrect + vLLM-style usage when fast_inference=False. + """ + + @functools.wraps(original_generate) + def _fast_generate_wrapper(*args, **kwargs): + # Check for vLLM-specific arguments + if "sampling_params" in kwargs: + raise ValueError( + "Unsloth: `sampling_params` is only supported when `fast_inference=True` (vLLM). " + "Since `fast_inference=False`, use HuggingFace generate arguments instead:\n" + " model.fast_generate(**tokens.to('cuda'), max_new_tokens=64, temperature=1.0, top_p=0.95)" + ) + + if "lora_request" in kwargs: + raise ValueError( + "Unsloth: `lora_request` is only supported when `fast_inference=True` (vLLM). " + "Since `fast_inference=False`, LoRA weights are already merged into the model." + ) + + # Check if first positional argument is a string or list of strings + if len(args) > 0: + first_arg = args[0] + is_string_input = False + + if isinstance(first_arg, str): + is_string_input = True + elif isinstance(first_arg, (list, tuple)) and len(first_arg) > 0: + if isinstance(first_arg[0], str): + is_string_input = True + + if is_string_input: + raise ValueError( + "Unsloth: Passing text strings to `fast_generate` is only supported " + "when `fast_inference=True` (vLLM). Since `fast_inference=False`, you must " + "tokenize the input first:\n\n" + " messages = tokenizer.apply_chat_template(\n" + ' [{"role": "user", "content": "Your prompt here"}],\n' + " tokenize=True, add_generation_prompt=True,\n" + ' return_tensors="pt", return_dict=True\n' + " )\n" + " output = model.fast_generate(\n" + " **messages.to('cuda'),\n" + " max_new_tokens=64,\n" + " temperature=1.0,\n" + " )" + ) + + # Call original generate + return original_generate(*args, **kwargs) + + return _fast_generate_wrapper diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index a091a0173f..c33317ee02 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -15,7 +15,7 @@ from .llama import * from ._utils import __version__ from unsloth_zoo.hf_utils import dtype_from_config -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from ..utils.packing import get_packed_info_from_kwargs from ..utils.attention_dispatch import ( AttentionConfig, @@ -35,8 +35,6 @@ try: repeat_kv, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.42"): raise ImportError( @@ -344,8 +342,8 @@ def CohereAttention_fast_forward_inference( Kn = Kn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2) Vn = Vn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2) if self.use_qk_norm: - Q = fast_layernorm_inference(self.q_norm, Q, self.q_norm_out_weight) - K = fast_layernorm_inference(self.k_norm, K, self.k_norm_out_weight) + Qn = fast_layernorm_inference(self.q_norm, Qn, self.q_norm_out_weight) + Kn = fast_layernorm_inference(self.k_norm, Kn, self.k_norm_out_weight) # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len) # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids) @@ -479,7 +477,7 @@ def CohereModel_fast_forward_inference( ) ) - hidden_states_mlp = fast_swiglu_inference(self.mlp, hidden_states) + hidden_states_mlp = fast_swiglu_inference(decoder_layer.mlp, hidden_states) residual += hidden_states_attention residual += hidden_states_mlp hidden_states = residual diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index fc5ea458a6..428f49d727 100644 --- a/unsloth/models/falcon_h1.py +++ b/unsloth/models/falcon_h1.py @@ -456,9 +456,9 @@ def FalconH1DecoderLayer_fast_forward( # Fully Connected residual = hidden_states hidden_states = fast_rms_layernorm_inference( - self.post_attention_layernorm, hidden_states + self.pre_ff_layernorm, hidden_states ) - hidden_states = fast_swiglu_inference(self.mlp, hidden_states) + hidden_states = fast_swiglu_inference(self.feed_forward, hidden_states) hidden_states += residual else: residual = hidden_states diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 291d442673..1789a9cd92 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -14,7 +14,7 @@ from .llama import * from ._utils import __version__ -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from unsloth_zoo.hf_utils import dtype_from_config from ..utils.packing import ( build_sdpa_packed_attention_mask, @@ -34,8 +34,6 @@ try: repeat_kv, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.38"): raise ImportError( diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 4b2503b8a1..16d04955d3 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -14,7 +14,7 @@ from .llama import * from ._utils import __version__ -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from unsloth_zoo.hf_utils import dtype_from_config from ..utils.packing import get_packed_info_from_kwargs from ..utils.attention_dispatch import ( @@ -41,8 +41,6 @@ try: repeat_kv, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.42"): raise ImportError( diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index 2632ab6914..aae746aed1 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -15,7 +15,7 @@ from .llama import * import os from ._utils import __version__ -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from unsloth_zoo.hf_utils import dtype_from_config from ..utils.packing import get_packed_info_from_kwargs from ..utils.attention_dispatch import ( @@ -41,14 +41,12 @@ try: GraniteForCausalLM, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.45.0"): raise ImportError( - f"Unsloth: Your transformers version of {transformers_version} does not support Gemma2.\n" - f"The minimum required version is 4.42.3.\n" - f'Try `pip install --upgrade "transformers>=4.42.3"`\n' + f"Unsloth: Your transformers version of {transformers_version} does not support Granite.\n" + f"The minimum required version is 4.45.0.\n" + f'Try `pip install --upgrade "transformers>=4.45.0"`\n' f"to obtain the latest transformers build, then restart this session." ) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 1d7695b9aa..39f2ba1460 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -146,6 +146,59 @@ torch_nn_functional_softmax = torch.nn.functional.softmax # SDPA has GQA internally SDPA_HAS_GQA = "enable_gqa" in scaled_dot_product_attention.__doc__ +from peft.utils.other import ModulesToSaveWrapper + + +def _offload_frozen_module_for_training( + module: ModulesToSaveWrapper, + device_type: str, + offload_device: str = "cpu", +) -> None: + """ + Offload frozen module to CPU and configure trainable copy for mixed precision training. + + This function optimizes memory usage by: + 1. Moving the trainable copy to the target device with appropriate precision + 2. Offloading the original frozen module to CPU/disk to free VRAM + 3. Converting float16 to float32 for compatibility with certain GPUs (e.g., Tesla T4) + + Args: + module: The module to configure. Must be a ModulesToSaveWrapper with a + `modules_to_save` attribute containing trainable and original modules. + device_type: Target device string for training (e.g., "cuda:0", "xpu:0") + offload_device: Device to offload frozen parameters (default: "cpu") + Note: Currently only "cpu" is supported; disk offloading is planned. + + Returns: + None (modifies module in-place) + + Note: + - Float16 weights are automatically promoted to float32 for GPU compatibility + - Original frozen parameters are moved to CPU to reduce active VRAM usage + - Future versions will support disk-based offloading for even larger models + + See Also: + - https://github.com/unslothai/unsloth/pull/1200 (Tesla T4 float32 requirement) + """ + # Early return with explicit None if module doesn't support mixed precision training + if not hasattr(module, "modules_to_save"): + return None + + new_dtype = module.modules_to_save.default.weight.dtype + if new_dtype == torch.float16: + # See https://github.com/unslothai/unsloth/pull/1200 + # Tesla T4 must use float32 and not float16 + new_dtype = torch.float32 + + module.modules_to_save.default.to( + device = device_type, dtype = new_dtype, non_blocking = True + ) + module.modules_to_save.default.requires_grad_(True) + + # [TODO] Move old module to CPU - should be disk! + module.original_module.to(device = offload_device, non_blocking = True) + module.original_module.requires_grad_(False) + # Fix new HF's inference code def _fast_prepare_inputs_for_generation( @@ -2326,7 +2379,7 @@ class FastLlamaModel: attn_implementation = "eager", **kwargs, ) - model.fast_generate = model.generate + model.fast_generate = make_fast_generate_wrapper(model.generate) model.fast_generate_batches = None else: from unsloth_zoo.vllm_utils import ( @@ -2600,6 +2653,7 @@ class FastLlamaModel: loftq_config = {}, temporary_location = "_unsloth_temporary_saved_buffers", qat_scheme = None, + ensure_weight_tying = False, **kwargs, ): if os.environ.get("UNSLOTH_USE_NEW_MODEL", "0") == "1": @@ -2629,6 +2683,7 @@ class FastLlamaModel: init_lora_weights = init_lora_weights, loftq_config = loftq_config, temporary_location = temporary_location, + ensure_weight_tying = ensure_weight_tying, **kwargs, ) if os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1": @@ -2709,46 +2764,16 @@ class FastLlamaModel: "Unsloth: Training embed_tokens in mixed precision to save VRAM" ) - new_dtype = model.get_input_embeddings().modules_to_save.default.weight.dtype - if new_dtype == torch.float16: - # See https://github.com/unslothai/unsloth/pull/1200 - # Tesla T4 must use float32 and not float16 - new_dtype = torch.float32 - - model.get_input_embeddings().modules_to_save.default.to( - device = DEVICE_TYPE_TORCH, dtype = new_dtype, non_blocking = True + _offload_frozen_module_for_training( + model.get_input_embeddings(), DEVICE_TYPE_TORCH ) - model.get_input_embeddings().modules_to_save.default.requires_grad_( - True - ) - - # [TODO] Move old embed_tokens to CPU - should be disk! - model.get_input_embeddings().original_module.to( - device = "cpu", non_blocking = True - ) - model.get_input_embeddings().original_module.requires_grad_(False) if "lm_head" in new_target_modules: print("Unsloth: Training lm_head in mixed precision to save VRAM") - new_dtype = model.get_output_embeddings().modules_to_save.default.weight.dtype - if new_dtype == torch.float16: - # See https://github.com/unslothai/unsloth/pull/1200 - # Tesla T4 must use float32 and not float16 - new_dtype = torch.float32 - - model.get_output_embeddings().modules_to_save.default.to( - device = DEVICE_TYPE_TORCH, dtype = new_dtype, non_blocking = True + _offload_frozen_module_for_training( + model.get_output_embeddings(), DEVICE_TYPE_TORCH ) - model.get_output_embeddings().modules_to_save.default.requires_grad_( - True - ) - - # [TODO] Move old lm_head to CPU - should be disk! - model.get_output_embeddings().original_module.to( - device = "cpu", non_blocking = True - ) - model.get_output_embeddings().original_module.requires_grad_(False) return model else: @@ -2779,9 +2804,10 @@ class FastLlamaModel: type(init_lora_weights) is bool or init_lora_weights == "gaussian" or init_lora_weights == "loftq" + or init_lora_weights == "corda" ): raise ValueError( - 'Unsloth: `init_lora_weights` must be either [True, False, "gaussian", "loftq"].' + 'Unsloth: `init_lora_weights` must be either [True, False, "gaussian", "loftq", "corda"].' ) if init_lora_weights == "loftq": @@ -2952,6 +2978,7 @@ class FastLlamaModel: loftq_config = loftq_config, use_rslora = use_rslora, modules_to_save = modules_to_save, + ensure_weight_tying = ensure_weight_tying, **kwargs, ) if not SUPPORTS_LOFTQ: @@ -3001,6 +3028,55 @@ class FastLlamaModel: model = FastLlamaModel.patch_peft_model(model, use_gradient_checkpointing) + if ensure_weight_tying: + try: + input_embeddings = model.get_input_embeddings() + output_embeddings = model.get_output_embeddings() + + if input_embeddings is not None and output_embeddings is not None: + + def _retie_parameter(target_module, source_module): + if not hasattr(source_module, "weight"): + return + weight = source_module.weight + # Remove existing registration to avoid "attribute already exists" + if "weight" in getattr(target_module, "_parameters", {}): + target_module._parameters.pop("weight") + if hasattr(target_module, "weight"): + try: + delattr(target_module, "weight") + except Exception as exc: + logger.warning_once( + f"Unsloth: Could not delete existing weight attr during retie on " + f"{type(target_module).__name__}: {exc}" + ) + target_module.register_parameter("weight", weight) + + # Tie trainable copies created by ModulesToSaveWrapper first (these are used in forward) + if hasattr(input_embeddings, "modules_to_save") and hasattr( + output_embeddings, "modules_to_save" + ): + if hasattr( + input_embeddings.modules_to_save, "default" + ) and hasattr(output_embeddings.modules_to_save, "default"): + _retie_parameter( + output_embeddings.modules_to_save.default, + input_embeddings.modules_to_save.default, + ) + + # Tie original_module references as well if present + if hasattr(input_embeddings, "original_module") and hasattr( + output_embeddings, "original_module" + ): + _retie_parameter( + output_embeddings.original_module, + input_embeddings.original_module, + ) + except Exception as e: + logger.warning_once( + f"Unsloth: Failed to ensure weight tying between embeddings and lm_head: {e}" + ) + if train_embed_tokens: print("Unsloth: Training embed_tokens in mixed precision to save VRAM") assert hasattr(model.get_input_embeddings(), "modules_to_save") diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 91016a13ba..eb3b21e206 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -151,8 +151,41 @@ class FastLanguageModel(FastLlamaModel): *args, **kwargs, ): + # Respect user-provided quantization_config (e.g. BitsAndBytesConfig) + quantization_config = kwargs.get("quantization_config", None) + if quantization_config is not None: + if isinstance(quantization_config, dict): + q_load_in_4bit = quantization_config.get("load_in_4bit", False) + q_load_in_8bit = quantization_config.get("load_in_8bit", False) + else: + q_load_in_4bit = getattr(quantization_config, "load_in_4bit", False) + q_load_in_8bit = getattr(quantization_config, "load_in_8bit", False) + if q_load_in_4bit: + load_in_4bit = True + load_in_8bit = False + if q_load_in_8bit: + load_in_8bit = True + load_in_4bit = False + # Login to allow private models token = hf_login(token) + # Align dtype with bnb_4bit_compute_dtype if provided and dtype is unset. + if dtype is None and quantization_config is not None: + bnb_compute_dtype = None + if isinstance(quantization_config, dict): + if quantization_config.get("load_in_4bit", False): + bnb_compute_dtype = quantization_config.get( + "bnb_4bit_compute_dtype", None + ) + else: + if getattr(quantization_config, "load_in_4bit", False): + bnb_compute_dtype = getattr( + quantization_config, "bnb_4bit_compute_dtype", None + ) + if isinstance(bnb_compute_dtype, str): + bnb_compute_dtype = getattr(torch, bnb_compute_dtype, None) + if isinstance(bnb_compute_dtype, torch.dtype): + dtype = bnb_compute_dtype if load_in_8bit or full_finetuning or qat_scheme is not None: return FastModel.from_pretrained( model_name = model_name, @@ -204,6 +237,17 @@ class FastLanguageModel(FastLlamaModel): "Unsloth: Please install vLLM before enabling `fast_inference`!\n" "You can do this in a terminal via `pip install vllm`" ) + if DEVICE_TYPE_TORCH == "cuda": + for i in range(DEVICE_COUNT): + # [TODO] DGX Spark vLLM breaks + if "NVIDIA GB10" in str(torch.cuda.get_device_name(i)).upper(): + print( + "Unsloth: DGX Spark detected - `fast_inference=True` is currently broken as of January 2026.\n" + "Defaulting to native Unsloth inference." + ) + fast_inference = False + break + # [TODO] For now fast_inference only works with fast_inference ie vLLM if load_in_fp8 != False: if not fast_inference: @@ -531,11 +575,17 @@ class FastLanguageModel(FastLlamaModel): if fast_inference: fast_inference, model_name = fast_inference_setup(model_name, model_config) + load_in_4bit_kwargs = load_in_4bit + load_in_8bit_kwargs = load_in_8bit + if quantization_config is not None and not fast_inference: + load_in_4bit_kwargs = False + load_in_8bit_kwargs = False + model, tokenizer = dispatch_model.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, dtype = _get_dtype(dtype), - load_in_4bit = load_in_4bit, + load_in_4bit = load_in_4bit_kwargs, token = token, device_map = device_map, rope_scaling = rope_scaling, @@ -572,22 +622,30 @@ class FastLanguageModel(FastLlamaModel): ) if load_in_4bit: - # Fix up bitsandbytes config - compute_dtype = dtype_from_config(model.config) - quantization_config = { - # Sometimes compute_dtype is not a string!! - "bnb_4bit_compute_dtype": compute_dtype, - "bnb_4bit_quant_type": "nf4", - "bnb_4bit_use_double_quant": True, - "llm_int8_enable_fp32_cpu_offload": False, - "llm_int8_has_fp16_weight": False, - "llm_int8_skip_modules": None, - "llm_int8_threshold": 6.0, - "load_in_4bit": True, - "load_in_8bit": False, - "quant_method": "bitsandbytes", - } - model.config.update({"quantization_config": quantization_config}) + # Fix up bitsandbytes config, but respect user-provided quantization_config + if quantization_config is None: + compute_dtype = dtype_from_config(model.config) + quantization_config = { + # Sometimes compute_dtype is not a string!! + "bnb_4bit_compute_dtype": compute_dtype, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_use_double_quant": True, + "llm_int8_enable_fp32_cpu_offload": False, + "llm_int8_has_fp16_weight": False, + "llm_int8_skip_modules": None, + "llm_int8_threshold": 6.0, + "load_in_4bit": True, + "load_in_8bit": False, + "quant_method": "bitsandbytes", + } + model.config.update({"quantization_config": quantization_config}) + else: + if hasattr(quantization_config, "to_dict"): + model.config.update( + {"quantization_config": quantization_config.to_dict()} + ) + elif isinstance(quantization_config, dict): + model.config.update({"quantization_config": quantization_config}) if load_in_fp8 != False: _tag_model_with_fp8_torchao_config(model, fp8_mode) @@ -679,12 +737,45 @@ class FastModel(FastBaseModel): *args, **kwargs, ): + # Respect user-provided quantization_config (e.g. BitsAndBytesConfig) + quantization_config = kwargs.get("quantization_config", None) + if quantization_config is not None: + if isinstance(quantization_config, dict): + q_load_in_4bit = quantization_config.get("load_in_4bit", False) + q_load_in_8bit = quantization_config.get("load_in_8bit", False) + else: + q_load_in_4bit = getattr(quantization_config, "load_in_4bit", False) + q_load_in_8bit = getattr(quantization_config, "load_in_8bit", False) + if q_load_in_4bit: + load_in_4bit = True + load_in_8bit = False + if q_load_in_8bit: + load_in_8bit = True + load_in_4bit = False + # Login to allow private models token = hf_login(token) if whisper_language is not None: assert type(whisper_language) is str if whisper_task is not None: assert type(whisper_task) is str + # Align dtype with bnb_4bit_compute_dtype if provided and dtype is unset. + if dtype is None and quantization_config is not None: + bnb_compute_dtype = None + if isinstance(quantization_config, dict): + if quantization_config.get("load_in_4bit", False): + bnb_compute_dtype = quantization_config.get( + "bnb_4bit_compute_dtype", None + ) + else: + if getattr(quantization_config, "load_in_4bit", False): + bnb_compute_dtype = getattr( + quantization_config, "bnb_4bit_compute_dtype", None + ) + if isinstance(bnb_compute_dtype, str): + bnb_compute_dtype = getattr(torch, bnb_compute_dtype, None) + if isinstance(bnb_compute_dtype, torch.dtype): + dtype = bnb_compute_dtype SUPPORTS_BFLOAT16 = is_bfloat16_supported() if dtype is None: dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 @@ -744,6 +835,17 @@ class FastModel(FastBaseModel): "Unsloth: Please install vLLM before enabling `fast_inference`!\n" "You can do this in a terminal via `pip install vllm`" ) + if DEVICE_TYPE_TORCH == "cuda": + for i in range(DEVICE_COUNT): + # [TODO] DGX Spark vLLM breaks + if "NVIDIA GB10" in str(torch.cuda.get_device_name(i)).upper(): + print( + "Unsloth: DGX Spark detected - `fast_inference=True` is currently broken as of January 2026.\n" + "Defaulting to native Unsloth inference." + ) + fast_inference = False + break + # [TODO] For now fast_inference only works with fast_inference ie vLLM if load_in_fp8 != False: if not fast_inference: @@ -1147,12 +1249,18 @@ class FastModel(FastBaseModel): if auto_model is None: auto_model = AutoModelForVision2Seq if is_vlm else AutoModelForCausalLM + load_in_4bit_kwargs = load_in_4bit + load_in_8bit_kwargs = load_in_8bit + if quantization_config is not None and not fast_inference: + load_in_4bit_kwargs = False + load_in_8bit_kwargs = False + model, tokenizer = FastBaseModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, dtype = _get_dtype(dtype), - load_in_4bit = load_in_4bit, - load_in_8bit = load_in_8bit, + load_in_4bit = load_in_4bit_kwargs, + load_in_8bit = load_in_8bit_kwargs, load_in_16bit = load_in_16bit, full_finetuning = full_finetuning, token = token, @@ -1198,22 +1306,30 @@ class FastModel(FastBaseModel): ) if load_in_4bit: - # Fix up bitsandbytes config - compute_dtype = dtype_from_config(model.config) - quantization_config = { - # Sometimes compute_dtype is not a string!! - "bnb_4bit_compute_dtype": compute_dtype, - "bnb_4bit_quant_type": "nf4", - "bnb_4bit_use_double_quant": True, - "llm_int8_enable_fp32_cpu_offload": False, - "llm_int8_has_fp16_weight": False, - "llm_int8_skip_modules": None, - "llm_int8_threshold": 6.0, - "load_in_4bit": True, - "load_in_8bit": False, - "quant_method": "bitsandbytes", - } - model.config.update({"quantization_config": quantization_config}) + # Fix up bitsandbytes config, but respect user-provided quantization_config + if quantization_config is None: + compute_dtype = dtype_from_config(model.config) + quantization_config = { + # Sometimes compute_dtype is not a string!! + "bnb_4bit_compute_dtype": compute_dtype, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_use_double_quant": True, + "llm_int8_enable_fp32_cpu_offload": False, + "llm_int8_has_fp16_weight": False, + "llm_int8_skip_modules": None, + "llm_int8_threshold": 6.0, + "load_in_4bit": True, + "load_in_8bit": False, + "quant_method": "bitsandbytes", + } + model.config.update({"quantization_config": quantization_config}) + else: + if hasattr(quantization_config, "to_dict"): + model.config.update( + {"quantization_config": quantization_config.to_dict()} + ) + elif isinstance(quantization_config, dict): + model.config.update({"quantization_config": quantization_config}) if load_in_fp8 != False: _tag_model_with_fp8_torchao_config(model, fp8_mode) diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 85332e1116..1e5533c25c 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -28,7 +28,6 @@ from .mapper import ( ) # https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading! -from packaging.version import Version from transformers import __version__ as transformers_version from unsloth.models._utils import TorchAOConfig from unsloth_zoo.utils import Version @@ -409,7 +408,7 @@ def _get_fp8_mode_and_check_settings( if Version(torchao.__version__) < Version("0.15.0"): raise ValueError(error_message) - # If fbgemm_gpu_genai is installed, check if it's >= 1.4.1 + # If fbgemm_gpu_genai is installed and old, disable FBGEMM and use Triton instead if ( importlib.util.find_spec("fbgemm_gpu") is not None and importlib.util.find_spec("fbgemm_gpu.experimental") is not None @@ -417,7 +416,12 @@ def _get_fp8_mode_and_check_settings( import fbgemm_gpu.experimental.gen_ai if Version(fbgemm_gpu.__version__) < Version("1.4.1"): - raise ValueError( - "Unsloth: On the fly `load_in_fp8` is only compatible with fbgemm_gpu_genai 1.4.1+. Try `unsloth/Qwen3-8B` instead." + # Old FBGEMM version - disable and use Triton kernels instead + os.environ["UNSLOTH_HAS_FBGEMM"] = "0" + from unsloth_zoo.log import logger + + logger.info( + f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu.__version__} is old for FP8 loading. " + f"Using Triton kernels instead." ) return fp8_mode diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 0eed45c5cd..5e893d2b6f 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -307,9 +307,9 @@ def MistralForCausalLM_fast_forward( RETURN_LOGITS = False if not RETURN_LOGITS and labels is not None: - n_items = kwargs.get("num_items_in_batch", None) or kwargs.get( - "n_items", None - ) + n_items = kwargs.get("num_items_in_batch", None) + if n_items is None: + n_items = kwargs.get("n_items", None) logit_softcapping = getattr(self.config, "final_logit_softcapping", 0) # loss = fused_linear_cross_entropy( @@ -363,11 +363,13 @@ def MistralForCausalLM_fast_forward( shift_labels, kwargs.get("packed_seq_lengths"), ) + n_items = kwargs.get("num_items_in_batch", None) + if n_items is None: + n_items = kwargs.get("n_items", None) loss = fast_cross_entropy_loss( logits = shift_logits, labels = shift_labels, - n_items = kwargs.get("num_items_in_batch", None) - or kwargs.get("n_items", None), + n_items = n_items, ) if not return_dict: diff --git a/unsloth/models/qwen3_moe.py b/unsloth/models/qwen3_moe.py index bec3fa7b0d..e1f8c71b6b 100644 --- a/unsloth/models/qwen3_moe.py +++ b/unsloth/models/qwen3_moe.py @@ -207,7 +207,7 @@ class FastQwen3MoeModel(FastQwen3Model): # https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py\ import transformers.models.qwen3_moe.modeling_qwen3_moe - transformers.models.Qwen3Moe.modeling_qwen3_moe.Qwen3MoeRotaryEmbedding = ( + transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeRotaryEmbedding = ( LlamaRotaryEmbedding ) return @@ -236,7 +236,7 @@ class FastQwen3MoeModel(FastQwen3Model): device_map = device_map, rope_scaling = rope_scaling, fix_tokenizer = fix_tokenizer, - model_patcher = FastQwen3Model, + model_patcher = FastQwen3MoeModel, tokenizer_name = tokenizer_name, trust_remote_code = trust_remote_code, **kwargs, diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 4ea36519d9..9788207c99 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -22,7 +22,6 @@ from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union import inspect import os import re -import torch from unsloth_zoo.compiler import create_new_function from unsloth_zoo.log import logger from unsloth_zoo.logging_utils import PatchRLStatistics @@ -44,10 +43,31 @@ torch_compile_options = { "triton.cudagraphs": False, } -from trl import __version__ as trl_version +# vLLM compatibility shim (TRL expects GuidedDecodingParams even if vLLM doesn't provide it) +try: + import vllm.sampling_params as _unsloth_vllm_sp + + if not hasattr(_unsloth_vllm_sp, "GuidedDecodingParams"): + + class GuidedDecodingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + + _unsloth_vllm_sp.GuidedDecodingParams = GuidedDecodingParams +except Exception: + pass + +from trl import __version__ as trl_version_raw +from importlib.metadata import version as importlib_version from unsloth_zoo.utils import Version -trl_version = Version(trl_version) +try: + trl_version = Version(trl_version_raw) +except Exception: + try: + trl_version = Version(importlib_version("trl")) + except Exception: + trl_version = Version("0.0.0") def vLLMSamplingParams(**kwargs): @@ -200,22 +220,24 @@ def PatchRL(FastLanguageModel): unwrap = "unwrap_model_for_generation" for trainer in trainers: try: - current_trainer = eval(f"trl.trainer.{trainer}") + current_trainer = getattr(trl.trainer, trainer) except: continue if hasattr(current_trainer, unwrap): try: - exec(f"trl.trainer.{trainer}.{unwrap} = unsloth_{unwrap}") + setattr(current_trainer, unwrap, unsloth_unwrap_model_for_generation) except: continue - exec(f"Trainer.prediction_step=unsloth_prediction_step") + Trainer.prediction_step = unsloth_prediction_step +grpo_selective_log_softmax = RL_REPLACEMENTS["grpo_selective_log_softmax"] selective_log_softmax = RL_REPLACEMENTS["selective_log_softmax"] calculate_pad_tokens_in_prompt = RL_REPLACEMENTS["calculate_pad_tokens_in_prompt"] create_completion_attention_mask = RL_REPLACEMENTS["create_completion_attention_mask"] left_pack_padding = RL_REPLACEMENTS["left_pack_padding"] align_logprobs_with_mask = RL_REPLACEMENTS["align_logprobs_with_mask"] +autotune_batch_and_chunks = RL_REPLACEMENTS["grpo_autotune_batch_and_chunks"] RLTrainer_replacement = ''' import os @@ -234,16 +256,33 @@ from transformers.training_args import ParallelMode # Also patches W&B since multiple runs must use wandb.finish() import functools from types import MethodType +try: + from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers +except: + def reset_unsloth_gradient_checkpointing_buffers(): pass def prepare_for_training_mode(f): @functools.wraps(f) def wrapper(self, *args, **kwargs): # Enable training mode + _was_training = None + # Get gradient checkpointing setting from training arguments + use_gc = getattr(self.args, 'gradient_checkpointing', True) + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training if hasattr(self, 'model') and hasattr(self.model, "for_training"): - self.model.for_training() + self.model.for_training(use_gradient_checkpointing=use_gc) output = f(self, *args, **kwargs) - # Return inference mode + # Restore previous mode when possible if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - self.model.for_inference() + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training(use_gradient_checkpointing=use_gc) + # Reset gradient checkpointing buffers to free memory while staying ready for next run + try: + reset_unsloth_gradient_checkpointing_buffers() + except: + pass # Patch W&B to enable logging on future runs, otherwise it'll overwrite the first run try: import wandb @@ -262,11 +301,13 @@ torch_compile_options = {{ "triton.cudagraphs" : False, }} +{grpo_selective_log_softmax_code} {selective_log_softmax_code} {calculate_pad_tokens_in_prompt_code} {create_completion_attention_mask_code} {left_pack_padding_code} {align_logprobs_with_mask_code} +{autotune_batch_and_chunks_code} {RL_pre} @@ -283,10 +324,20 @@ class Unsloth{RLConfig_name}({RLConfig_name}): default = -1, metadata = {{'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}}, ) + unsloth_logit_chunk_multiplier : Optional[int] = field( + default = None, + metadata = {{'help': 'Multiplier for chunked logit computations.'}}, + ) + unsloth_grpo_mini_batch : Optional[int] = field( + default = None, + metadata = {{'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}}, + ) {max_seq_length_pre} def __init__({RLConfig_arguments}, vllm_sampling_params = None, unsloth_num_chunks = -1, + unsloth_logit_chunk_multiplier = None, + unsloth_grpo_mini_batch = None, {max_seq_length_call} **kwargs, ): @@ -294,6 +345,15 @@ class Unsloth{RLConfig_name}({RLConfig_name}): super().__init__({RLConfig_call_args}{RLConfig_kwargs}) self.vllm_sampling_params = vllm_sampling_params self.unsloth_num_chunks = unsloth_num_chunks + if unsloth_grpo_mini_batch is not None: + if self.generation_batch_size >= unsloth_grpo_mini_batch: + self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch + else: + raise ValueError( + f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, " + f"which is self.per_device_train_batch_size * gradient_accumulation_steps." + ) + self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier {max_seq_length_post} pass @@ -323,6 +383,32 @@ pass ''' +def _wrap_grpo_generate_and_score(trainer_cls): + if not hasattr(trainer_cls, "_generate_and_score_completions"): + return + original = trainer_cls._generate_and_score_completions + if getattr(original, "_unsloth_restore_training_wrapped", False): + return + + def wrapped(self, *args, **kwargs): + was_training = getattr(getattr(self, "model", None), "training", None) + try: + return original(self, *args, **kwargs) + finally: + if ( + was_training is False + and hasattr(self, "model") + and hasattr(self.model, "for_inference") + ): + try: + self.model.for_inference() + except Exception: + pass + + wrapped._unsloth_restore_training_wrapped = True + trainer_cls._generate_and_score_completions = wrapped + + def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Patch for vLLM and Unsloth PEFT import trl @@ -559,8 +645,12 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): " if args_max_seq_length is None and model_max_seq_length is not None:\n" " max_seq_length = model.max_seq_length\n" " if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length\n" + " elif args_max_seq_length is not None and model_max_seq_length is not None:\n" + " if args_max_seq_length > model_max_seq_length:\n" + " print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but '\n" + " 'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.')\n" + " args.max_seq_length = model_max_seq_length\n" ) - " elif args_max_seq_length is not None and model_max_seq_length is not None:\n" " if args_max_seq_length > model_max_seq_length:\n" " print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but \n" " the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.')\n" " args.max_seq_length = model_max_seq_length\n" extra_args += length_check # At this point max_seq_length might be set, but trl is moving to max_length @@ -681,6 +771,19 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): ) RLTrainer_post += training_check + # Sync chat_template from processing_class to vLLM's tokenizer + # This fixes base models that have custom chat templates applied after loading + if "model" in call_args: + vllm_chat_template_sync = ( + "if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'):\n" + " _vllm_tok = self.llm.get_tokenizer()\n" + " _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None)\n" + " if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None:\n" + " _vllm_tok.chat_template = _pc.chat_template\n" + "pass\n" + ) + RLTrainer_post += vllm_chat_template_sync + # Edit optional metrics other_metrics_processor = "" if trainer_file in RL_METRICS_CHANGES: @@ -813,7 +916,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): num_proc_check = ( "if dataset_num_proc is None:\n" " import psutil\n" - " dataset_num_proc = min(max(psutil.cpu_count()+4, 2), 64)\n" + " dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)\n" " memory_gb_left = psutil.virtual_memory().available / (1024**3)\n" " if memory_gb_left <= 4: dataset_num_proc = 1 # Too risky, so set to 1\n" " elif memory_gb_left <= 6: dataset_num_proc = min(2, dataset_num_proc)\n" @@ -900,9 +1003,9 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): if "temperature" in call_args: check_temperature = ( "if temperature <= 0:\n" - " raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')\n" + " raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')\n" "elif temperature >= 10:\n" - " raise MathError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.')\n" + " raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.')\n" "\n" ) extra_args += check_temperature @@ -948,6 +1051,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Selective log softmax and other functions selective_log_softmax_code = inspect.getsource(selective_log_softmax) + grpo_selective_log_softmax_code = inspect.getsource(grpo_selective_log_softmax) calculate_pad_tokens_in_prompt_code = inspect.getsource( calculate_pad_tokens_in_prompt ) @@ -956,6 +1060,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): ) left_pack_padding_code = inspect.getsource(left_pack_padding) align_logprobs_with_mask_code = inspect.getsource(align_logprobs_with_mask) + autotune_batch_and_chunks_code = inspect.getsource(autotune_batch_and_chunks) # Get final source code RLTrainer_source = RLTrainer_replacement.format( RLTrainer_name = RLTrainer_name, @@ -977,8 +1082,10 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): max_seq_length_call = max_seq_length_call, max_seq_length_post = max_seq_length_post, selective_log_softmax_code = selective_log_softmax_code, + grpo_selective_log_softmax_code = grpo_selective_log_softmax_code, calculate_pad_tokens_in_prompt_code = calculate_pad_tokens_in_prompt_code, create_completion_attention_mask_code = create_completion_attention_mask_code, + autotune_batch_and_chunks_code = autotune_batch_and_chunks_code, left_pack_padding_code = left_pack_padding_code, align_logprobs_with_mask_code = align_logprobs_with_mask_code, ) @@ -990,10 +1097,10 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Temporary patch _is_vlm to False # as of 0.22 it only exists in sfttrainer - oriignal_is_vlm_text = "self._is_vlm = True" + original_is_vlm_text = "self._is_vlm = True" new_is_vlm_text = "self._is_vlm = False" RLTrainer_source = RLTrainer_source.replace( - oriignal_is_vlm_text, new_is_vlm_text + original_is_vlm_text, new_is_vlm_text ) # Remove multiple doc strings @@ -1046,6 +1153,16 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): globals(), ) + if trainer_file == "grpo_trainer": + try: + _wrap_grpo_generate_and_score( + getattr(created_module, f"Unsloth{RLTrainer_name}") + ) + except Exception as e: + logger.info( + f"Unsloth: Could not wrap _generate_and_score_completions for {RLTrainer_name}: {e}" + ) + def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports): init = inspect.getsource(RLTrainer.__init__) @@ -1077,6 +1194,41 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import "model = self._prepare_peft_model(model, peft_config, args)\n", "pass\n" ) + # Skip add_adapter("ref") for reference model computation + # Unsloth: We comment out the "ref" adapter creation because: + # 1. We want to use the original BASE MODEL as the reference model, not the SFT/LoRA model + # 2. PEFT doesn't allow multiple adapters when target_parameters is used (MoE models) + # When "ref" is not in peft_config, GRPO/RLOO fallback uses disable_adapter() + # which gives the base model logits - exactly what we want + add_adapter_block_pattern = ( + r"([ \t]*)" # Capture leading indentation + r"if\s+is_peft_available\(\)\s+and\s+is_peft_model\(model\)\s+and\s+args\.beta\s*!=\s*0\.0\s*:" + r"(.*?)" # Match the entire block until ref_param.data.copy_ + r"ref_param\.data\.copy_\(param\.data\)" + ) + + def comment_out_block(match): + """Comment out each line in the matched block, preserving indentation.""" + full_match = match.group(0) + indent = match.group(1) + lines = full_match.split("\n") + commented_lines = [] + # Add explanation comment first + commented_lines.append( + f"{indent}# Unsloth: Commented out - use base model as reference, not SFT/LoRA model" + ) + # Comment out each line - insert # after leading whitespace to preserve indentation + for line in lines: + if line.strip(): + stripped = line.lstrip() + leading_ws = line[: len(line) - len(stripped)] + commented_lines.append(f"{leading_ws}# {stripped}") + else: + commented_lines.append(line) + return "\n".join(commented_lines) + + init = re.sub(add_adapter_block_pattern, comment_out_block, init, flags = re.DOTALL) + # Set use_vllm if not set if "args.use_vllm" in init and "model" in init and "args" in init: # .*? matches first match. .+? matches final match. diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 7d4d520c1f..ff36da125d 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -50,7 +50,7 @@ RL_ADDITIONAL_FUNCTIONS = defaultdict(list) torch_compile_options = { "epilogue_fusion": True, - "max_autotune": True, + "max_autotune": False, # I saw speedups, but not sure if this has issues in collab "shape_padding": True, "trace.enabled": False, "triton.cudagraphs": False, @@ -258,18 +258,20 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # The new multi-line string that will replace the line above replacement_lines = """ + max_left_pad = None batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size try: # TRL 0.23.1 and below path if not has_images: # Left pad prompt before calculation old and ref hidden states - prompt_completion_ids = left_pack_padding(prompt_completion_ids, self.processing_class.pad_token_id) - self.model.for_training() + left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) + max_left_pad = torch.max(left_pad_tokens_per_prompt).item() except: # TRL 0.24.0 and below path if images is None: # Left pad prompt before calculation old and ref hidden states - prompt_completion_ids = left_pack_padding(prompt_completion_ids, self.processing_class.pad_token_id) + left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) + max_left_pad = torch.max(left_pad_tokens_per_prompt).item() self.model.for_training()""" function = function.replace(line_to_replace, replacement_lines) @@ -346,17 +348,45 @@ def grpo_trainer__generate_and_score_completions(function_name, function): if self.use_vllm:""" function = function.replace(replace_part, new_replacement) + # Important note: we disable TRL's importance sampling logic + # It is disabled because the LLM path moves left padding to the right. + # We must adjust the vLLM sampling_logprob tensor in Unsloth to account for this. + string_to_find = "if self.use_vllm and self.vllm_importance_sampling_correction:" + + replacement_string = ( + "if False and self.use_vllm and self.vllm_importance_sampling_correction:" + ) + + function = function.replace(string_to_find, replacement_string) + string_to_find = """ if "image_sizes" in prompt_inputs: output["image_sizes"] = prompt_inputs["image_sizes"]""" replacement_string = """ if "image_sizes" in prompt_inputs: output["image_sizes"] = prompt_inputs["image_sizes"] - - if self.use_vllm: - try: + if max_left_pad is not None: + output["max_left_pad"] = torch.tensor(prompt_ids.shape[0] * [max_left_pad]).unsqueeze(-1) + try: + if self.use_vllm and getattr(self, "vllm_importance_sampling_correction", False): output["sampling_per_token_logps"] = sampling_per_token_logps - except NameError: - output["sampling_per_token_logps"] = None""" + except NameError: + output["sampling_per_token_logps"] = None""" + + function = function.replace(string_to_find, replacement_string) + + # This path is for TRL 0.24.0 images is a variable exclusive to this version + string_to_find = """ if images is not None: + output["num_images"] = num_images""" + + replacement_string = """ if images is not None: + output["num_images"] = num_images + if max_left_pad is not None: + output["max_left_pad"] = torch.tensor(prompt_ids.shape[0] * [max_left_pad]).unsqueeze(-1) + try: + if self.use_vllm and getattr(self, "vllm_importance_sampling_correction", False): + output["sampling_per_token_logps"] = sampling_per_token_logps + except NameError: + output["sampling_per_token_logps"] = None""" function = function.replace(string_to_find, replacement_string) @@ -532,12 +562,12 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): *args, **kwargs, ): + # All Unsloth code here in this function is licensed under AGPL3 # if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0': # return None, None # logps, entropies Unsloth efficient GRPO if compute_efficient: return None, None else: - # Otherwise, calculate normally: if not hasattr(self, "_autocast_dtype"): self._autocast_dtype = ( torch.float16 @@ -556,47 +586,199 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): kwargs.get("image_sizes", None), ) - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - unwrapped_model = self.accelerator.unwrap_model( model, keep_fp32_wrapper = False ) - with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): - with _get_inference_mode_context_manager(model): - if pixel_values is None: - attention_mask = input_ids != self.processing_class.pad_token_id - attention_mask = attention_mask.to(attention_mask.dtype) - # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded - logits = unwrapped_model( - input_ids = input_ids, - attention_mask = attention_mask, - pixel_values = pixel_values, - image_grid_thw = image_grid_thw, - pixel_attention_mask = pixel_attention_mask, - image_sizes = image_sizes, - # logits_to_keep = logits_to_keep + 1, - ).logits + lm_head = self.model.get_output_embeddings().weight + + dtype_bytes = ( + 16 if self._autocast_dtype in [torch.float16, torch.bfloat16] else 32 + ) + total_rows = input_ids.shape[0] + seq_len = input_ids.shape[1] + hidden_dim = lm_head.shape[1] + vocab_dim = lm_head.shape[0] + + if self.args.unsloth_grpo_mini_batch is None: + B, multiplier = autotune_batch_and_chunks( + total_rows, + seq_len, + hidden_dim, + vocab_dim, + dtype_bytes, + self.args.unsloth_logit_chunk_multiplier, + ) + B = total_rows // B + else: + B = self.args.unsloth_grpo_mini_batch + + if self.args.unsloth_logit_chunk_multiplier is None: + multiplier = max(4, seq_len // 4096) + else: + multiplier = self.args.unsloth_logit_chunk_multiplier + + all_logprobs_list = [] + if pixel_values is None: + left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt( + input_ids, logits_to_keep, self.processing_class.pad_token_id + ) + max_left_pad = torch.max(left_pad_tokens_per_prompt).item() + input_ids = left_pack_padding( + input_ids, self.processing_class.pad_token_id + ) + attention_mask = input_ids != self.processing_class.pad_token_id + attention_mask = attention_mask.to(attention_mask.dtype) + else: + max_left_pad = 0 + + # input_ids_chunks = torch.chunk(input_ids, chunks = B, dim = 0) + attention_mask_chunks = torch.chunk(attention_mask, chunks = B, dim = 0) + + def chunk_optional(tensor, chunks): + if tensor is None: + return [None] * chunks + return torch.chunk(tensor, chunks = chunks, dim = 0) + + import math + + total_samples = input_ids.shape[0] + batch_size = math.ceil(total_samples / B) + + input_ids_chunks = [] + attention_mask_chunks = [] + pixel_values_chunks = [] + image_grid_thw_chunks = [] + pixel_attention_mask_chunks = [] + + current_pixel_idx = 0 + # TRL 0.23.0 batching logic + for start in range(0, total_samples, batch_size): + end = start + batch_size + + input_ids_chunks.append(input_ids[start:end]) + attention_mask_chunks.append(attention_mask[start:end]) + + if image_grid_thw is not None and pixel_values is not None: + grid_slice = image_grid_thw[start:end] + image_grid_thw_chunks.append(grid_slice) + + batch_pixel_count = grid_slice.prod(dim = -1).sum().item() + + start_pixel_idx = current_pixel_idx + end_pixel_idx = current_pixel_idx + batch_pixel_count + + pixel_values_chunks.append( + pixel_values[start_pixel_idx:end_pixel_idx] + ) + + if pixel_attention_mask is not None: + pixel_attention_mask_chunks.append( + pixel_attention_mask[start_pixel_idx:end_pixel_idx] + ) else: - logits = unwrapped_model( - input_ids = input_ids, - attention_mask = attention_mask, - pixel_values = pixel_values, - image_grid_thw = image_grid_thw, - pixel_attention_mask = pixel_attention_mask, - image_sizes = image_sizes, - logits_to_keep = logits_to_keep + 1, - ).logits + pixel_attention_mask_chunks.append(None) + current_pixel_idx = end_pixel_idx + + else: + pixel_values_chunks.append(None) + image_grid_thw_chunks.append(None) + pixel_attention_mask_chunks.append(None) + + if image_sizes is not None and not isinstance(image_sizes, torch.Tensor): + image_sizes_chunks = [[size] for size in image_sizes] + else: + image_sizes_chunks = chunk_optional(image_sizes, B) + + temperature = self.temperature + logit_softcapping = getattr(model.config, "final_logit_softcapping", 0) + if logit_softcapping is None: + logit_softcapping = 0 + logit_scale_multiply = getattr(model.config, "logit_scale", 0) + if logit_scale_multiply is None: + logit_scale_multiply = 0 + logit_scale_divide = getattr(model.config, "logits_scaling", 0) + if logit_scale_divide is None: + logit_scale_divide = 0 + + zipped_inputs = zip( + input_ids_chunks, + attention_mask_chunks, + pixel_values_chunks, + image_grid_thw_chunks, + pixel_attention_mask_chunks, + image_sizes_chunks, + ) + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + + with _get_inference_mode_context_manager(model): + for ( + input_ids_chunk, + attention_mask_chunk, + pixel_values_chunk, + image_grid_thw_chunk, + pixel_attention_mask_chunk, + image_sizes_chunk, + ) in zipped_inputs: + with torch.amp.autocast( + device_type = "cuda", dtype = self._autocast_dtype + ): + if pixel_values is None: + logits_chunk = unwrapped_model( + input_ids = input_ids_chunk, + attention_mask = attention_mask_chunk, + pixel_values = pixel_values_chunk, + image_grid_thw = image_grid_thw_chunk, + pixel_attention_mask = pixel_attention_mask_chunk, + image_sizes = image_sizes_chunk, + ).logits + + completion_input_ids_chunk = input_ids_chunk[ + :, -(logits_to_keep + max_left_pad) : + ] + logits_chunk = logits_chunk[ + :, -(logits_to_keep + max_left_pad + 1) :, : + ] + logits_chunk = logits_chunk[:, :-1, :] + else: + # Essentially, for VLMs we do not go via the optimized path in models/, + # so we don't encounter the Flash Attn left-padding issue. + logits_chunk = unwrapped_model( + input_ids = input_ids_chunk, + attention_mask = attention_mask_chunk, + pixel_values = pixel_values_chunk, + image_grid_thw = image_grid_thw_chunk, + pixel_attention_mask = pixel_attention_mask_chunk, + image_sizes = image_sizes_chunk, + logits_to_keep = logits_to_keep + 1, + ).logits + + logits_chunk = logits_chunk[:, :-1, :] + completion_input_ids_chunk = input_ids_chunk[ + :, -logits_to_keep: + ] + + logprobs_chunk = chunked_hidden_states_selective_log_softmax( + logits_chunk, + lm_head, + completion_input_ids_chunk, + chunks = input_ids_chunk.shape[0] * multiplier, + logit_scale_multiply = logit_scale_multiply, + logit_scale_divide = logit_scale_divide, + logit_softcapping = logit_softcapping, + temperature = temperature, + ) + # This is needed to avoid race conditions with GPT OSS offload_embbed=True + # However, it seems that this line does not slow down or disrupt models. + torch.cuda.synchronize() + all_logprobs_list.append(logprobs_chunk) + logprobs = torch.cat(all_logprobs_list, dim = 0) entropies = None - if compute_entropy: - from trl.trainer.utils import entropy_from_logits - - entropies = entropy_from_logits(logits) os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" - # logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred - return logits.detach(), entropies # logps, entropies + + return logprobs.detach(), entropies # logps, entropies # input_ids = input_ids[:, -logits_to_keep:] # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. # See https://github.com/huggingface/trl/issues/2770 @@ -708,14 +890,14 @@ def grpo_trainer_compute_loss(function_name, function): # ref_per_token_logps = per_token_logps = get_logps_func(model, input_ids, attention_mask, logits_to_keep) # else: # ref_per_token_logps = None - ref_hidden_states = inputs.get("ref_per_token_logps", None) + ref_logps = inputs.get("ref_per_token_logps", None) # per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 # x - x.detach() allows for preserving gradients from x advantages = inputs["advantages"] # per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) # per_token_loss = -(per_token_loss - self.beta * per_token_kl) # loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean() - old_hidden_states = inputs.get("old_per_token_logps", None) + old_logps = inputs.get("old_per_token_logps", None) input_ids = input_ids[:, -logits_to_keep:] @@ -730,24 +912,13 @@ def grpo_trainer_compute_loss(function_name, function): if logit_scale_divide is None: logit_scale_divide = 0 + max_left_pad = inputs.get("max_left_pad", 0) if per_token_logps is not None: - if ref_hidden_states is not None: - ref_hidden_states = ref_hidden_states[ - :, :-1, : - ] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred - if old_hidden_states is not None: - old_hidden_states = old_hidden_states[ - :, :-1, : - ] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred - per_token_logps = per_token_logps[ - :, :-1, : - ] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred - loss, completion_length, mean_kl, delta, flat_is_ratio = ( grpo_compute_loss_slow( - ref_hidden_states, + ref_logps, per_token_logps, - old_hidden_states, + old_logps, input_ids, completion_mask, self.beta, @@ -761,6 +932,7 @@ def grpo_trainer_compute_loss(function_name, function): max_completion_length = self.args.max_completion_length, delta = self.args.delta, temperature = self.args.temperature, + max_left_pad = max_left_pad, logit_softcapping = logit_softcapping, logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, @@ -781,8 +953,8 @@ def grpo_trainer_compute_loss(function_name, function): logits_to_keep = logits_to_keep, completion_mask = completion_mask, advantages = advantages, - old_hidden_states = old_hidden_states, - ref_hidden_states = ref_hidden_states, + old_logps = old_logps, + ref_logps = ref_logps, n_chunks = self.args.unsloth_num_chunks, loss_type = self.args.loss_type, importance_sampling_level = self.importance_sampling_level, @@ -791,6 +963,7 @@ def grpo_trainer_compute_loss(function_name, function): max_completion_length = self.args.max_completion_length, delta = self.args.delta, temperature = self.args.temperature, + max_left_pad = max_left_pad, logit_softcapping = logit_softcapping, logit_scale_multiply = logit_scale_multiply, logit_scale_divide = logit_scale_divide, @@ -809,8 +982,8 @@ def grpo_trainer_compute_loss(function_name, function): logits_to_keep = logits_to_keep, completion_mask = completion_mask, advantages = advantages, - old_hidden_states = old_hidden_states, - ref_hidden_states = ref_hidden_states, + old_logps = old_logps, + ref_logps = ref_logps, n_chunks = self.args.unsloth_num_chunks, temperature = self.args.temperature, logit_softcapping = logit_softcapping, @@ -827,7 +1000,11 @@ def grpo_trainer_compute_loss(function_name, function): self._metrics["completion_length"].append(completion_length.item()) self._metrics["kl"].append(mean_kl.item()) - if self.use_vllm and delta is not None: + if ( + self.use_vllm + and delta is not None + and getattr(self, "vllm_importance_sampling_correction", False) + ): mean_delta = ( torch.mean(delta) if delta.numel() > 0 @@ -861,13 +1038,19 @@ def grpo_trainer_compute_loss(function_name, function): else torch.tensor(0.0, device = self.model.device) ) self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( - nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item() + self.accelerator.gather(min_importance_sampling_ratio) + .nan_to_num(nan = float("inf")) + .min() + .item() ) self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() ) self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( - nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item() + self.accelerator.gather(max_importance_sampling_ratio) + .nan_to_num(nan = float("-inf")) + .max() + .item() ) return loss @@ -949,11 +1132,15 @@ def openenv_vllm_reload_weights(): return if Version(importlib_version("trl")) < Version("0.26.0"): return + try: import trl.experimental.openenv.utils as openenv_utils import trl.experimental.openenv as openenv except ImportError as e: logger.info(f"Unsloth: Failed to import trl openenv: {e}") + logger.info( + "Unsloth: trl.experimental.openenv not available — skipping RL openenv patches." + ) return src = inspect.getsource(openenv_utils.generate_rollout_completions) diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py new file mode 100644 index 0000000000..b66ac7cf8a --- /dev/null +++ b/unsloth/models/sentence_transformer.py @@ -0,0 +1,1853 @@ +# Copyright 2025 electroglyph. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging + +from .loader import FastModel +from ._utils import SUPPORTS_BFLOAT16 +import inspect +import json +import os +import types +from huggingface_hub import hf_hub_download +from typing import Optional +import torch +from transformers.modeling_outputs import BaseModelOutput +from collections import OrderedDict +from transformers.models.distilbert import modeling_distilbert +from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask_for_sdpa +import transformers +from packaging.version import Version +import re +from transformers import AutoModel, AutoConfig +from transformers.models.auto.auto_factory import _get_model_class +import tempfile +from huggingface_hub import HfApi, get_token +from ..save import unsloth_save_pretrained_torchao, unsloth_save_pretrained_gguf +import contextlib +import shutil + + +def _save_pretrained_torchao( + self, + save_directory, + tokenizer = None, + torchao_config = None, + push_to_hub = False, + token = None, +): + self.save_pretrained(save_directory) + + # grab inner model + inner_model = self[0].auto_model + if hasattr(inner_model, "_orig_mod"): + inner_model = inner_model._orig_mod + + # merge LoRA first + if hasattr(inner_model, "merge_and_unload"): + inner_model = inner_model.merge_and_unload() + + # confirm Transformer path + transformer_path = "0_Transformer" + modules_path = os.path.join(save_directory, "modules.json") + if os.path.exists(modules_path): + try: + with open(modules_path, "r") as f: + modules = json.load(f) + for m in modules: + if m.get("type", "").endswith("Transformer"): + transformer_path = m.get("path", "") + break + except: + pass + + transformer_dir = os.path.join(save_directory, transformer_path) + transformer_dir = os.path.abspath(transformer_dir) + + if tokenizer is None: + tokenizer = self.tokenizer + + @contextlib.contextmanager + def patch_unsloth_save(): + original_causal = transformers.AutoModelForCausalLM + original_rmtree = shutil.rmtree + # unsloth_save_pretrained_torchao expects AutoModelForCausalLM + transformers.AutoModelForCausalLM = transformers.AutoModel + # prevent unsloth from deleting the unquantized model directory + shutil.rmtree = lambda *args, **kwargs: None + try: + yield + finally: + # unpatch + transformers.AutoModelForCausalLM = original_causal + shutil.rmtree = original_rmtree + + with patch_unsloth_save(): + unsloth_save_pretrained_torchao( + inner_model, + transformer_dir, + tokenizer = tokenizer, + torchao_config = torchao_config, + push_to_hub = push_to_hub, + token = token, + ) + + # avoid `0_Transformer-torchao`, it was either this or fix modules.json + torchao_dir = transformer_dir + "-torchao" + if os.path.exists(torchao_dir): + if not os.path.exists(transformer_dir): + os.makedirs(transformer_dir, exist_ok = True) + + # move contents + for item in os.listdir(torchao_dir): + s = os.path.join(torchao_dir, item) + d = os.path.join(transformer_dir, item) + if os.path.isdir(s): + shutil.copytree(s, d, dirs_exist_ok = True) + else: + shutil.copy2(s, d) + + # remove torchao dir + shutil.rmtree(torchao_dir) + + # remove conflicting safetensors if we brought in bin + if os.path.exists(os.path.join(transformer_dir, "pytorch_model.bin")): + safetensors_path = os.path.join(transformer_dir, "model.safetensors") + if os.path.exists(safetensors_path): + try: + os.remove(safetensors_path) + except: + pass + + try: + FastSentenceTransformer._add_unsloth_branding(save_directory) + except: + pass + + +# Thanks Etherl: +def _save_pretrained_gguf( + self, + save_directory, + tokenizer = None, + quantization_method = "fast_quantized", + first_conversion = None, + push_to_hub = False, + token = None, + max_shard_size = "5GB", + temporary_location = "_unsloth_temporary_saved_buffers", + maximum_memory_usage = 0.85, + **kwargs, +): + """ + Saves the SentenceTransformer model to GGUF format by saving the inner transformer model, + converting it, and placing the resulting GGUF files in the save directory. + """ + # 1. Save standard SentenceTransformer structure (configs, modules.json, etc.) + self.save_pretrained(save_directory) + + # 2. Extract inner transformer model + inner_model = self[0].auto_model + if hasattr(inner_model, "_orig_mod"): + inner_model = inner_model._orig_mod + + # If it's a PEFT model, unsloth_save_pretrained_gguf handles merging, + # but we pass the inner model wrapper. + + # 3. Identify where the transformer weights are stored + transformer_path = "0_Transformer" + modules_path = os.path.join(save_directory, "modules.json") + if os.path.exists(modules_path): + try: + with open(modules_path, "r") as f: + modules = json.load(f) + for m in modules: + if m.get("type", "").endswith("Transformer"): + transformer_path = m.get("path", "") + break + except: + pass + + # This is where Unsloth will perform the save + conversion operations + transformer_dir = os.path.join(save_directory, transformer_path) + # Ensure this path is absolute for consistent comparison later + transformer_dir = os.path.abspath(transformer_dir) + + if tokenizer is None: + tokenizer = self.tokenizer + + # 4. Patch environment to ensure Unsloth treats this embedding model correctly + @contextlib.contextmanager + def patch_unsloth_gguf_save(): + # Prevent deletion of the directory we just created via self.save_pretrained + original_rmtree = shutil.rmtree + try: + yield + finally: + shutil.rmtree = original_rmtree + + # 5. Call Unsloth's GGUF saver on the inner model targeting the transformer subdirectory + with patch_unsloth_gguf_save(): + result = unsloth_save_pretrained_gguf( + inner_model, + save_directory = transformer_dir, + tokenizer = tokenizer, + quantization_method = quantization_method, + first_conversion = first_conversion, + push_to_hub = False, # Force local first to move files + token = token, + max_shard_size = max_shard_size, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + ) + + # 6. Move GGUF files from the subdirectory (0_Transformer) to the root save_directory + gguf_files = result.get("gguf_files", []) + + new_gguf_locations = [] + + for gguf_file in gguf_files: + if os.path.exists(gguf_file): + filename = os.path.basename(gguf_file) + dest_path = os.path.join(save_directory, filename) + + # Convert to absolute path to avoid mixing relative/absolute in commonpath + abs_gguf_file = os.path.abspath(gguf_file) + + # Check if file is inside transformer_dir (subpath) + try: + is_subpath = ( + os.path.commonpath([abs_gguf_file, transformer_dir]) + == transformer_dir + ) + except ValueError: + # Can happen on Windows with different drives, or mix of absolute/relative (handled by abspath above) + is_subpath = False + + if is_subpath: + # If the GGUF file is inside the transformer_dir, move it out to root + shutil.move(gguf_file, dest_path) + new_gguf_locations.append(dest_path) + else: + # If it's elsewhere, move it to root if not already there + if os.path.abspath(dest_path) != abs_gguf_file: + shutil.move(gguf_file, dest_path) + new_gguf_locations.append(dest_path) + + # Update result with new locations + result["gguf_files"] = new_gguf_locations + + # 7. Add branding + try: + FastSentenceTransformer._add_unsloth_branding(save_directory) + + # Add GGUF details to README + readme_path = os.path.join(save_directory, "README.md") + if os.path.exists(readme_path): + with open(readme_path, "a", encoding = "utf-8") as f: + f.write("\n## GGUF Quantization\n") + f.write( + f"This model contains GGUF quantized versions in: {', '.join([os.path.basename(f) for f in new_gguf_locations])}\n" + ) + except: + pass + + # 8. Handle Push to Hub if requested + if push_to_hub: + if token is None: + token = get_token() + + api = HfApi(token = token) + repo_id = save_directory # Assuming save_directory is the repo name if pushing + + print(f"Unsloth: Uploading to {repo_id}...") + try: + api.create_repo( + repo_id = repo_id, exist_ok = True, private = kwargs.get("private", False) + ) + api.upload_folder( + folder_path = save_directory, + repo_id = repo_id, + commit_message = "Upload GGUF and SentenceTransformer model", + ) + print(f"Unsloth: Uploaded to https://huggingface.co/{repo_id}") + except Exception as e: + print(f"Unsloth: Upload failed: {e}") + + return result + + +class FastSentenceTransformer(FastModel): + @staticmethod + def _read_pooling_mode(model_name, token): + """ + Read the pooling mode from the modules.json file if it exists, otherwise return "mean". + """ + try: + if os.path.exists(model_name) and os.path.exists( + os.path.join(model_name, "modules.json") + ): + modules_json_path = os.path.join(model_name, "modules.json") + else: + modules_json_path = hf_hub_download( + model_name, "modules.json", token = token + ) + + with open(modules_json_path, "r") as f: + modules_config = json.load(f) + + pooling_config_path = None + for module in modules_config: + if module.get("type", "") == "sentence_transformers.models.Pooling": + pooling_path = module.get("path", "") + if pooling_path: + # try to find config.json for pooling module + if os.path.exists(model_name) and os.path.exists( + os.path.join(model_name, pooling_path, "config.json") + ): + pooling_config_path = os.path.join( + model_name, pooling_path, "config.json" + ) + else: + pooling_config_path = hf_hub_download( + model_name, + os.path.join(pooling_path, "config.json"), + token = token, + ) + break + + if pooling_config_path: + with open(pooling_config_path, "r") as f: + pooling_config = json.load(f) + # from here: + # https://github.com/huggingface/sentence-transformers/blob/main/sentence_transformers/models/Pooling.py#L43 + pooling_map = { + "pooling_mode_cls_token": "cls", + "pooling_mode_mean_tokens": "mean", + "pooling_mode_max_tokens": "max", + "pooling_mode_mean_sqrt_len_tokens": "mean_sqrt_len", + "pooling_mode_weightedmean_tokens": "weightedmean", + "pooling_mode_lasttoken": "lasttoken", + } + for config_key, mode in pooling_map.items(): + if pooling_config.get(config_key): + if mode != "mean": + print(f"Pooling mode detected as {mode}, updating...") + return mode + + except Exception as e: + print( + f"Failed to detect pooling mode, not a sentence-transformers model. Using default pooling mode 'mean', this may or may not work." + ) + return "mean" + + # should prolly be done upstream instead of this hackfest here + @staticmethod + def _patch_mpnet_v4(): + """ + Patch the MPNetModel to support gradient checkpointing. + Supports transformers 4. + """ + from transformers.models.mpnet import modeling_mpnet + + # add supports_gradient_checkpointing flag + modeling_mpnet.MPNetModel.supports_gradient_checkpointing = True + + # add _set_gradient_checkpointing method + def _set_gradient_checkpointing(self, module = None, value = True): + if module is None: + module = self.encoder + if isinstance(module, modeling_mpnet.MPNetEncoder): + module.gradient_checkpointing = value + + modeling_mpnet.MPNetModel._set_gradient_checkpointing = ( + _set_gradient_checkpointing + ) + + # patch MPNetEncoder.forward to support checkpointing + # based on: + # https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/mpnet/modeling_mpnet.py#L321 + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = False, + **kwargs, + ): + position_bias = self.compute_position_bias(hidden_states) + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # do gradient checkpointing if enabled and training + if getattr(self, "gradient_checkpointing", False) and self.training: + + def create_custom_forward(module): + # bog standard checkpoint + def custom_forward(*inputs): + return module(*inputs, output_attentions = output_attentions) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(layer_module), + hidden_states, + attention_mask, + head_mask[i] if head_mask is not None else None, + position_bias, + use_reentrant = True, # fix for torch 2.9 + ) + else: + # original code from here on + layer_outputs = layer_module( + hidden_states, + attention_mask, + head_mask[i] if head_mask is not None else None, + position_bias, + output_attentions = output_attentions, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, all_hidden_states, all_attentions] + if v is not None + ) + return BaseModelOutput( + last_hidden_state = hidden_states, + hidden_states = all_hidden_states, + attentions = all_attentions, + ) + + # assign the patched forward + modeling_mpnet.MPNetEncoder.forward = forward + + @staticmethod + def _patch_mpnet_v5(): + """ + Patch the MPNetModel to support gradient checkpointing. + Supports transformers 5. + """ + from transformers.models.mpnet import modeling_mpnet + + # add supports_gradient_checkpointing flag + modeling_mpnet.MPNetModel.supports_gradient_checkpointing = True + + # add _set_gradient_checkpointing method + def _set_gradient_checkpointing(self, module = None, value = True): + if module is None: + module = self.encoder + if isinstance(module, modeling_mpnet.MPNetEncoder): + module.gradient_checkpointing = value + + modeling_mpnet.MPNetModel._set_gradient_checkpointing = ( + _set_gradient_checkpointing + ) + + # patch MPNetEncoder.forward to support checkpointing + # based on: + # https://github.com/huggingface/transformers/blob/v5.0.0rc1/src/transformers/models/mpnet/modeling_mpnet.py#L284 + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = False, + **kwargs, + ): + position_bias = self.compute_position_bias(hidden_states) + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # do gradient checkpointing if enabled and training + if getattr(self, "gradient_checkpointing", False) and self.training: + + def create_custom_forward(module): + # checkpoint + def custom_forward(*inputs): + return module(*inputs, output_attentions = output_attentions) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(layer_module), + hidden_states, + attention_mask, + position_bias, + use_reentrant = True, # required for torch >= 2.9 + ) + else: + # original code from here on + layer_outputs = layer_module( + hidden_states, + attention_mask, + position_bias, + output_attentions, + **kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, all_hidden_states, all_attentions] + if v is not None + ) + return BaseModelOutput( + last_hidden_state = hidden_states, + hidden_states = all_hidden_states, + attentions = all_attentions, + ) + + modeling_mpnet.MPNetEncoder.forward = forward + + @staticmethod + def _patch_distilbert_v4(): + # change kwargs to positional args to be compatible with peft_utils + """ + Patch the forward method of the DistilBertModel to use positional arguments instead of keyword arguments. + Transformers 4 version. + """ + + # based on: + # https://github.com/huggingface/transformers/blob/v4.57.3/src/transformers/models/distilbert/modeling_distilbert.py#L666 + # original code from here on: + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time" + ) + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError( + "You have to specify either input_ids or inputs_embeds" + ) + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + head_mask_is_none = head_mask is None + # Prepare head mask if needed + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embeddings = self.embeddings( + input_ids, inputs_embeds + ) # (bs, seq_length, dim) + + if self.config._attn_implementation == "flash_attention_2": + attention_mask = ( + attention_mask + if (attention_mask is not None and 0 in attention_mask) + else None + ) + else: + if attention_mask is None: + attention_mask = torch.ones( + input_shape, device = device + ) # (bs, seq_length) + + if ( + self.config._attn_implementation == "sdpa" + and head_mask_is_none + and not output_attentions + ): + attention_mask = _prepare_4d_attention_mask_for_sdpa( + attention_mask, embeddings.dtype, tgt_len = input_shape[1] + ) + # patch here, change kwargs to positional args: + return self.transformer( + embeddings, + attention_mask, + head_mask, + output_attentions, + output_hidden_states, + return_dict, + ) + + modeling_distilbert.DistilBertModel.forward = forward + + @staticmethod + def _has_add_pooling_layer(config, auto_model_class = None): + """ + Checks if the model class supports the `add_pooling_layer` argument + """ + try: + if auto_model_class is None: + auto_model_class = AutoModel + # try to resolve the class + model_class = _get_model_class(config, auto_model_class._model_mapping) + + if model_class: + sig = inspect.signature(model_class.__init__) + return "add_pooling_layer" in sig.parameters + except: + pass + + return False + + @staticmethod + def _patch_distilbert_v5(): + """ + Patch the forward method of the DistilBertModel to use positional arguments instead of keyword arguments. + Transformers 5 version. + """ + # based on: + # https://github.com/huggingface/transformers/blob/v5.0.0rc1/src/transformers/models/distilbert/modeling_distilbert.py#L386 + # original code from here on: + from transformers.masking_utils import create_bidirectional_mask + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + **kwargs, + ): + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + + embeddings = self.embeddings(input_ids, inputs_embeds, position_ids) + + attention_mask = create_bidirectional_mask( + config = self.config, + input_embeds = embeddings, + attention_mask = attention_mask, + ) + + # patch here: unsloth gradient checkpointing hook needs positional arguments + return self.transformer( + embeddings, + attention_mask, + **kwargs, + ) + + modeling_distilbert.DistilBertModel.forward = forward + + @staticmethod + def _add_unsloth_tags(repo_id, token, tags = None): + """ + Add Unsloth and sentence-transformers tags to the Hugging Face Hub repository. + """ + from huggingface_hub import HfApi + + api = HfApi(token = token) + if tags is None: + tags = [] + tags.extend(["unsloth", "sentence-transformers"]) + try: + api.add_tags( + repo_id = repo_id, + tags = tags, + repo_type = "model", + ) + except: + pass + + @staticmethod + def _add_unsloth_branding(save_directory): + """ + Add Unsloth branding to the README.md file generated by sentence-transformers. + """ + readme_path = os.path.join(save_directory, "README.md") + if not os.path.exists(readme_path): + return + + with open(readme_path, "r", encoding = "utf-8") as f: + content = f.read() + + # add unsloth tag to frontmatter + if "---\ntags:\n" in content: + content = content.replace("---\ntags:\n", "---\ntags:\n- unsloth\n") + else: + # if tags exist but not right at start, use regex to append + pattern = r"(^tags:\s*\n)" + if re.search(pattern, content, re.MULTILINE): + content = re.sub( + pattern, r"\1- unsloth\n", content, count = 1, flags = re.MULTILINE + ) + + # add branding badge and text + branding = ( + "\n\nThis model was finetuned with [Unsloth](https://github.com/unslothai/unsloth).\n\n" + '[](https://github.com/unslothai/unsloth)\n' + ) + + # add to description + if "# SentenceTransformer" in content: + parts = content.split("# SentenceTransformer", 1) + content = parts[0] + "# SentenceTransformer" + branding + parts[1] + else: + content += branding + + with open(readme_path, "w", encoding = "utf-8") as f: + f.write(content) + + @staticmethod + def _module_path(model_name, token = None): + """ + Returns the path to the modules.json file or None + """ + try: + if os.path.exists(model_name) and os.path.isdir(model_name): + path = os.path.join(model_name, "modules.json") + return path if os.path.exists(path) else None + else: + try: + return hf_hub_download(model_name, "modules.json", token = token) + except: + return None + except: + return None + + @staticmethod + def _create_transformer_module( + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + ): + """Helper to create and configure a Transformer module.""" + from sentence_transformers.models import Transformer + + # prevents sentence-transformers from loading the model a second time, thanks Etherl + original_from_pretrained = AutoModel.from_pretrained + + def return_existing_model(*args, **kwargs): + return model + + try: + # Temporarily redirect AutoModel loading to return our pre-loaded model + AutoModel.from_pretrained = return_existing_model + + # Initialize Transformer + transformer_module = Transformer( + model_name, + max_seq_length = max_seq_length, + model_args = {"trust_remote_code": trust_remote_code}, + config_args = {"trust_remote_code": trust_remote_code}, + ) + finally: + # Restore original functionality immediately + AutoModel.from_pretrained = original_from_pretrained + + transformer_module.tokenizer = tokenizer + transformer_module.do_lower_case = getattr(tokenizer, "do_lower_case", False) + + # sentence-transformers only passes along known keys to model.forward + model_forward_params = list(inspect.signature(model.forward).parameters) + transformer_module.model_forward_params = set(model_forward_params) | { + "input_ids", + "attention_mask", + "token_type_ids", + "inputs_embeds", + } + + # determine max_seq_length if not provided + if max_seq_length is None: + if hasattr(model, "config") and hasattr( + model.config, "max_position_embeddings" + ): + max_seq_length = model.config.max_position_embeddings + elif hasattr(tokenizer, "model_max_length"): + max_seq_length = tokenizer.model_max_length + else: + max_seq_length = 512 + + transformer_module.max_seq_length = max_seq_length + transformer_module.config_keys = ["max_seq_length", "do_lower_case"] + transformer_module.save_in_root = True + + if hasattr(model, "config"): + model.config.tokenizer_class = tokenizer.__class__.__name__ + + return transformer_module + + @staticmethod + def _load_modules( + model_name, + token, + model, + tokenizer, + max_seq_length, + pooling_mode, + trust_remote_code = False, + ) -> tuple[OrderedDict, bool]: + """ + Load modules from modules.json if available, otherwise fallback to hard-coded modules. + + Returns: + tuple[OrderedDict, bool]: (modules, no_modules_json) + """ + from sentence_transformers.util import import_from_string, load_dir_path + from sentence_transformers.models import Pooling, Normalize + + modules = OrderedDict() + modules_json_path = FastSentenceTransformer._module_path(model_name, token) + + if modules_json_path: + with open(modules_json_path, encoding = "utf8") as f: + modules_config = json.load(f) + + for module_config in modules_config: + class_ref = module_config["type"] + name = module_config.get( + "name", str(module_config.get("idx", len(modules))) + ) + + if class_ref == "sentence_transformers.models.Transformer": + transformer_module = ( + FastSentenceTransformer._create_transformer_module( + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + ) + ) + modules[name] = transformer_module + else: + # load other modules (Pooling, Normalize, etc.) + module_path = module_config["path"] + if os.path.isdir(model_name): + load_path = os.path.join(model_name, module_path) + else: + try: + load_path = load_dir_path( + model_name, module_path, token = token + ) + except Exception as e: + print( + f"Unsloth Warning: Could not download module {module_path}: {e}" + ) + continue + + module_class = import_from_string(class_ref) + try: + module = module_class.load(load_path) + modules[name] = module + except Exception as e: + print( + f"Unsloth Warning: Failed to load module {name} ({class_ref}): {e}" + ) + + return modules, False + + # fallback if no modules.json (non sentence-transformers models) + print( + "Unsloth: No modules.json found, falling back to [Transformer, Pooling, Normalize]. This may or may not work." + ) + + transformer_module = FastSentenceTransformer._create_transformer_module( + model_name, model, tokenizer, max_seq_length, trust_remote_code + ) + modules["0"] = transformer_module + + hidden_size = getattr(model.config, "hidden_size", 768) + + if pooling_mode == "mean": + pooling_mode = FastSentenceTransformer._read_pooling_mode(model_name, token) + + modules["1"] = Pooling( + word_embedding_dimension = hidden_size, pooling_mode = pooling_mode + ) + modules["2"] = Normalize() + + return modules, True + + # Encoder model types that benefit from native torch.compile instead of Unsloth patching + ENCODER_MODEL_TYPES = { + "mpnet", + "bert", + "distilbert", + "modernbert", + "roberta", + "xlm-roberta", + "albert", + "electra", + } + + @staticmethod + def _estimate_compile_threshold( + model, + batch_size = None, + grad_accum = None, + max_seq_length = None, + ): + """ + Estimate the minimum training steps needed for torch.compile to be beneficial. + Returns the threshold with a 1.2x safety margin built in. + + Based on empirical benchmarks: + - Larger models have lower breakeven (more time saved per step) + - Warmup time scales with model size but speedup also increases + + Optional inputs (batch_size, grad_accum, max_seq_length) allow + a coarse pre-run adjustment. These are intentionally conservative + and avoid any runtime measurements. + """ + # Get parameter count from inner model + if hasattr(model, "__getitem__"): + try: + inner = model[0].auto_model + params = sum(p.numel() for p in inner.parameters()) + except: + params = 100_000_000 # Default to 100M if can't determine + else: + params = sum(p.numel() for p in model.parameters()) + + model_type = None + try: + if "inner" in locals(): + model_type = getattr(getattr(inner, "config", None), "model_type", None) + except Exception: + model_type = None + if isinstance(model_type, str): + model_type = model_type.lower() + + params_m = params / 1e6 + + # Empirical formula based on benchmarks with batch_size=2, grad_accum=4 + # Small models: high fixed overhead, lower speedup + # Large models: warmup scales but speedup is significant + if params_m < 50: + estimated_warmup = 35 + params_m * 0.3 + base_speedup = 1.35 + elif params_m < 200: + estimated_warmup = 12 + params_m * 0.03 + base_speedup = 1.75 + else: + estimated_warmup = 15 + params_m * 0.04 + base_speedup = 1.60 + + # Estimate time per step (ms) and time saved + naive_ms = 50 + params_m * 1.0 + compiled_ms = naive_ms / base_speedup + time_saved_per_step_s = (naive_ms - compiled_ms) / 1000 + + if time_saved_per_step_s > 0: + breakeven = estimated_warmup / time_saved_per_step_s + else: + breakeven = float("inf") + + # Return threshold with 1.2x safety margin + threshold = breakeven * 1.2 + + # Optional adjustment based on expected work per step. + # This uses only pre-run information (batch size, grad accum, seq length). + generic_scale = 1.0 + fast_scale = 1.0 + if ( + batch_size is not None + or grad_accum is not None + or max_seq_length is not None + ): + try: + bs = int(batch_size) if batch_size is not None else 2 + ga = int(grad_accum) if grad_accum is not None else 4 + seq = int(max_seq_length) if max_seq_length is not None else 512 + except Exception: + bs, ga, seq = 2, 4, 512 + + bs = max(1, bs) + ga = max(1, ga) + # Guard against unbounded tokenizer.model_max_length + seq = max(64, min(seq, 8192)) + + ref_bs, ref_ga, ref_seq = 2, 4, 512 + + # Generic path: lighter scaling, less conservative than params-only. + ga_scale = (ref_ga / ga) ** 1.0 + bs_seq_scale = ((ref_bs * ref_seq) / (bs * seq)) ** 0.15 + generic_scale = 0.35 * ga_scale * bs_seq_scale + generic_scale = max(0.05, min(generic_scale, 5.0)) + + # Fast encoder path: stronger scaling based on observed behavior. + fast_ga_scale = (ref_ga / ga) ** 1.5 + fast_bs_seq_scale = ((ref_bs * ref_seq) / (bs * seq)) ** 0.25 + fast_scale = 0.2 * fast_ga_scale * fast_bs_seq_scale + fast_scale = max(0.05, min(fast_scale, 5.0)) + + # Conservative safety factors: generic is less conservative than fast. + generic_threshold = threshold * generic_scale * 1.25 + + is_fast_type = ( + isinstance(model_type, str) + and model_type in FastSentenceTransformer.ENCODER_MODEL_TYPES + ) + if is_fast_type: + fast_threshold = threshold * fast_scale * 1.5 + # Prefer the smaller (less conservative) of the two estimates. + final_threshold = min(generic_threshold, fast_threshold) + else: + final_threshold = generic_threshold + + # Reduce mpnet overestimation slightly. + if model_type == "mpnet": + final_threshold *= 0.7 + + # Lower bound to avoid compiling on extremely short runs. + return int(max(20, final_threshold)) + + @staticmethod + def _apply_torch_compile(model, mode = "default"): + """ + Apply torch.compile to a SentenceTransformer model. + Includes workaround for accelerate's unwrap_model bug. + """ + if hasattr(model, "__getitem__"): + inner_model = model[0].auto_model + compiled = torch.compile(inner_model, mode = mode) + model[0].auto_model = compiled + # Fix for accelerate unwrap_model bug: + # When SentenceTransformer contains a compiled inner model, + # accelerate checks has_compiled_regions() which returns True, + # then tries to access model.__dict__["_orig_mod"] which fails. + # This workaround sets _orig_mod to satisfy accelerate. + model.__dict__["_orig_mod"] = model + else: + model = torch.compile(model, mode = mode) + return model + + @staticmethod + def from_pretrained( + model_name, + max_seq_length = None, + dtype = None, + load_in_4bit = False, # Changed default: 4-bit is slow for encoders + load_in_8bit = False, + load_in_16bit = True, # Changed default: 16-bit is optimal for encoders + full_finetuning = False, + token = None, + device_map = "sequential", + rope_scaling = None, + fix_tokenizer = True, + trust_remote_code = False, + use_gradient_checkpointing = False, # Changed default: conflicts with torch.compile + resize_model_vocab = None, + revision = None, + use_exact_model_name = False, + offload_embedding = False, + random_state = 3407, + max_lora_rank = 64, + disable_log_stats = True, + qat_scheme = None, + unsloth_tiled_mlp = False, + pooling_mode = "mean", + for_inference = False, + **kwargs, + ): + try: + from sentence_transformers import SentenceTransformer + from sentence_transformers.models import Transformer, Pooling, Normalize + except ImportError: + raise ImportError( + "Unsloth: To use `FastSentenceTransformer`, you must install `sentence-transformers`.\n" + "Run `pip install sentence-transformers` to install it." + ) + + # if for_inference == True, skip Unsloth optimizations to avoid torch compile issues + if for_inference: + st_device = device_map + if isinstance(st_device, dict) or ( + isinstance(st_device, str) and st_device in ["auto", "sequential"] + ): + st_device = None + + # this was added because when loading for inference it was defaulting to float32 + # propagate dtype to model_kwargs, default to "auto" + model_kwargs = kwargs.get("model_kwargs", {}) + model_kwargs["dtype"] = dtype if dtype is not None else "auto" + + # filter kwargs for SentenceTransformer + st_kwargs = { + "device": st_device, + "trust_remote_code": trust_remote_code, + "token": token, + "revision": revision, + "model_kwargs": model_kwargs, + } + + # add other known kwargs if present + known_keys = [ + "cache_folder", + "truncate_dim", + "tokenizer_kwargs", + "config_kwargs", + ] + for k in known_keys: + if k in kwargs: + st_kwargs[k] = kwargs[k] + + st_model = SentenceTransformer(model_name, **st_kwargs) + return st_model + + # sanity check, thanks Etherl: + if full_finetuning and (load_in_4bit or load_in_8bit): + print( + "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA." + ) + load_in_4bit = False + load_in_8bit = False + load_in_fp8 = False + load_in_16bit = False + + if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2: + raise RuntimeError( + "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n" + "Also, we by default set `load_in_16bit = True`.\n" + "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n" + "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`" + ) + + if "auto_model" not in kwargs: + kwargs["auto_model"] = AutoModel + + transformers4 = Version(transformers.__version__).major < 5 + model_type = "" + config = None + try: + config = AutoConfig.from_pretrained( + model_name, token = token, trust_remote_code = trust_remote_code + ) + model_type = getattr(config, "model_type", "") + except: + pass + + # Fast encoder path: Use native torch.compile for encoder models (6x speedup) + # This bypasses Unsloth's auto-compiler which adds @torch.compiler.disable decorators + # that interfere with torch.compile and cause runtime errors for encoder models. + # NOTE: The old Unsloth path is BROKEN for encoder models with torch 2.9+ due to + # conflicting @torch.compile and @torch.compiler.disable decorators. + # Set UNSLOTH_COMPILE_DISABLE=1 to disable torch.compile and use the old path. + is_encoder_model = ( + model_type.lower() in FastSentenceTransformer.ENCODER_MODEL_TYPES + ) + use_fast_encoder = os.environ.get("UNSLOTH_COMPILE_DISABLE", "0") != "1" + if use_fast_encoder and is_encoder_model: + # torch.compile mode: "default" is safest for PEFT/LoRA training + # Note: "reduce-overhead" uses CUDA Graphs which is incompatible with PEFT + compile_mode = "default" + + # Determine dtype - handle float16 machines that don't support bfloat16 + if dtype is None: + if load_in_16bit: + dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 + else: + dtype = torch.float32 + elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16: + print( + "Unsloth: Device does not support bfloat16. Using float16 instead." + ) + dtype = torch.float16 + + # Determine device + st_device = device_map + if isinstance(st_device, dict) or ( + isinstance(st_device, str) and st_device in ["auto", "sequential"] + ): + st_device = "cuda" + + # Check if model supports SDPA (Scaled Dot Product Attention) for extra speedup + supports_sdpa = False + if config is not None: + try: + model_class = _get_model_class( + config, kwargs.get("auto_model", AutoModel)._model_mapping + ) + supports_sdpa = getattr(model_class, "_supports_sdpa", False) + except: + pass + + # Build model_kwargs for SentenceTransformer + model_kwargs = {"torch_dtype": dtype} + + # Enable SDPA if supported (1.2x extra speedup on top of torch.compile) + if supports_sdpa: + model_kwargs["attn_implementation"] = "sdpa" + + # Print optimization status + sdpa_str = " + SDPA" if supports_sdpa else "" + if load_in_4bit: + print( + f"Unsloth: Using fast encoder path for {model_type} with 4-bit quantization{sdpa_str}" + ) + else: + print( + f"Unsloth: Using fast encoder path for {model_type} (torch.compile{sdpa_str})" + ) + + # Handle 4-bit quantization via BitsAndBytesConfig + if load_in_4bit: + from transformers import BitsAndBytesConfig + + bnb_config = BitsAndBytesConfig( + load_in_4bit = True, + bnb_4bit_compute_dtype = dtype, + bnb_4bit_quant_type = "nf4", + bnb_4bit_use_double_quant = True, + ) + model_kwargs["quantization_config"] = bnb_config + # When using quantization, device must be handled by accelerate + st_device = None + + # Handle gradient checkpointing - warn user it conflicts with torch.compile + _use_gc = use_gradient_checkpointing + if _use_gc and _use_gc != False: + print( + "Unsloth Warning: Gradient checkpointing is incompatible with torch.compile." + ) + print("Disabling torch.compile to enable gradient checkpointing.") + compile_mode = None # Disable compilation + + is_mpnet = "mpnet" == model_type.lower() + + if is_mpnet and transformers4: + FastSentenceTransformer._patch_mpnet_v4() + elif is_mpnet: + FastSentenceTransformer._patch_mpnet_v5() + + # Load via native SentenceTransformer (bypasses Unsloth patching) + st_model = SentenceTransformer( + model_name, + device = st_device, + trust_remote_code = trust_remote_code, + token = token, + revision = revision, + model_kwargs = model_kwargs, + ) + + # Store metadata for get_peft_model + st_model._unsloth_fast_encoder = True + st_model._compile_mode = compile_mode + st_model._dtype = dtype + st_model._load_in_4bit = load_in_4bit + st_model.no_modules = False + + # Add save methods + def _save_pretrained_merged(self, save_directory, **save_kwargs): + self.save_pretrained(save_directory) + tokenizer = save_kwargs.pop("tokenizer", self.tokenizer) + if hasattr(self[0], "auto_model"): + inner = self[0].auto_model + # Handle compiled model + if hasattr(inner, "_orig_mod"): + inner = inner._orig_mod + if hasattr(inner, "merge_and_unload"): + merged = inner.merge_and_unload() + merged.save_pretrained(save_directory) + elif hasattr(inner, "save_pretrained"): + inner.save_pretrained(save_directory) + if tokenizer is not None: + tokenizer.save_pretrained(save_directory) + FastSentenceTransformer._add_unsloth_branding(save_directory) + + st_model.save_pretrained_merged = types.MethodType( + _save_pretrained_merged, st_model + ) + + st_model.save_pretrained_torchao = types.MethodType( + _save_pretrained_torchao, st_model + ) + + st_model.save_pretrained_gguf = types.MethodType( + _save_pretrained_gguf, st_model + ) + + def _push_to_hub_merged(self, repo_id, **push_kwargs): + hub_token = push_kwargs.get("token", None) or get_token() + if hub_token is None: + raise ValueError("No HF token provided") + api = HfApi(token = hub_token) + try: + api.create_repo( + repo_id = repo_id, + private = push_kwargs.get("private"), + exist_ok = True, + repo_type = "model", + ) + except: + pass + FastSentenceTransformer._add_unsloth_tags(repo_id, hub_token) + with tempfile.TemporaryDirectory() as temp_dir: + self.save_pretrained_merged(temp_dir, **push_kwargs) + api.upload_folder( + folder_path = temp_dir, + repo_id = repo_id, + commit_message = push_kwargs.get( + "commit_message", "Upload model" + ), + ) + print(f"Unsloth: Pushed to https://huggingface.co/{repo_id}") + + st_model.push_to_hub_merged = types.MethodType( + _push_to_hub_merged, st_model + ) + + return st_model + + # Warn if using 4-bit with encoder (slow due to dequantization overhead) + if is_encoder_model and load_in_4bit: + print( + "Unsloth Warning: 4-bit quantization adds ~2.3x overhead for encoder models." + ) + print("Consider using load_in_16bit=True for better performance.") + + # check if the model supports add_pooling_layer + if "add_pooling_layer" not in kwargs: + supported = FastSentenceTransformer._has_add_pooling_layer( + config, kwargs.get("auto_model", AutoModel) + ) + if supported: + kwargs["add_pooling_layer"] = False + + # forces fp8 to be False since it's not supported + fp8 = kwargs.pop("load_in_fp8", None) + if fp8: + logging.info("Unsloth: Disabling fp8 for model") + load_in_fp8 = False + + # this is a fix for Snowflake/snowflake-arctic-embed-l-v2.0 + # it has pooler weights which we don't care about for training, + # however unsloth throws an exception if "UNSLOTH_WARN_UNINITIALIZED" == 1 and it sees unused weights + old_environ = os.environ.get("UNSLOTH_WARN_UNINITIALIZED", "1") + os.environ["UNSLOTH_WARN_UNINITIALIZED"] = "0" + + is_distilbert = "distilbert" == model_type.lower() + is_mpnet = "mpnet" == model_type.lower() + + if is_distilbert and transformers4: + FastSentenceTransformer._patch_distilbert_v4() + elif is_distilbert: + FastSentenceTransformer._patch_distilbert_v5() + elif is_mpnet and transformers4: + FastSentenceTransformer._patch_mpnet_v4() + elif is_mpnet: + FastSentenceTransformer._patch_mpnet_v5() + + # check if modules.json exists - if not, force 16-bit training + # why? because i have to implement saving myself for these models, and i don't feel like adding dequantization + # to the save_pretrained_merged for a model that really should be trained in 16-bit anyway + has_modules_json = ( + FastSentenceTransformer._module_path(model_name, token) is not None + ) + + if not has_modules_json and load_in_4bit: + print( + "Unsloth: No modules.json found. This is not a sentence-transformers model.\n" + "Forcing 16-bit loading to simplify merged model saving." + ) + load_in_4bit = False + load_in_16bit = True + + try: + model, tokenizer = FastModel.from_pretrained( + model_name = model_name, + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = load_in_4bit, + load_in_8bit = load_in_8bit, + load_in_16bit = load_in_16bit, + full_finetuning = full_finetuning, + token = token, + device_map = device_map, + rope_scaling = rope_scaling, + fix_tokenizer = fix_tokenizer, + trust_remote_code = trust_remote_code, + use_gradient_checkpointing = use_gradient_checkpointing, + resize_model_vocab = resize_model_vocab, + revision = revision, + return_logits = False, + use_exact_model_name = use_exact_model_name, + offload_embedding = offload_embedding, + random_state = random_state, + max_lora_rank = max_lora_rank, + disable_log_stats = disable_log_stats, + qat_scheme = qat_scheme, + load_in_fp8 = load_in_fp8, + unsloth_tiled_mlp = unsloth_tiled_mlp, + **kwargs, + ) + finally: + os.environ["UNSLOTH_WARN_UNINITIALIZED"] = old_environ + + # try to load modules, otherwise fallback to old hard-coded modules + from sentence_transformers import SentenceTransformer + + modules, no_modules = FastSentenceTransformer._load_modules( + model_name, + token, + model, + tokenizer, + max_seq_length, + pooling_mode, + trust_remote_code = trust_remote_code, + ) + + st_device = device_map + if isinstance(st_device, dict) or ( + isinstance(st_device, str) and st_device in ["auto", "sequential"] + ): + st_device = None + + st_model = SentenceTransformer(modules = modules, device = st_device) + st_model.no_modules = no_modules + + def _save_pretrained_merged(self, save_directory, **kwargs): + # check which adapter files exist before save_pretrained + adapter_files = ["adapter_model.safetensors", "adapter_config.json"] + existing_before = { + f + for f in adapter_files + if os.path.exists(os.path.join(save_directory, f)) + } + + # sentence-transformers config and modules only get saved if we call save_pretrained + self.save_pretrained(save_directory) + + # remove LoRA adapters only if they were created by save_pretrained (not pre-existing) + for file in adapter_files: + if file not in existing_before: + try: + os.remove(os.path.join(save_directory, file)) + except: + pass + + tokenizer = kwargs.pop("tokenizer", self.tokenizer) + if self.no_modules: + # fallback for non-sentence-transformers models + print( + "Unsloth: No modules detected. Using standard merge_and_unload for saving..." + ) + safe_kwargs = kwargs.copy() + # filter out Unsloth-specific args that are not in huggingface's save_pretrained + unsloth_args = [ + "save_method", + "temporary_location", + "maximum_memory_usage", + ] + for k in unsloth_args: + safe_kwargs.pop(k, None) + + merged_model = self[0].auto_model.merge_and_unload() + merged_model.save_pretrained(save_directory, **safe_kwargs) + if tokenizer is not None: + tokenizer.save_pretrained(save_directory) + else: + self[0].auto_model.save_pretrained_merged( + save_directory, tokenizer = tokenizer, **kwargs + ) + + # add Unsloth branding to the generated README + try: + FastSentenceTransformer._add_unsloth_branding(save_directory) + except Exception as e: + print(f"Unsloth Warning: Failed to add branding to README: {e}") + + st_model.save_pretrained_merged = types.MethodType( + _save_pretrained_merged, st_model + ) + + st_model.save_pretrained_torchao = types.MethodType( + _save_pretrained_torchao, st_model + ) + + st_model.save_pretrained_gguf = types.MethodType( + _save_pretrained_gguf, st_model + ) + + def _push_to_hub_merged(self, repo_id, **kwargs): + token = kwargs.get("token", None) or get_token() + if token is None: + raise ValueError( + "No HF token provided. Please provide a token or login with `hf auth login`" + ) + private = kwargs.get("private", None) + commit_message = kwargs.get("commit_message", "Upload model") + + from huggingface_hub import HfApi + + api = HfApi(token = token) + try: + api.create_repo( + repo_id = repo_id, + private = private, + exist_ok = True, + repo_type = "model", + ) + except: + pass + + # order doesn't seem to matter for this after repo creation... + FastSentenceTransformer._add_unsloth_tags(repo_id, token) + + with tempfile.TemporaryDirectory() as temp_dir: + self.save_pretrained_merged(temp_dir, **kwargs) + api.upload_folder( + folder_path = temp_dir, + repo_id = repo_id, + commit_message = commit_message, + ) + print( + f"Unsloth: Successfully pushed merged model to https://huggingface.co/{repo_id}" + ) + + st_model.push_to_hub_merged = types.MethodType(_push_to_hub_merged, st_model) + return st_model + + @staticmethod + def get_peft_model( + model, + r = 16, + target_modules = [ + "query", + "key", + "value", + "dense", + ], + lora_alpha = 16, + lora_dropout = 0.0, + bias = "none", + layers_to_transform = None, + layers_pattern = None, + use_gradient_checkpointing = False, # Changed default: conflicts with torch.compile + random_state = 3407, + max_seq_length = 2048, + use_rslora = False, + modules_to_save = None, + init_lora_weights = True, + loftq_config = {}, + **kwargs, + ): + from sentence_transformers import SentenceTransformer + from peft import LoraConfig, get_peft_model as peft_get_peft_model + + if "task_type" not in kwargs: + kwargs["task_type"] = "FEATURE_EXTRACTION" + print("Setting task_type to FEATURE_EXTRACTION") + + if isinstance(model, SentenceTransformer): + # Check if this is a fast encoder model (uses torch.compile instead of Unsloth patching) + is_fast_encoder = getattr(model, "_unsloth_fast_encoder", False) + + if is_fast_encoder: + # Fast encoder path: Use native PEFT + torch.compile (6x speedup) + transformer_module = model[0] + inner_model = transformer_module.auto_model + + # Check if model is quantized (4-bit/8-bit) + is_quantized = ( + getattr(inner_model, "is_quantized", False) + or getattr(inner_model.config, "quantization_config", None) + is not None + ) + + # Track if gradient checkpointing was actually enabled + gc_enabled = False + + # this is needed when from_pretrained was called without gradient + # checkpointing but get_peft_model requests it + if use_gradient_checkpointing and use_gradient_checkpointing != False: + import transformers + from packaging.version import Version + + transformers4 = Version(transformers.__version__).major < 5 + model_type = getattr(inner_model.config, "model_type", "").lower() + + if model_type == "mpnet" and transformers4: + FastSentenceTransformer._patch_mpnet_v4() + elif model_type == "mpnet": + FastSentenceTransformer._patch_mpnet_v5() + + # Prepare for k-bit training if quantized + if is_quantized: + from ._utils import prepare_model_for_kbit_training + + _gc_for_kbit = ( + use_gradient_checkpointing + if use_gradient_checkpointing + else False + ) + try: + inner_model = prepare_model_for_kbit_training( + inner_model, + use_gradient_checkpointing = _gc_for_kbit, + ) + print("Unsloth: Prepared quantized model for k-bit training") + gc_enabled = bool(_gc_for_kbit) + except ValueError as e: + if "does not support gradient checkpointing" in str(e): + # Model doesn't support gradient checkpointing, disable it + print( + f"Unsloth Warning: {inner_model.__class__.__name__} does not support gradient checkpointing. Skipping." + ) + inner_model = prepare_model_for_kbit_training( + inner_model, + use_gradient_checkpointing = False, + ) + print( + "Unsloth: Prepared quantized model for k-bit training (without gradient checkpointing)" + ) + else: + raise + + # Enable gradient checkpointing if requested (only for non-quantized, since prepare_model handles it) + elif use_gradient_checkpointing and use_gradient_checkpointing != False: + if hasattr(inner_model, "gradient_checkpointing_enable"): + try: + inner_model.gradient_checkpointing_enable() + print("Unsloth: Enabled gradient checkpointing") + gc_enabled = True + except ValueError as e: + if "does not support gradient checkpointing" in str(e): + print( + f"Unsloth Warning: {inner_model.__class__.__name__} does not support gradient checkpointing. Skipping." + ) + + # Create LoRA config + lora_config = LoraConfig( + r = r, + lora_alpha = lora_alpha, + target_modules = target_modules, + lora_dropout = lora_dropout, + bias = bias, + task_type = kwargs.get("task_type", "FEATURE_EXTRACTION"), + ) + + # Apply PEFT directly (not through FastModel) + peft_model = peft_get_peft_model(inner_model, lora_config) + + # Apply QAT if specified + qat_scheme = kwargs.get("qat_scheme", None) + if qat_scheme is not None: + from ._utils import _prepare_model_for_qat + + peft_model = _prepare_model_for_qat(peft_model, qat_scheme) + + # Determine compile mode (only if not using gradient checkpointing) + compile_mode = getattr(model, "_compile_mode", "default") + # Re-enable torch.compile if gradient checkpointing was requested but couldn't be enabled + if compile_mode is None and not gc_enabled: + compile_mode = "default" + print( + "Unsloth: Re-enabling torch.compile since gradient checkpointing is not supported" + ) + + # Re-assign the peft model back to the transformer module + transformer_module.auto_model = peft_model + + # Store compile info for auto-compile at trainer time + # torch.compile is deferred until training starts so we can check max_steps + if compile_mode is not None: + model._compile_mode = compile_mode + model._compile_threshold = ( + FastSentenceTransformer._estimate_compile_threshold(model) + ) + # Flag to indicate compile has not been applied yet + model._compile_pending = True + print( + f"Unsloth: torch.compile will be applied automatically if max_steps > {model._compile_threshold}" + ) + else: + model._compile_mode = None + model._compile_pending = False + print( + "Unsloth: torch.compile disabled (gradient checkpointing enabled)" + ) + + return model + + # Original path for non-fast-encoder models + transformer_module = model[0] + inner_model = transformer_module.auto_model + + peft_model = FastModel.get_peft_model( + model = inner_model, + r = r, + target_modules = target_modules, + lora_alpha = lora_alpha, + lora_dropout = lora_dropout, + bias = bias, + layers_to_transform = layers_to_transform, + layers_pattern = layers_pattern, + use_gradient_checkpointing = use_gradient_checkpointing, + random_state = random_state, + max_seq_length = max_seq_length, + use_rslora = use_rslora, + modules_to_save = modules_to_save, + init_lora_weights = init_lora_weights, + loftq_config = loftq_config, + **kwargs, + ) + + # re-assign the peft model back to the transformer module + transformer_module.auto_model = peft_model + return model + else: + return FastModel.get_peft_model( + model = model, + r = r, + target_modules = target_modules, + lora_alpha = lora_alpha, + lora_dropout = lora_dropout, + bias = bias, + layers_to_transform = layers_to_transform, + layers_pattern = layers_pattern, + use_gradient_checkpointing = use_gradient_checkpointing, + random_state = random_state, + max_seq_length = max_seq_length, + use_rslora = use_rslora, + modules_to_save = modules_to_save, + init_lora_weights = init_lora_weights, + loftq_config = loftq_config, + **kwargs, + ) + + +def _patch_sentence_transformer_trainer(): + """ + Patch SentenceTransformerTrainer to automatically apply torch.compile + when training steps exceed the breakeven threshold. + + This is called automatically when this module is imported. + """ + try: + from sentence_transformers import SentenceTransformerTrainer + except ImportError: + return # sentence_transformers not installed + + if getattr(SentenceTransformerTrainer, "_unsloth_auto_compile_patched", False): + return # Already patched + + from functools import wraps + + _original_init = SentenceTransformerTrainer.__init__ + + @wraps(_original_init) + def _patched_init(self, *args, **kwargs): + # Extract model and training_args + model = kwargs.get("model") or (args[0] if args else None) + training_args = kwargs.get("args") or (args[1] if len(args) > 1 else None) + + # Check if model has pending compile + if ( + model is not None + and training_args is not None + and getattr(model, "_compile_pending", False) + ): + max_steps = getattr(training_args, "max_steps", -1) + compile_mode = getattr(model, "_compile_mode", "default") + + # Re-estimate threshold now that training args are available + batch_size = getattr(training_args, "per_device_train_batch_size", None) + grad_accum = getattr(training_args, "gradient_accumulation_steps", None) + max_seq_length = getattr(model, "max_seq_length", None) + if max_seq_length is None and hasattr(model, "__getitem__"): + try: + max_seq_length = getattr(model[0], "max_seq_length", None) + except Exception: + max_seq_length = None + if max_seq_length is None: + tokenizer = getattr(model, "tokenizer", None) + max_seq_length = ( + getattr(tokenizer, "model_max_length", None) + if tokenizer is not None + else None + ) + + threshold = FastSentenceTransformer._estimate_compile_threshold( + model, + batch_size = batch_size, + grad_accum = grad_accum, + max_seq_length = max_seq_length, + ) + model._compile_threshold = threshold + + if max_steps > 0 and max_steps >= threshold: + print( + f"Unsloth: Auto-compiling model ({max_steps} steps >= {threshold} threshold)" + ) + FastSentenceTransformer._apply_torch_compile(model, mode = compile_mode) + model._compile_pending = False + elif max_steps > 0: + print( + f"Unsloth: Skipping torch.compile ({max_steps} steps < {threshold} threshold)" + ) + model._compile_pending = False + + # Call original __init__ + _original_init(self, *args, **kwargs) + + SentenceTransformerTrainer.__init__ = _patched_init + SentenceTransformerTrainer._unsloth_auto_compile_patched = True + + +# Auto-patch trainer on module import +_patch_sentence_transformer_trainer() diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index b78b190bcb..6835f2e986 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -68,11 +68,9 @@ import functools import os import gc import math -import functools from typing import Optional, Tuple, List, Union import re, inspect, sys import contextlib -import types try: from huggingface_hub.utils import get_token @@ -108,7 +106,7 @@ PRE_COMPILE_INFERENCE = [ "gpt_oss", ] -from transformers import GenerationConfig, CompileConfig, HybridCache, AutoConfig +from transformers import GenerationConfig, CompileConfig, AutoConfig try: from transformers import PreTrainedConfig @@ -119,8 +117,6 @@ except: HAS_TORCH_DTYPE = "torch_dtype" in PretrainedConfig.__doc__ -from transformers import GenerationConfig, CompileConfig, HybridCache - _compile_config = CompileConfig( fullgraph = False, dynamic = None, @@ -149,7 +145,7 @@ def unsloth_base_fast_generate( elif "input_ids" in kwargs: input_ids = kwargs["input_ids"] elif "input" in kwargs: - input_ids = kwargs["input_ids"] + input_ids = kwargs["input"] elif "input_features" in kwargs: input_ids = kwargs["input_features"] elif "input_embeds" in kwargs: @@ -158,7 +154,7 @@ def unsloth_base_fast_generate( input_ids = kwargs["inputs"] else: key = next(iter(kwargs.keys())) - if type(kwargs["key"]) is not torch.Tensor: + if type(kwargs[key]) is not torch.Tensor: raise TypeError("Unsloth: You need to pass in input_ids to .generate!") input_ids = kwargs[key] assert type(input_ids) is torch.Tensor @@ -531,6 +527,7 @@ class FastBaseModel: del kwargs["attn_implementation"] bnb_config = None + user_quantization_config = kwargs.get("quantization_config", None) if full_finetuning and (load_in_4bit or load_in_8bit): print( "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA." @@ -598,7 +595,8 @@ class FastBaseModel: ): pass else: - kwargs["quantization_config"] = bnb_config + if user_quantization_config is None: + kwargs["quantization_config"] = bnb_config else: if auto_config is None: auto_config = AutoConfig.from_pretrained( @@ -643,7 +641,8 @@ class FastBaseModel: ) except: pass - kwargs["quantization_config"] = quantization_config + if user_quantization_config is None: + kwargs["quantization_config"] = quantization_config # Check if using forced float32 - we load it in bfloat16, then cast to float16! torch_dtype = dtype @@ -675,7 +674,7 @@ class FastBaseModel: **kwargs, ) if hasattr(model, "generate"): - model.fast_generate = model.generate + model.fast_generate = make_fast_generate_wrapper(model.generate) model.fast_generate_batches = error_out_no_vllm if offload_embedding: if bool( @@ -718,9 +717,13 @@ class FastBaseModel: if full_finetuning: max_lora_rank = max(get_lora_supported_ranks()) raise NotImplementedError( - f"Unsloth: `fast_inference = True` does not yet support `full_finetuning = True`.\n" - f"Use LoRA rank `r = {max_lora_rank}` as the closest replacement for full finetuning with Unsloth for RL." + "Unsloth: `fast_inference=True` cannot be used together with `full_finetuning=True`.\n" + "Reason: fast_inference is optimized for inference-only workflows and " + "does not currently support full fine-tuning.\n" + "Workaround: disable fast_inference, or use parameter-efficient fine-tuning " + f"(e.g. LoRA with rank r={max_lora_rank})." ) + model_config.model_name = model_name if fast_inference: @@ -936,6 +939,7 @@ class FastBaseModel: task_type = TaskType.CAUSAL_LM, temporary_location = "_unsloth_temporary_saved_buffers", qat_scheme = None, + ensure_weight_tying = False, # [TODO] Add `ensure_weight_tying` for `modules_to_save` for vision models **kwargs, ): if os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1": @@ -1267,7 +1271,7 @@ class FastBaseModel: # Since transformers 4.53, must turn on explicitly for module in model.modules(): if hasattr(module, "gradient_checkpointing"): - module.gradient_checkpointing = True + module.gradient_checkpointing = use_gradient_checkpointing # Also re-enable training for embeddings for NEFTune if hasattr(model, "get_input_embeddings"): diff --git a/unsloth/save.py b/unsloth/save.py index 3a275cf0c3..071e032c53 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -131,6 +131,13 @@ ALLOWED_QUANTS = { } +def has_curl(): + return shutil.which("curl") is not None + + +CURL_FLAG = "-DLLAMA_CURL=ON" if has_curl() else "-DLLAMA_CURL=OFF" + + def print_quantization_methods(): for key, value in ALLOWED_QUANTS.items(): print(f'"{key}" ==> {value}') @@ -547,7 +554,7 @@ def unsloth_save_model( elif mb_found: sharded_ram_usage = int(mb_found.group(1)) * 1024 * 1024 elif type(max_shard_size) is int: - sharded_ram_usage = sharded_ram_usage + sharded_ram_usage = max_shard_size # Switch to our fast saving modules if it's a slow PC! n_cpus = psutil.cpu_count(logical = False) @@ -872,15 +879,16 @@ def install_llama_cpp_make_non_blocking(): IS_CMAKE = False if check == 0: # Uses old MAKE - n_jobs = max(int(psutil.cpu_count() * 1.5), 1) + n_jobs = max(int((psutil.cpu_count() or 1) * 1.5), 1) full_command = ["make", "all", "-j" + str(n_jobs), "-C", "llama.cpp"] IS_CMAKE = False else: # Uses new CMAKE - n_jobs = max(int(psutil.cpu_count()), 1) # Use less CPUs since 1.5x faster + n_jobs = max(int(psutil.cpu_count() or 1), 1) # Use less CPUs since 1.5x faster check = os.system( - "cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF -DLLAMA_CURL=ON" + f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}" ) + if check != 0: raise RuntimeError( f"*** Unsloth: Failed compiling llama.cpp using os.system(...) with error {check}. Please report this ASAP!" @@ -986,16 +994,17 @@ def install_llama_cpp_old(version = -10): # Try using MAKE commands = [ "make clean -C llama.cpp", - f"make all -j{psutil.cpu_count()*2} -C llama.cpp", + f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp", ] if try_execute(commands) == "CMAKE": # Instead use CMAKE commands = [ - "cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF -DLLAMA_CURL=ON", - f"cmake --build llama.cpp/build --config Release -j{psutil.cpu_count()*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}", + f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}", + f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}", "cp llama.cpp/build/bin/llama-* llama.cpp", "rm -rf llama.cpp/build", ] + try_execute(commands) # Check if successful @@ -1031,14 +1040,14 @@ def install_llama_cpp_blocking(use_cuda = False): "make clean -C llama.cpp", # https://github.com/ggerganov/llama.cpp/issues/7062 # Weirdly GPU conversion for GGUF breaks?? - # f"{use_cuda} make all -j{psutil.cpu_count()*2} -C llama.cpp", - f"make all -j{psutil.cpu_count()*2} -C llama.cpp", + # f"{use_cuda} make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp", + f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp", ] if try_execute(commands) == "CMAKE": # Instead use CMAKE commands = [ - "cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF -DLLAMA_CURL=ON", - f"cmake --build llama.cpp/build --config Release -j{psutil.cpu_count()*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}", + f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}", + f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}", "cp llama.cpp/build/bin/llama-* llama.cpp", "rm -rf llama.cpp/build", ] diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 99651643a8..0136e3498e 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -25,6 +25,7 @@ import collections import numpy as np import gc import subprocess +import psutil from unsloth_zoo.tokenizer_utils import ( mean_of_trained_tokens, diff --git a/unsloth/trainer.py b/unsloth/trainer.py index c0b2dd03b6..858dcf2cd3 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -14,6 +14,7 @@ import logging import os +import psutil import warnings from dataclasses import dataclass, field from typing import Optional @@ -210,7 +211,7 @@ def _backwards_compatible_trainer(trainer_class, config_class): if "processing_class" in trainer_params and "tokenizer" in kwargs: kwargs["processing_class"] = kwargs.pop("tokenizer") - if ("args" in kwargs) and (Version(trl.__version__) >= Version("0.13.0.dev0")): + if ("args" in kwargs) and (Version(trl) >= Version("0.13.0.dev0")): training_args = kwargs.pop("args", None) # Get parameters that Trainer.__init__ actually expects @@ -411,7 +412,7 @@ def _patch_trl_trainer(): if hasattr(trl, "__UNSLOTH_BACKWARDS_COMPATIBLE__"): return - if Version(trl.__version__) <= Version("0.11.0"): + if Version(trl) <= Version("0.11.0"): return import trl.trainer diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index ccd49dada8..a7620549be 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -32,9 +32,6 @@ from ..utils.packing import ( if HAS_FLASH_ATTENTION: from flash_attn import flash_attn_func, flash_attn_varlen_func HAS_XFORMERS = xformers is not None -BlockDiagonalCausalMask = None -if HAS_XFORMERS: - BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask SDPA_HAS_GQA = "enable_gqa" in (scaled_dot_product_attention.__doc__ or "") FLASH_VARLEN = "flash_varlen" @@ -219,16 +216,10 @@ def run_attention( ) if config.n_groups != 1 and not requires_grad: - if has_block: - out = out.view(bsz, q_len, config.n_kv_heads, config.n_groups, head_dim) - else: - out = out.view(bsz, q_len, config.n_kv_heads, config.n_groups, head_dim) + out = out.view(bsz, q_len, config.n_kv_heads, config.n_groups, head_dim) out = out.reshape(bsz, q_len, n_heads, head_dim) else: - if has_block: - out = out.view(bsz, q_len, n_heads, head_dim) - else: - out = out.view(bsz, q_len, n_heads, head_dim) + out = out.view(bsz, q_len, n_heads, head_dim) return out else: local_mask = context.attention_mask diff --git a/unsloth/utils/hf_hub.py b/unsloth/utils/hf_hub.py index 75df00fbf0..e3960ba0ce 100644 --- a/unsloth/utils/hf_hub.py +++ b/unsloth/utils/hf_hub.py @@ -19,7 +19,9 @@ def formatted_int(value: int) -> str: elif value < MILLION: return f"{float(value) / 1000:,.1f}K" elif value < BILLION: - return f"{float(value) // 1000000:,.1f}M" + return f"{float(value) / 1000000:,.1f}M" + else: + return f"{float(value) / 1000000000:,.1f}B" def get_model_info(