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..cc188674b7 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.11 hooks: - id: ruff args: diff --git a/README.md b/README.md index 43c09381fc..ae1fccfbba 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@
@@ -8,7 +8,7 @@
-
+
### Train gpt-oss, DeepSeek, Gemma, Qwen & Llama 2x faster with 70% less VRAM!
@@ -18,7 +18,7 @@
## ✨ 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 |
|-----------|---------|--------|----------|
@@ -34,9 +34,9 @@ Notebooks are beginner friendly. Read our [guide](https://docs.unsloth.ai/get-st
| **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) & [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,35 @@ 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.
+- New RoPE & MLP **Triton Kernels** & **Padding Free + Packing**: 3x faster training & 30% less VRAM. [Blog](https://unsloth.ai/docs/new/3x-faster-training-packing)
+- **New Mistral**: 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)
+- **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)
+- **gpt-oss RL**: Introducing the fastest possible inference for gpt-oss RL! [Read blog](https://unsloth.ai/docs/models/gpt-oss-how-to-run-and-fine-tune/gpt-oss-reinforcement-learning)
+- **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 [Unsloth 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.
**r/unsloth Reddit** | [Join Reddit community](https://reddit.com/r/unsloth)|
-| 📚 **Documentation & Wiki** | [Read Our Docs](https://docs.unsloth.ai) |
-|
**r/unsloth Reddit** | [Join Reddit community](https://reddit.com/r/unsloth) |
+| 📚 **Documentation & Wiki** | [Read Our Docs](https://unsloth.ai/docs) |
+|
+
+* 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, [BERT](https://unsloth.ai/docs/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://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 +128,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 +143,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 +260,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 +275,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 +302,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 +343,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 +352,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 +420,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 decc0e9f5f..05fb690da5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -51,7 +51,7 @@ 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",
@@ -60,7 +60,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
- "unsloth_zoo>=2025.12.7",
+ "unsloth_zoo>=2026.1.3",
"torchvision",
"unsloth[triton]",
]
@@ -523,7 +523,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
- "unsloth_zoo>=2025.12.7",
+ "unsloth_zoo>=2026.1.3",
"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",
@@ -542,7 +542,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/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..89824a25d1 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,7 @@ from .import_fixes import (
fix_xformers_performance_issue,
fix_vllm_aimv2_issue,
fix_vllm_guided_decoding_params,
+ fix_vllm_pdl_blackwell,
ignore_logger_messages,
patch_ipykernel_hf_xet,
patch_trackio,
@@ -135,6 +139,7 @@ from .import_fixes import (
fix_xformers_performance_issue()
fix_vllm_aimv2_issue()
fix_vllm_guided_decoding_params()
+fix_vllm_pdl_blackwell()
ignore_logger_messages()
patch_ipykernel_hf_xet()
patch_trackio()
@@ -146,6 +151,7 @@ fix_executorch()
del fix_xformers_performance_issue
del fix_vllm_aimv2_issue
del fix_vllm_guided_decoding_params
+del fix_vllm_pdl_blackwell
del ignore_logger_messages
del patch_ipykernel_hf_xet
del patch_trackio
@@ -273,6 +279,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 68038de679..0f924bfdfd 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
diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index efc7a7f4cd..27e5342e20 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'
@@ -316,10 +341,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 +568,128 @@ 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_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/_utils.py b/unsloth/models/_utils.py
index c0ca942995..c52c3d24f7 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.3"
__all__ = [
"SUPPORTS_BFLOAT16",
@@ -75,6 +75,7 @@ __all__ = [
"hf_login",
"is_moe_model",
"get_moe_target_parameters",
+ "make_fast_generate_wrapper",
]
import torch
@@ -1107,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
@@ -1185,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()
@@ -1198,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
@@ -1983,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":
@@ -2199,6 +2217,18 @@ def _prepare_model_for_qat(
from torchao.quantization.granularity import PerGroup, PerAxis
from torchao.quantization.qat import QATConfig
+ # 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":
@@ -2368,7 +2398,6 @@ def hf_login(token: Optional[str] = None) -> Optional[str]:
logger.info(f"Failed to login to huggingface using token with error: {e}")
return token
-
# =============================================
# MoE (Mixture of Experts) Detection and LoRA Utilities
@@ -2446,4 +2475,59 @@ def get_moe_target_parameters(model, target_modules=None) -> Optional[List[str]]
return moe_params
return None
-# =============================================
+
+
+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 deb31ec0dc..038fd855c7 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 (
@@ -2601,6 +2654,7 @@ class FastLlamaModel:
temporary_location = "_unsloth_temporary_saved_buffers",
qat_scheme = None,
target_parameters = None, # For MoE expert layers (nn.Parameter)
+ ensure_weight_tying = False,
**kwargs,
):
if os.environ.get("UNSLOTH_USE_NEW_MODEL", "0") == "1":
@@ -2631,6 +2685,7 @@ class FastLlamaModel:
loftq_config = loftq_config,
temporary_location = temporary_location,
target_parameters = target_parameters,
+ ensure_weight_tying = ensure_weight_tying,
**kwargs,
)
if os.environ.get("UNSLOTH_ENABLE_FULL_FINETUNING", "0") == "1":
@@ -2711,46 +2766,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:
@@ -2781,9 +2806,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":
@@ -2959,6 +2985,7 @@ class FastLlamaModel:
use_rslora = use_rslora,
modules_to_save = modules_to_save,
target_parameters = target_parameters,
+ ensure_weight_tying = ensure_weight_tying,
**kwargs,
)
if not SUPPORTS_LOFTQ:
@@ -3008,6 +3035,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 2c98a27035..ae3b002bd1 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)
@@ -680,12 +738,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
@@ -745,6 +836,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:
@@ -1152,12 +1254,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,
@@ -1203,22 +1311,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/vision.py b/unsloth/models/vision.py
index 3f25392e20..c0f8a1c9a3 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
@@ -150,7 +148,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:
@@ -159,7 +157,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
@@ -532,6 +530,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."
@@ -599,7 +598,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(
@@ -644,7 +644,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
@@ -676,7 +677,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(
@@ -719,9 +720,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:
@@ -938,6 +943,7 @@ class FastBaseModel:
temporary_location = "_unsloth_temporary_saved_buffers",
qat_scheme = None,
target_parameters = None, # For MoE expert layers (nn.Parameter)
+ 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":
@@ -1273,7 +1279,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(