Merge branch 'unslothai:main' into main
This commit is contained in:
electron271
GitHub
commit
3ac0abf7d5
30 changed files with 2928 additions and 612 deletions
|
|
@ -1,6 +1,6 @@
|
|||
repos:
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.14.6
|
||||
rev: v0.14.8
|
||||
hooks:
|
||||
- id: ruff
|
||||
args:
|
||||
|
|
|
|||
12
README.md
12
README.md
|
|
@ -23,14 +23,15 @@ Notebooks are beginner friendly. Read our [guide](https://docs.unsloth.ai/get-st
|
|||
| Model | Free Notebooks | Performance | Memory use |
|
||||
|-----------|---------|--------|----------|
|
||||
| **gpt-oss (20B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-Fine-tuning.ipynb) | 1.5x faster | 70% less |
|
||||
| **Qwen3 (14B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_(14B)-Reasoning-Conversational.ipynb) | 2x faster | 70% less |
|
||||
| **Mistral Ministral 3 (3B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Ministral_3_VL_(3B)_Vision.ipynb) | 1.5x faster | 60% less |
|
||||
| **gpt-oss (20B): GRPO** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-GRPO.ipynb) | 2x faster | 80% less |
|
||||
| **Qwen3: Advanced GRPO** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_(4B)-GRPO.ipynb) | 2x faster | 50% less |
|
||||
| **Qwen3-VL (8B): GSPO** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_VL_(8B)-Vision-GRPO.ipynb) | 1.5x faster | 80% less |
|
||||
| **Qwen3-VL (8B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_VL_(8B)-Vision.ipynb) | 2x faster | 50% less |
|
||||
| **Gemma 3 (270M)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma3_(270M).ipynb) | 1.7x faster | 60% less |
|
||||
| **Gemma 3n (4B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma3N_(4B)-Conversational.ipynb) | 1.5x faster | 50% less |
|
||||
| **DeepSeek-OCR (3B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Deepseek_OCR_(3B).ipynb) | 1.5x faster | 30% less |
|
||||
| **Llama 3.1 (8B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-Alpaca.ipynb) | 2x faster | 70% less |
|
||||
| **Llama 3.1 (8B) Alpaca** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-Alpaca.ipynb) | 2x faster | 70% less |
|
||||
| **Llama 3.2 Conversational** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.2_(1B_and_3B)-Conversational.ipynb) | 2x faster | 70% less |
|
||||
| **Orpheus-TTS (3B)** | [▶️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Orpheus_(3B)-TTS.ipynb) | 1.5x faster | 50% less |
|
||||
|
||||
- See all our notebooks for: [Kaggle](https://github.com/unslothai/notebooks?tab=readme-ov-file#-kaggle-notebooks), [GRPO](https://docs.unsloth.ai/get-started/unsloth-notebooks#grpo-reasoning-rl-notebooks), **[TTS](https://docs.unsloth.ai/get-started/unsloth-notebooks#text-to-speech-tts-notebooks)** & [Vision](https://docs.unsloth.ai/get-started/unsloth-notebooks#vision-multimodal-notebooks)
|
||||
|
|
@ -50,17 +51,20 @@ Use our official [Unsloth Docker image](https://hub.docker.com/r/unsloth/unsloth
|
|||
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.
|
||||
|
||||
## 🦥 Unsloth News
|
||||
- New RoPE & MLP **Triton Kernels** & **Auto 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 our vision or RL sodoku notebook. [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)
|
||||
- **Quantization-Aware Training**: We collabed with Pytorch, recovering ~70% accuracy. [Read blog](https://docs.unsloth.ai/new/quantization-aware-training-qat)
|
||||
- **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.
|
||||
|
||||
<details>
|
||||
<summary>Click for more news</summary>
|
||||
|
||||
- **Quantization-Aware Training**: We collabed with Pytorch, recovering ~70% accuracy. [Read blog](https://docs.unsloth.ai/new/quantization-aware-training-qat)
|
||||
- **Memory-efficient RL**: We're introducing even better RL. Our new kernels & algos allows faster RL with 50% less VRAM & 10× more context. [Read blog](https://docs.unsloth.ai/new/memory-efficient-rl)
|
||||
- **Gemma 3n** by Google: [Read Blog](https://docs.unsloth.ai/basics/gemma-3n-how-to-run-and-fine-tune). We [uploaded GGUFs, 4-bit models](https://huggingface.co/collections/unsloth/gemma-3n-685d3874830e49e1c93f9339).
|
||||
- **[Text-to-Speech (TTS)](https://docs.unsloth.ai/basics/text-to-speech-tts-fine-tuning)** is now supported, including `sesame/csm-1b` and STT `openai/whisper-large-v3`.
|
||||
|
|
|
|||
|
|
@ -55,12 +55,12 @@ huggingfacenotorch = [
|
|||
"huggingface_hub>=0.34.0",
|
||||
"hf_transfer",
|
||||
"diffusers",
|
||||
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,<=4.57.2",
|
||||
"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",
|
||||
"trl>=0.18.2,!=0.19.0,<=0.24.0",
|
||||
]
|
||||
huggingface = [
|
||||
"unsloth[huggingfacenotorch]",
|
||||
"unsloth_zoo>=2025.11.5",
|
||||
"unsloth_zoo>=2025.12.4",
|
||||
"torchvision",
|
||||
"unsloth[triton]",
|
||||
]
|
||||
|
|
@ -285,10 +285,16 @@ cu130onlytorch290 = [
|
|||
"xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')",
|
||||
]
|
||||
cu126onlytorch291 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.33.post2-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
|
||||
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.33.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')",
|
||||
]
|
||||
cu128onlytorch291 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.33.post2-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
|
||||
"xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.33.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')",
|
||||
]
|
||||
cu130onlytorch291 = [
|
||||
"xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post2-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
|
||||
"xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post2-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')",
|
||||
]
|
||||
cu118 = [
|
||||
"unsloth[huggingface]",
|
||||
|
|
@ -517,10 +523,10 @@ colab-ampere-torch220 = [
|
|||
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
|
||||
]
|
||||
colab-new = [
|
||||
"unsloth_zoo>=2025.11.5",
|
||||
"unsloth_zoo>=2025.12.4",
|
||||
"packaging",
|
||||
"tyro",
|
||||
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,<=4.57.2",
|
||||
"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",
|
||||
"datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0",
|
||||
"sentencepiece>=0.2.0",
|
||||
"tqdm",
|
||||
|
|
|
|||
272
tests/utils/test_attention_masks.py
Normal file
272
tests/utils/test_attention_masks.py
Normal file
|
|
@ -0,0 +1,272 @@
|
|||
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Lesser General Public License as published by
|
||||
# the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Lesser General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""Unit tests for packed-attention mask helpers with sliding-window logic."""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
from unsloth.utils import attention_dispatch
|
||||
from unsloth.utils import packing as packing_utils
|
||||
|
||||
|
||||
def _make_seq_info(lengths):
|
||||
lengths = torch.tensor(lengths, dtype = torch.int32)
|
||||
cu = torch.cat(
|
||||
[
|
||||
torch.zeros(1, dtype = torch.int32),
|
||||
torch.cumsum(lengths, dim = 0, dtype = torch.int32),
|
||||
]
|
||||
)
|
||||
max_len = int(lengths.max().item())
|
||||
return lengths, cu, max_len
|
||||
|
||||
|
||||
def test_sdpa_packed_attention_mask_sliding_window():
|
||||
seq_info = _make_seq_info([5, 3])
|
||||
mask = packing_utils.build_sdpa_packed_attention_mask(
|
||||
seq_info,
|
||||
dtype = torch.float32,
|
||||
device = torch.device("cpu"),
|
||||
sliding_window = 3,
|
||||
)
|
||||
|
||||
assert mask.shape == (1, 1, 8, 8)
|
||||
|
||||
block_first = mask[0, 0, :5, :5]
|
||||
upper = torch.triu(torch.ones_like(block_first), diagonal = 1).bool()
|
||||
assert torch.all(block_first[upper] == float("-inf"))
|
||||
assert block_first[3, 0].item() == float("-inf")
|
||||
assert block_first[4, 1].item() == float("-inf")
|
||||
assert block_first[4, 2].item() > -math.inf
|
||||
assert mask[0, 0, 0, 6].item() == float("-inf")
|
||||
|
||||
|
||||
def test_xformers_block_mask_sliding_window(monkeypatch):
|
||||
class _FakeMask:
|
||||
def __init__(self, lengths, window = None):
|
||||
self.lengths = lengths
|
||||
self.window = window
|
||||
|
||||
@classmethod
|
||||
def from_seqlens(cls, lengths):
|
||||
return cls(tuple(lengths))
|
||||
|
||||
def make_local_attention(self, window_size):
|
||||
return _FakeMask(self.lengths, window = window_size)
|
||||
|
||||
monkeypatch.setattr(packing_utils, "_XFormersBlockMask", _FakeMask, raising = False)
|
||||
|
||||
seq_info = _make_seq_info([4, 4])
|
||||
mask = packing_utils.build_xformers_block_causal_mask(
|
||||
seq_info,
|
||||
sliding_window = 2,
|
||||
)
|
||||
|
||||
assert isinstance(mask, _FakeMask)
|
||||
assert mask.window == 2
|
||||
|
||||
|
||||
def test_run_attention_sdpa_passes_sliding_window(monkeypatch):
|
||||
seq_info = _make_seq_info([3, 2])
|
||||
sliding_window = 2
|
||||
|
||||
original_builder = attention_dispatch.build_sdpa_packed_attention_mask
|
||||
captured = {}
|
||||
|
||||
def _capture_builder(seq_info_arg, *, dtype, device, sliding_window = None):
|
||||
captured["window"] = sliding_window
|
||||
return original_builder(
|
||||
seq_info_arg,
|
||||
dtype = dtype,
|
||||
device = device,
|
||||
sliding_window = sliding_window,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
attention_dispatch,
|
||||
"build_sdpa_packed_attention_mask",
|
||||
_capture_builder,
|
||||
)
|
||||
|
||||
def _fake_sdpa(Q, K, V, **kwargs):
|
||||
captured["mask"] = kwargs.get("attn_mask")
|
||||
return torch.zeros_like(Q)
|
||||
|
||||
monkeypatch.setattr(attention_dispatch, "scaled_dot_product_attention", _fake_sdpa)
|
||||
|
||||
config = attention_dispatch.AttentionConfig(
|
||||
backend = attention_dispatch.SDPA,
|
||||
n_kv_heads = 1,
|
||||
n_groups = 1,
|
||||
)
|
||||
|
||||
context = attention_dispatch.AttentionContext(
|
||||
bsz = 1,
|
||||
q_len = 5,
|
||||
kv_seq_len = 5,
|
||||
n_heads = 1,
|
||||
head_dim = 1,
|
||||
requires_grad = False,
|
||||
seq_info = seq_info,
|
||||
attention_mask = None,
|
||||
causal_mask = None,
|
||||
sliding_window = sliding_window,
|
||||
)
|
||||
|
||||
Q = torch.zeros(1, 1, 5, 1)
|
||||
K = torch.zeros_like(Q)
|
||||
V = torch.zeros_like(Q)
|
||||
|
||||
attention_dispatch.run_attention(
|
||||
config = config,
|
||||
context = context,
|
||||
Q = Q,
|
||||
K = K,
|
||||
V = V,
|
||||
)
|
||||
|
||||
assert captured["window"] == sliding_window
|
||||
mask = captured["mask"]
|
||||
assert mask is not None and mask.shape == (1, 1, 5, 5)
|
||||
assert mask[0, 0, 4, 1].item() == float("-inf")
|
||||
|
||||
|
||||
def test_run_attention_xformers_passes_sliding_window(monkeypatch):
|
||||
seq_info = _make_seq_info([4])
|
||||
sliding_window = 3
|
||||
|
||||
class _FakeBias:
|
||||
pass
|
||||
|
||||
captured = {}
|
||||
|
||||
def _fake_builder(seq_info_arg, *, sliding_window = None, base_mask = None):
|
||||
captured["window"] = sliding_window
|
||||
captured["base"] = base_mask
|
||||
return _FakeBias()
|
||||
|
||||
def _fake_attention(Q, K, V, attn_bias = None, **_):
|
||||
captured["bias"] = attn_bias
|
||||
return torch.zeros_like(Q)
|
||||
|
||||
monkeypatch.setattr(
|
||||
attention_dispatch, "build_xformers_block_causal_mask", _fake_builder
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
attention_dispatch, "xformers_attention", _fake_attention, raising = False
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
attention_dispatch, "XFORMERS_BLOCK_DIAG_CLS", _FakeBias, raising = False
|
||||
)
|
||||
|
||||
config = attention_dispatch.AttentionConfig(
|
||||
backend = attention_dispatch.XFORMERS,
|
||||
n_kv_heads = 1,
|
||||
n_groups = 1,
|
||||
)
|
||||
|
||||
context = attention_dispatch.AttentionContext(
|
||||
bsz = 1,
|
||||
q_len = 4,
|
||||
kv_seq_len = 4,
|
||||
n_heads = 1,
|
||||
head_dim = 1,
|
||||
requires_grad = False,
|
||||
seq_info = seq_info,
|
||||
attention_mask = None,
|
||||
causal_mask = None,
|
||||
sliding_window = sliding_window,
|
||||
)
|
||||
|
||||
Q = torch.zeros(1, 1, 4, 1)
|
||||
K = torch.zeros_like(Q)
|
||||
V = torch.zeros_like(Q)
|
||||
|
||||
attention_dispatch.run_attention(
|
||||
config = config,
|
||||
context = context,
|
||||
Q = Q,
|
||||
K = K,
|
||||
V = V,
|
||||
)
|
||||
|
||||
assert captured["window"] == sliding_window
|
||||
assert isinstance(captured["bias"], _FakeBias)
|
||||
|
||||
|
||||
def test_run_attention_flash_varlen_receives_window_and_softcap(monkeypatch):
|
||||
seq_info = _make_seq_info([4])
|
||||
sliding_window = 3
|
||||
softcap = 0.5
|
||||
window_tuple = (sliding_window, sliding_window)
|
||||
|
||||
captured = {}
|
||||
|
||||
def _fake_flash_varlen(Q, K, V, cu_q, cu_k, max_q, max_k, **kwargs):
|
||||
captured["kwargs"] = kwargs
|
||||
return torch.zeros_like(Q)
|
||||
|
||||
monkeypatch.setattr(
|
||||
attention_dispatch,
|
||||
"flash_attn_varlen_func",
|
||||
_fake_flash_varlen,
|
||||
)
|
||||
monkeypatch.setattr(attention_dispatch, "HAS_FLASH_ATTENTION", True)
|
||||
|
||||
config = attention_dispatch.AttentionConfig(
|
||||
backend = attention_dispatch.FLASH_VARLEN,
|
||||
n_kv_heads = 1,
|
||||
n_groups = 1,
|
||||
flash_varlen_kwargs = {
|
||||
"dropout_p": 0.0,
|
||||
"softmax_scale": 1.0,
|
||||
"causal": True,
|
||||
"softcap": softcap,
|
||||
"window_size": window_tuple,
|
||||
},
|
||||
)
|
||||
|
||||
context = attention_dispatch.AttentionContext(
|
||||
bsz = 1,
|
||||
q_len = 4,
|
||||
kv_seq_len = 4,
|
||||
n_heads = 1,
|
||||
head_dim = 2,
|
||||
requires_grad = False,
|
||||
seq_info = seq_info,
|
||||
attention_mask = None,
|
||||
causal_mask = None,
|
||||
sliding_window = sliding_window,
|
||||
)
|
||||
|
||||
Q = torch.zeros(1, 1, 4, 2)
|
||||
K = torch.zeros_like(Q)
|
||||
V = torch.zeros_like(Q)
|
||||
|
||||
attention_dispatch.run_attention(
|
||||
config = config,
|
||||
context = context,
|
||||
Q = Q,
|
||||
K = K,
|
||||
V = V,
|
||||
)
|
||||
|
||||
assert captured["kwargs"]["softcap"] == softcap
|
||||
assert captured["kwargs"]["window_size"] == window_tuple
|
||||
|
||||
|
||||
"""Unit tests for packed-attention mask helpers with sliding-window logic."""
|
||||
391
tests/utils/test_packing.py
Normal file
391
tests/utils/test_packing.py
Normal file
|
|
@ -0,0 +1,391 @@
|
|||
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Lesser General Public License as published by
|
||||
# the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Lesser General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
from unsloth import FastLanguageModel
|
||||
from unsloth.utils import attention_dispatch as attention_dispatch_utils
|
||||
from unsloth.utils.packing import (
|
||||
configure_padding_free,
|
||||
configure_sample_packing,
|
||||
enable_padding_free_metadata,
|
||||
enable_sample_packing,
|
||||
mask_packed_sequence_boundaries,
|
||||
)
|
||||
|
||||
from contextlib import ExitStack
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from datasets import Dataset
|
||||
from trl import SFTConfig, SFTTrainer
|
||||
from trl.trainer.sft_trainer import DataCollatorForLanguageModeling
|
||||
|
||||
|
||||
def _build_packed_training_setup(tmp_path, device):
|
||||
dtype = None
|
||||
if device.type == "cuda":
|
||||
if torch.cuda.is_bf16_supported():
|
||||
dtype = torch.bfloat16
|
||||
else:
|
||||
dtype = torch.float16
|
||||
|
||||
try:
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name = "hf-internal-testing/tiny-random-LlamaForCausalLM",
|
||||
max_seq_length = 64,
|
||||
load_in_4bit = False,
|
||||
dtype = dtype,
|
||||
)
|
||||
except OSError as exc: # pragma: no cover - offline CI
|
||||
pytest.skip(f"Requires access to tiny llama checkpoint: {exc}")
|
||||
|
||||
model.to(device)
|
||||
|
||||
dataset = Dataset.from_dict(
|
||||
{
|
||||
"text": [
|
||||
"Hello world!",
|
||||
"Short sample.",
|
||||
"This is a slightly longer packed example to test batching.",
|
||||
"Another response to include in the batch.",
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
training_args = SFTConfig(
|
||||
per_device_train_batch_size = 1,
|
||||
per_device_eval_batch_size = 1,
|
||||
gradient_accumulation_steps = 1,
|
||||
dataset_text_field = "text",
|
||||
max_length = 64,
|
||||
logging_steps = 1,
|
||||
max_steps = 1,
|
||||
fp16 = device.type == "cuda" and not torch.cuda.is_bf16_supported(),
|
||||
bf16 = device.type == "cuda" and torch.cuda.is_bf16_supported(),
|
||||
dataset_num_proc = 1,
|
||||
output_dir = str(tmp_path),
|
||||
packing = True,
|
||||
)
|
||||
|
||||
trainer = SFTTrainer(
|
||||
model = model,
|
||||
processing_class = tokenizer,
|
||||
train_dataset = dataset,
|
||||
args = training_args,
|
||||
)
|
||||
|
||||
enable_sample_packing(model, trainer)
|
||||
|
||||
dataloader = trainer.get_train_dataloader()
|
||||
batch = next(iter(dataloader))
|
||||
|
||||
model_device = next(model.parameters()).device
|
||||
|
||||
for key, value in list(batch.items()):
|
||||
if torch.is_tensor(value):
|
||||
batch[key] = value.to(model_device)
|
||||
|
||||
from unsloth.models import llama as llama_mod
|
||||
|
||||
return model, batch, trainer, llama_mod
|
||||
|
||||
|
||||
def _trim_batch_to_total_tokens(data, total_tokens):
|
||||
def _trim_tensor(t: torch.Tensor):
|
||||
if t.ndim >= 2 and t.size(1) > total_tokens:
|
||||
return t[:, :total_tokens].contiguous()
|
||||
return t
|
||||
|
||||
trimmed = {}
|
||||
for key, value in data.items():
|
||||
if torch.is_tensor(value):
|
||||
trimmed[key] = _trim_tensor(value)
|
||||
else:
|
||||
trimmed[key] = value
|
||||
return trimmed
|
||||
|
||||
|
||||
def test_mask_packed_sequence_boundaries_marks_single_row():
|
||||
shift_labels = torch.arange(6, dtype = torch.long).view(1, 6)
|
||||
changed = mask_packed_sequence_boundaries(
|
||||
shift_labels,
|
||||
torch.tensor([2, 1, 3], dtype = torch.int32),
|
||||
)
|
||||
assert changed is True
|
||||
flat = shift_labels.view(-1)
|
||||
assert flat[1].item() == -100
|
||||
assert flat[2].item() == -100
|
||||
assert flat[5].item() == -100
|
||||
assert flat[0].item() != -100
|
||||
|
||||
|
||||
def test_mask_packed_sequence_boundaries_across_multiple_rows():
|
||||
shift_labels = torch.arange(10, dtype = torch.long).view(2, 5)
|
||||
lengths = torch.tensor([3, 2, 4, 1], dtype = torch.int32)
|
||||
changed = mask_packed_sequence_boundaries(shift_labels, lengths)
|
||||
assert changed is True
|
||||
flat = shift_labels.view(-1)
|
||||
for idx in (2, 4, 8, 9):
|
||||
assert flat[idx].item() == -100
|
||||
assert torch.any(flat != -100)
|
||||
|
||||
|
||||
def test_configure_sample_packing():
|
||||
config = SimpleNamespace()
|
||||
configure_sample_packing(config)
|
||||
|
||||
assert config.packing is True
|
||||
assert config.padding_free is True
|
||||
assert config.remove_unused_columns is False
|
||||
|
||||
|
||||
def test_configure_padding_free():
|
||||
config = SimpleNamespace(remove_unused_columns = True)
|
||||
configure_padding_free(config)
|
||||
|
||||
assert config.padding_free is True
|
||||
assert config.remove_unused_columns is False
|
||||
|
||||
|
||||
class _DummyChild(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.max_seq_length = 8
|
||||
|
||||
|
||||
class _DummyModel(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.max_seq_length = 16
|
||||
self.child = _DummyChild()
|
||||
self.config = SimpleNamespace(_attn_implementation = "sdpa")
|
||||
self.generation_config = SimpleNamespace(attn_implementation = "sdpa")
|
||||
|
||||
|
||||
class _DummyTrainer:
|
||||
def __init__(self):
|
||||
self.args = SimpleNamespace(remove_unused_columns = True)
|
||||
self.data_collator = DataCollatorForLanguageModeling(
|
||||
pad_token_id = 0,
|
||||
completion_only_loss = False,
|
||||
padding_free = True,
|
||||
return_position_ids = False,
|
||||
return_tensors = "pt",
|
||||
)
|
||||
|
||||
|
||||
class _PaddingFreeCollator:
|
||||
def __init__(self):
|
||||
self.padding_free = True
|
||||
self.return_position_ids = False
|
||||
self.calls = 0
|
||||
|
||||
def torch_call(self, examples):
|
||||
self.calls += 1
|
||||
return {
|
||||
"input_ids": torch.tensor([[0]], dtype = torch.long),
|
||||
"examples_seen": self.calls,
|
||||
}
|
||||
|
||||
|
||||
def test_enable_sample_packing():
|
||||
model = _DummyModel()
|
||||
trainer = _DummyTrainer()
|
||||
|
||||
enable_sample_packing(model, trainer)
|
||||
|
||||
# model hierarchy should now allow packed overlength inputs
|
||||
assert getattr(model, "_unsloth_allow_packed_overlength") is True
|
||||
assert getattr(model.child, "_unsloth_allow_packed_overlength") is True
|
||||
|
||||
collator = trainer.data_collator
|
||||
assert collator.return_position_ids is True
|
||||
assert getattr(collator, "_unsloth_packing_wrapped") is True
|
||||
|
||||
examples = [
|
||||
{
|
||||
"input_ids": [0, 1, 2],
|
||||
"labels": [0, 1, 2],
|
||||
"seq_lengths": [2, 1],
|
||||
},
|
||||
{
|
||||
"input_ids": [3, 4, 5],
|
||||
"labels": [3, 4, 5],
|
||||
"seq_lengths": [3],
|
||||
},
|
||||
]
|
||||
batch = collator.torch_call(examples)
|
||||
|
||||
# packed lengths are aggregated into a single tensor
|
||||
assert "packed_seq_lengths" in batch
|
||||
assert torch.equal(
|
||||
batch["packed_seq_lengths"],
|
||||
torch.tensor([2, 1, 3], dtype = torch.int32),
|
||||
)
|
||||
|
||||
assert batch["input_ids"].shape == (1, 6)
|
||||
expected_positions = torch.tensor([0, 1, 0, 0, 1, 2], dtype = torch.long)
|
||||
assert torch.equal(batch["position_ids"].view(-1)[:6], expected_positions)
|
||||
|
||||
|
||||
def test_enable_sample_packing_trl_collator(tmp_path):
|
||||
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
|
||||
model, _, trainer, _ = _build_packed_training_setup(tmp_path, device)
|
||||
|
||||
enable_sample_packing(model, trainer)
|
||||
|
||||
examples = [
|
||||
{
|
||||
"input_ids": [0, 1, 2],
|
||||
"labels": [0, 1, 2],
|
||||
"seq_lengths": [2, 1],
|
||||
},
|
||||
{
|
||||
"input_ids": [3, 4, 5],
|
||||
"labels": [3, 4, 5],
|
||||
"seq_lengths": [3],
|
||||
},
|
||||
]
|
||||
|
||||
batch = trainer.data_collator.torch_call(examples)
|
||||
|
||||
assert batch["input_ids"].shape == (1, 6)
|
||||
assert torch.equal(
|
||||
batch["packed_seq_lengths"],
|
||||
torch.tensor([2, 1, 3], dtype = torch.int32),
|
||||
)
|
||||
|
||||
expected_positions = torch.tensor([0, 1, 0, 0, 1, 2], dtype = torch.long)
|
||||
assert torch.equal(batch["position_ids"].view(-1)[:6], expected_positions)
|
||||
|
||||
if hasattr(trainer, "accelerator"):
|
||||
trainer.accelerator.free_memory()
|
||||
|
||||
|
||||
def test_enable_padding_free_metadata():
|
||||
model = _DummyModel()
|
||||
trainer = SimpleNamespace(
|
||||
args = SimpleNamespace(remove_unused_columns = True),
|
||||
data_collator = _PaddingFreeCollator(),
|
||||
)
|
||||
|
||||
enable_padding_free_metadata(model, trainer)
|
||||
|
||||
assert getattr(model, "_unsloth_allow_packed_overlength") is True
|
||||
assert getattr(model.child, "_unsloth_allow_packed_overlength") is True
|
||||
|
||||
collator = trainer.data_collator
|
||||
assert collator.return_position_ids is True
|
||||
assert getattr(collator, "_unsloth_padding_free_lengths_wrapped") is True
|
||||
|
||||
examples = [
|
||||
{"input_ids": [0, 1, 2]},
|
||||
{"input_ids": [3, 4]},
|
||||
]
|
||||
batch = collator.torch_call(examples)
|
||||
assert torch.equal(
|
||||
batch["packed_seq_lengths"],
|
||||
torch.tensor([3, 2], dtype = torch.int32),
|
||||
)
|
||||
assert trainer.args.remove_unused_columns is False
|
||||
|
||||
|
||||
def test_packing_sdpa(tmp_path):
|
||||
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
|
||||
model, batch, trainer, llama_mod = _build_packed_training_setup(tmp_path, device)
|
||||
|
||||
assert "packed_seq_lengths" in batch
|
||||
assert "attention_mask" not in batch
|
||||
assert batch["packed_seq_lengths"].dtype == torch.int32
|
||||
|
||||
total_tokens = batch["input_ids"].size(-1)
|
||||
assert int(batch["packed_seq_lengths"].sum().item()) == total_tokens
|
||||
|
||||
packed_tokens = int(batch["packed_seq_lengths"].sum().item())
|
||||
assert "position_ids" in batch
|
||||
flat_positions = batch["position_ids"].reshape(-1)[:packed_tokens]
|
||||
expected_positions = torch.cat(
|
||||
[
|
||||
torch.arange(length, dtype = torch.long)
|
||||
for length in batch["packed_seq_lengths"].tolist()
|
||||
]
|
||||
)
|
||||
assert torch.equal(flat_positions.cpu(), expected_positions)
|
||||
inputs = _trim_batch_to_total_tokens(batch, packed_tokens)
|
||||
|
||||
seq_info = llama_mod.get_packed_info_from_kwargs(
|
||||
{"packed_seq_lengths": batch["packed_seq_lengths"]},
|
||||
inputs["input_ids"].device,
|
||||
)
|
||||
assert seq_info is not None
|
||||
|
||||
original_mask = attention_dispatch_utils.build_sdpa_packed_attention_mask
|
||||
mask_calls = []
|
||||
captured_loss_labels = {}
|
||||
|
||||
def _capture_mask(seq_info, dtype, device, *, sliding_window = None):
|
||||
mask_calls.append(tuple(seq_info[0].tolist()))
|
||||
return original_mask(
|
||||
seq_info,
|
||||
dtype = dtype,
|
||||
device = device,
|
||||
sliding_window = sliding_window,
|
||||
)
|
||||
|
||||
def _capture_loss(*, logits, labels, **loss_kwargs):
|
||||
captured_loss_labels["labels"] = labels.detach().to("cpu")
|
||||
return torch.zeros((), device = logits.device, dtype = logits.dtype)
|
||||
|
||||
with ExitStack() as stack:
|
||||
stack.enter_context(
|
||||
patch.object(attention_dispatch_utils, "HAS_FLASH_ATTENTION", False)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch.object(attention_dispatch_utils, "HAS_XFORMERS", False)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch.object(
|
||||
attention_dispatch_utils,
|
||||
"build_sdpa_packed_attention_mask",
|
||||
side_effect = _capture_mask,
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch.object(
|
||||
llama_mod,
|
||||
"fast_cross_entropy_loss",
|
||||
side_effect = _capture_loss,
|
||||
)
|
||||
)
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
|
||||
assert mask_calls, "SDPA packed mask was not constructed"
|
||||
assert outputs.loss is not None
|
||||
assert "labels" in captured_loss_labels
|
||||
flat_loss_labels = captured_loss_labels["labels"].reshape(-1)
|
||||
boundaries = (
|
||||
torch.cumsum(
|
||||
batch["packed_seq_lengths"].to(device = "cpu", dtype = torch.long), dim = 0
|
||||
)
|
||||
- 1
|
||||
)
|
||||
for idx in boundaries.tolist():
|
||||
assert flat_loss_labels[idx].item() == -100
|
||||
assert torch.any(flat_loss_labels != -100)
|
||||
|
||||
if hasattr(trainer, "accelerator"):
|
||||
trainer.accelerator.free_memory()
|
||||
|
|
@ -34,23 +34,24 @@ import os
|
|||
|
||||
|
||||
def run(args):
|
||||
import torch
|
||||
from unsloth import FastLanguageModel
|
||||
from datasets import load_dataset
|
||||
from transformers.utils import strtobool
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
from transformers import TrainingArguments
|
||||
from unsloth import is_bfloat16_supported
|
||||
from unsloth.models.loader_utils import prepare_device_map
|
||||
import logging
|
||||
|
||||
logging.getLogger("hf-to-gguf").setLevel(logging.WARNING)
|
||||
|
||||
# Load model and tokenizer
|
||||
device_map, distributed = prepare_device_map()
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name = args.model_name,
|
||||
max_seq_length = args.max_seq_length,
|
||||
dtype = args.dtype,
|
||||
load_in_4bit = args.load_in_4bit,
|
||||
device_map = device_map,
|
||||
)
|
||||
|
||||
# Configure PEFT model
|
||||
|
|
@ -112,6 +113,7 @@ def run(args):
|
|||
# Configure training arguments
|
||||
training_args = SFTConfig(
|
||||
per_device_train_batch_size = args.per_device_train_batch_size,
|
||||
per_device_eval_batch_size = args.per_device_eval_batch_size,
|
||||
gradient_accumulation_steps = args.gradient_accumulation_steps,
|
||||
warmup_steps = args.warmup_steps,
|
||||
max_steps = args.max_steps,
|
||||
|
|
@ -127,7 +129,8 @@ def run(args):
|
|||
report_to = args.report_to,
|
||||
max_length = args.max_seq_length,
|
||||
dataset_num_proc = 2,
|
||||
packing = False,
|
||||
ddp_find_unused_parameters = False if distributed else None,
|
||||
packing = args.packing,
|
||||
)
|
||||
|
||||
# Initialize trainer
|
||||
|
|
@ -138,8 +141,7 @@ def run(args):
|
|||
args = training_args,
|
||||
)
|
||||
|
||||
# Train model
|
||||
trainer_stats = trainer.train()
|
||||
trainer.train()
|
||||
|
||||
# Save model
|
||||
if args.save_model:
|
||||
|
|
@ -164,13 +166,15 @@ def run(args):
|
|||
else:
|
||||
print(f"Saving model with quantization method: {args.quantization}")
|
||||
model.save_pretrained_gguf(
|
||||
args.save_path, tokenizer, quantization_method = args.quantization
|
||||
args.save_path,
|
||||
tokenizer,
|
||||
quantization_method = args.quantization,
|
||||
)
|
||||
if args.push_model:
|
||||
model.push_to_hub_gguf(
|
||||
hub_path = args.hub_path,
|
||||
hub_token = args.hub_token,
|
||||
quantization_method = quantization_method,
|
||||
quantization_method = args.quantization,
|
||||
)
|
||||
else:
|
||||
model.save_pretrained_merged(args.save_path, tokenizer, args.save_method)
|
||||
|
|
@ -181,7 +185,6 @@ def run(args):
|
|||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Define argument parser
|
||||
parser = argparse.ArgumentParser(
|
||||
description = "🦥 Fine-tune your llm faster using unsloth!"
|
||||
)
|
||||
|
|
@ -218,7 +221,8 @@ if __name__ == "__main__":
|
|||
)
|
||||
|
||||
lora_group = parser.add_argument_group(
|
||||
"🧠 LoRA Options", "These options are used to configure the LoRA model."
|
||||
"🧠 LoRA Options",
|
||||
"These options are used to configure the LoRA model.",
|
||||
)
|
||||
lora_group.add_argument(
|
||||
"--r",
|
||||
|
|
@ -239,7 +243,10 @@ if __name__ == "__main__":
|
|||
help = "LoRA dropout rate, default is 0.0 which is optimized.",
|
||||
)
|
||||
lora_group.add_argument(
|
||||
"--bias", type = str, default = "none", help = "Bias setting for LoRA"
|
||||
"--bias",
|
||||
type = str,
|
||||
default = "none",
|
||||
help = "Bias setting for LoRA",
|
||||
)
|
||||
lora_group.add_argument(
|
||||
"--use_gradient_checkpointing",
|
||||
|
|
@ -254,10 +261,15 @@ if __name__ == "__main__":
|
|||
help = "Random state for reproducibility, default is 3407.",
|
||||
)
|
||||
lora_group.add_argument(
|
||||
"--use_rslora", action = "store_true", help = "Use rank stabilized LoRA"
|
||||
"--use_rslora",
|
||||
action = "store_true",
|
||||
help = "Use rank stabilized LoRA",
|
||||
)
|
||||
lora_group.add_argument(
|
||||
"--loftq_config", type = str, default = None, help = "Configuration for LoftQ"
|
||||
"--loftq_config",
|
||||
type = str,
|
||||
default = None,
|
||||
help = "Configuration for LoftQ",
|
||||
)
|
||||
|
||||
training_group = parser.add_argument_group("🎓 Training Options")
|
||||
|
|
@ -267,6 +279,12 @@ if __name__ == "__main__":
|
|||
default = 2,
|
||||
help = "Batch size per device during training, default is 2.",
|
||||
)
|
||||
training_group.add_argument(
|
||||
"--per_device_eval_batch_size",
|
||||
type = int,
|
||||
default = 4,
|
||||
help = "Batch size per device during evaluation, default is 4.",
|
||||
)
|
||||
training_group.add_argument(
|
||||
"--gradient_accumulation_steps",
|
||||
type = int,
|
||||
|
|
@ -280,7 +298,10 @@ if __name__ == "__main__":
|
|||
help = "Number of warmup steps, default is 5.",
|
||||
)
|
||||
training_group.add_argument(
|
||||
"--max_steps", type = int, default = 400, help = "Maximum number of training steps."
|
||||
"--max_steps",
|
||||
type = int,
|
||||
default = 400,
|
||||
help = "Maximum number of training steps.",
|
||||
)
|
||||
training_group.add_argument(
|
||||
"--learning_rate",
|
||||
|
|
@ -289,7 +310,10 @@ if __name__ == "__main__":
|
|||
help = "Learning rate, default is 2e-4.",
|
||||
)
|
||||
training_group.add_argument(
|
||||
"--optim", type = str, default = "adamw_8bit", help = "Optimizer type."
|
||||
"--optim",
|
||||
type = str,
|
||||
default = "adamw_8bit",
|
||||
help = "Optimizer type.",
|
||||
)
|
||||
training_group.add_argument(
|
||||
"--weight_decay",
|
||||
|
|
@ -309,8 +333,12 @@ if __name__ == "__main__":
|
|||
default = 3407,
|
||||
help = "Seed for reproducibility, default is 3407.",
|
||||
)
|
||||
training_group.add_argument(
|
||||
"--packing",
|
||||
action = "store_true",
|
||||
help = "Enable padding-free sample packing via TRL's bin packer.",
|
||||
)
|
||||
|
||||
# Report/Logging arguments
|
||||
report_group = parser.add_argument_group("📊 Report Options")
|
||||
report_group.add_argument(
|
||||
"--report_to",
|
||||
|
|
@ -331,19 +359,31 @@ if __name__ == "__main__":
|
|||
"all",
|
||||
"none",
|
||||
],
|
||||
help = "The list of integrations to report the results and logs to. Supported platforms are: \n\t\t 'azure_ml', 'clearml', 'codecarbon', 'comet_ml', 'dagshub', 'dvclive', 'flyte', 'mlflow', 'neptune', 'tensorboard', and 'wandb'. Use 'all' to report to all integrations installed, 'none' for no integrations.",
|
||||
help = (
|
||||
"The list of integrations to report the results and logs to. Supported platforms are:\n\t\t "
|
||||
"'azure_ml', 'clearml', 'codecarbon', 'comet_ml', 'dagshub', 'dvclive', 'flyte', "
|
||||
"'mlflow', 'neptune', 'tensorboard', and 'wandb'. Use 'all' to report to all integrations "
|
||||
"installed, 'none' for no integrations."
|
||||
),
|
||||
)
|
||||
report_group.add_argument(
|
||||
"--logging_steps", type = int, default = 1, help = "Logging steps, default is 1"
|
||||
"--logging_steps",
|
||||
type = int,
|
||||
default = 1,
|
||||
help = "Logging steps, default is 1",
|
||||
)
|
||||
|
||||
# Saving and pushing arguments
|
||||
save_group = parser.add_argument_group("💾 Save Model Options")
|
||||
save_group.add_argument(
|
||||
"--output_dir", type = str, default = "outputs", help = "Output directory"
|
||||
"--output_dir",
|
||||
type = str,
|
||||
default = "outputs",
|
||||
help = "Output directory",
|
||||
)
|
||||
save_group.add_argument(
|
||||
"--save_model", action = "store_true", help = "Save the model after training"
|
||||
"--save_model",
|
||||
action = "store_true",
|
||||
help = "Save the model after training",
|
||||
)
|
||||
save_group.add_argument(
|
||||
"--save_method",
|
||||
|
|
@ -358,14 +398,20 @@ if __name__ == "__main__":
|
|||
help = "Convert the model to GGUF after training",
|
||||
)
|
||||
save_group.add_argument(
|
||||
"--save_path", type = str, default = "model", help = "Path to save the model"
|
||||
"--save_path",
|
||||
type = str,
|
||||
default = "model",
|
||||
help = "Path to save the model",
|
||||
)
|
||||
save_group.add_argument(
|
||||
"--quantization",
|
||||
type = str,
|
||||
default = "q8_0",
|
||||
nargs = "+",
|
||||
help = "Quantization method for saving the model. common values ('f16', 'q4_k_m', 'q8_0'), Check our wiki for all quantization methods https://github.com/unslothai/unsloth/wiki#saving-to-gguf ",
|
||||
help = (
|
||||
"Quantization method for saving the model. common values ('f16', 'q4_k_m', 'q8_0'), "
|
||||
"Check our wiki for all quantization methods https://github.com/unslothai/unsloth/wiki#saving-to-gguf"
|
||||
),
|
||||
)
|
||||
|
||||
push_group = parser.add_argument_group("🚀 Push Model Options")
|
||||
|
|
@ -386,7 +432,9 @@ if __name__ == "__main__":
|
|||
help = "Path on Hugging Face hub to push the model",
|
||||
)
|
||||
push_group.add_argument(
|
||||
"--hub_token", type = str, help = "Token for pushing the model to Hugging Face hub"
|
||||
"--hub_token",
|
||||
type = str,
|
||||
help = "Token for pushing the model to Hugging Face hub",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
|
|
|||
|
|
@ -17,16 +17,27 @@ from packaging.version import Version
|
|||
import os, re, subprocess, inspect, functools
|
||||
import numpy as np
|
||||
|
||||
# Fix some issues before importing other packages
|
||||
from .import_fixes import fix_message_factory_issue
|
||||
|
||||
fix_message_factory_issue()
|
||||
del fix_message_factory_issue
|
||||
# Log Unsloth is being used
|
||||
os.environ["UNSLOTH_IS_PRESENT"] = "1"
|
||||
|
||||
# Check if modules that need patching are already imported
|
||||
critical_modules = ["trl", "transformers", "peft"]
|
||||
already_imported = [mod for mod in critical_modules if mod in sys.modules]
|
||||
|
||||
# Fix some issues before importing other packages
|
||||
from .import_fixes import (
|
||||
fix_message_factory_issue,
|
||||
check_fbgemm_gpu_version,
|
||||
torchvision_compatibility_check,
|
||||
)
|
||||
|
||||
fix_message_factory_issue()
|
||||
check_fbgemm_gpu_version()
|
||||
torchvision_compatibility_check()
|
||||
del fix_message_factory_issue
|
||||
del check_fbgemm_gpu_version
|
||||
del torchvision_compatibility_check
|
||||
|
||||
# This check is critical because Unsloth optimizes these libraries by modifying
|
||||
# their code at import time. If they're imported first, the original (slower,
|
||||
# more memory-intensive) implementations will be used instead of Unsloth's
|
||||
|
|
@ -35,7 +46,7 @@ if already_imported:
|
|||
# stacklevel=2 makes warning point to user's import line rather than this library code,
|
||||
# showing them exactly where to fix the import order in their script
|
||||
warnings.warn(
|
||||
f"WARNING: Unsloth should be imported before {', '.join(already_imported)} "
|
||||
f"WARNING: Unsloth should be imported before [{', '.join(already_imported)}] "
|
||||
f"to ensure all optimizations are applied. Your code may run slower or encounter "
|
||||
f"memory issues without these optimizations.\n\n"
|
||||
f"Please restructure your imports with 'import unsloth' at the top of your file.",
|
||||
|
|
@ -55,8 +66,6 @@ os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
|
|||
# "pinned_use_cuda_host_register:True,"\
|
||||
# "pinned_num_register_threads:8"
|
||||
|
||||
# Log Unsloth is being used
|
||||
os.environ["UNSLOTH_IS_PRESENT"] = "1"
|
||||
|
||||
from importlib.metadata import version as importlib_version
|
||||
from importlib.metadata import PackageNotFoundError
|
||||
|
|
@ -64,7 +73,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.11.2"):
|
||||
if Version(unsloth_zoo_version) < Version("2025.12.4"):
|
||||
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`"
|
||||
|
|
@ -110,25 +119,34 @@ from unsloth_zoo.device_type import (
|
|||
from .import_fixes import (
|
||||
fix_xformers_performance_issue,
|
||||
fix_vllm_aimv2_issue,
|
||||
fix_vllm_guided_decoding_params,
|
||||
ignore_logger_messages,
|
||||
patch_ipykernel_hf_xet,
|
||||
patch_trackio,
|
||||
patch_datasets,
|
||||
patch_enable_input_require_grads,
|
||||
fix_openenv_no_vllm,
|
||||
)
|
||||
|
||||
fix_xformers_performance_issue()
|
||||
fix_vllm_aimv2_issue()
|
||||
fix_vllm_guided_decoding_params()
|
||||
ignore_logger_messages()
|
||||
patch_ipykernel_hf_xet()
|
||||
patch_trackio()
|
||||
patch_datasets()
|
||||
patch_enable_input_require_grads()
|
||||
fix_openenv_no_vllm()
|
||||
|
||||
del fix_xformers_performance_issue
|
||||
del fix_vllm_aimv2_issue
|
||||
del fix_vllm_guided_decoding_params
|
||||
del ignore_logger_messages
|
||||
del patch_ipykernel_hf_xet
|
||||
del patch_trackio
|
||||
del patch_datasets
|
||||
del patch_enable_input_require_grads
|
||||
del fix_openenv_no_vllm
|
||||
|
||||
# Torch 2.4 has including_emulation
|
||||
if DEVICE_TYPE == "cuda":
|
||||
|
|
|
|||
|
|
@ -16,10 +16,31 @@ import os
|
|||
import importlib.util
|
||||
from pathlib import Path
|
||||
from importlib.metadata import version as importlib_version
|
||||
from packaging.version import Version
|
||||
from packaging.version import Version as TrueVersion
|
||||
import re
|
||||
import logging
|
||||
# Cannot import logger here since it'll import transformers
|
||||
# from unsloth_zoo.log import logger
|
||||
|
||||
UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
|
||||
|
||||
def Version(version):
|
||||
try:
|
||||
new_version = str(version)
|
||||
new_version = re.match(r"[0-9\.]{1,}", new_version)
|
||||
if new_version is None:
|
||||
raise Exception(str(e))
|
||||
new_version = new_version.group(0).rstrip(".")
|
||||
if new_version != version:
|
||||
new_version += ".1" # Add .1 for dev / alpha / beta / rc
|
||||
return TrueVersion(new_version)
|
||||
except:
|
||||
from inspect import getframeinfo, stack
|
||||
|
||||
caller = getframeinfo(stack()[1][0])
|
||||
raise RuntimeError(
|
||||
f"Unsloth: Could not get version for `{version}`\n"
|
||||
f"File name = [{caller.filename}] Line number = [{caller.lineno}]"
|
||||
)
|
||||
|
||||
|
||||
# Ignore logging messages
|
||||
|
|
@ -49,9 +70,10 @@ def fix_message_factory_issue():
|
|||
def GetPrototype(self, *args, **kwargs):
|
||||
return
|
||||
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
if not hasattr(google.protobuf.message_factory, "MessageFactory"):
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
||||
logger.info("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
||||
google.protobuf.message_factory.MessageFactory = MessageFactory
|
||||
elif (
|
||||
hasattr(google.protobuf.message_factory, "MessageFactory")
|
||||
|
|
@ -61,8 +83,7 @@ def fix_message_factory_issue():
|
|||
and not hasattr(google.protobuf.message_factory, "GetMessageClass")
|
||||
):
|
||||
google.protobuf.message_factory.MessageFactory = MessageFactory
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
||||
logger.info("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
||||
elif (
|
||||
hasattr(google.protobuf.message_factory, "MessageFactory")
|
||||
and not hasattr(
|
||||
|
|
@ -76,8 +97,7 @@ def fix_message_factory_issue():
|
|||
return GetMessageClass(descriptor)
|
||||
|
||||
google.protobuf.message_factory.MessageFactory.GetPrototype = GetPrototype
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print("Unsloth: Patching protobuf.MessageFactory.GetPrototype")
|
||||
logger.info("Unsloth: Patching protobuf.MessageFactory.GetPrototype")
|
||||
pass
|
||||
except:
|
||||
pass
|
||||
|
|
@ -89,6 +109,8 @@ def fix_xformers_performance_issue():
|
|||
return
|
||||
xformers_version = importlib_version("xformers")
|
||||
if Version(xformers_version) < Version("0.0.29"):
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
xformers_location = importlib.util.find_spec("xformers").origin
|
||||
xformers_location = os.path.split(xformers_location)[0]
|
||||
cutlass = Path(xformers_location) / "ops" / "fmha" / "cutlass.py"
|
||||
|
|
@ -105,13 +127,11 @@ def fix_xformers_performance_issue():
|
|||
f.seek(0)
|
||||
f.write(text)
|
||||
f.truncate()
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print(
|
||||
"Unsloth: Patching Xformers to fix some performance issues."
|
||||
)
|
||||
logger.info(
|
||||
"Unsloth: Patching Xformers to fix some performance issues."
|
||||
)
|
||||
except Exception as e:
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print(f"Unsloth: Failed patching Xformers with error = {str(e)}")
|
||||
logger.info(f"Unsloth: Failed patching Xformers with error = {str(e)}")
|
||||
|
||||
|
||||
# ValueError: 'aimv2' is already used by a Transformers config, pick another name.
|
||||
|
|
@ -120,6 +140,8 @@ def fix_vllm_aimv2_issue():
|
|||
return
|
||||
vllm_version = importlib_version("vllm")
|
||||
if Version(vllm_version) < Version("0.10.1"):
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
vllm_version = importlib.util.find_spec("vllm").origin
|
||||
vllm_version = os.path.split(vllm_version)[0]
|
||||
ovis_config = Path(vllm_version) / "transformers_utils" / "configs" / "ovis.py"
|
||||
|
|
@ -146,13 +168,27 @@ def fix_vllm_aimv2_issue():
|
|||
f.seek(0)
|
||||
f.write(text)
|
||||
f.truncate()
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print(
|
||||
"Unsloth: Patching vLLM to fix `'aimv2' is already used by a Transformers config, pick another name.`"
|
||||
)
|
||||
logger.info(
|
||||
"Unsloth: Patching vLLM to fix `'aimv2' is already used by a Transformers config, pick another name.`"
|
||||
)
|
||||
except Exception as e:
|
||||
if UNSLOTH_ENABLE_LOGGING:
|
||||
print(f"Unsloth: Failed patching vLLM with error = {str(e)}")
|
||||
logger.info(f"Unsloth: Failed patching vLLM with error = {str(e)}")
|
||||
|
||||
|
||||
def fix_vllm_guided_decoding_params():
|
||||
if importlib.util.find_spec("vllm") is None:
|
||||
return
|
||||
# GuidedDecodingParmas is renamed to StructuredOutputsParams in vLLM
|
||||
# https://github.com/vllm-project/vllm/pull/22772/files
|
||||
# trl still wants to use GuidedDecodingParams. This is a temporary patch till trl updates
|
||||
import vllm
|
||||
|
||||
try:
|
||||
from vllm.sampling_params import GuidedDecodingParams
|
||||
except ImportError:
|
||||
vllm.sampling_params.GuidedDecodingParams = (
|
||||
vllm.sampling_params.StructuredOutputsParams
|
||||
)
|
||||
|
||||
|
||||
def ignore_logger_messages():
|
||||
|
|
@ -223,3 +259,168 @@ def patch_datasets():
|
|||
f"#### Unsloth: Using `datasets = {str(datasets_version)}` will cause recursion errors.\n"
|
||||
"Please downgrade datasets to `datasets==4.3.0"
|
||||
)
|
||||
|
||||
|
||||
def check_fbgemm_gpu_version():
|
||||
if importlib.util.find_spec("fbgemm_gpu") is None:
|
||||
return
|
||||
try:
|
||||
fbgemm_gpu_version = importlib_version("fbgemm_gpu_genai")
|
||||
except:
|
||||
return
|
||||
# We noticed some SegFault or bad alloc errors on lower versions of fbgemm_gpu.
|
||||
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!"
|
||||
)
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
logger.info(f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected.")
|
||||
|
||||
|
||||
def patch_enable_input_require_grads():
|
||||
"""
|
||||
Patch transformers PreTrainedModel.enable_input_require_grads to handle vision models
|
||||
that raise NotImplementedError from get_input_embeddings().
|
||||
|
||||
"""
|
||||
import inspect
|
||||
from transformers import PreTrainedModel
|
||||
|
||||
# Check if the original function iterates over self.modules() instead of just returning the enable_input_require_grads
|
||||
# Ref: https://github.com/huggingface/transformers/pull/41993/files#diff-6b72b98c4c2dcfc6cc606843917733f5d858374fbc22a735ff483bbc0c1e63eaL1979-R1996
|
||||
try:
|
||||
original_source = inspect.getsource(PreTrainedModel.enable_input_require_grads)
|
||||
except:
|
||||
return
|
||||
|
||||
# Only patch if the new pattern exists (iterating over self.modules())
|
||||
if "for module in self.modules()" not in original_source:
|
||||
return
|
||||
|
||||
def _patched_enable_input_require_grads(self):
|
||||
def make_inputs_require_grads(module, input, output):
|
||||
output.requires_grad_(True)
|
||||
|
||||
hooks = []
|
||||
seen_modules = set()
|
||||
|
||||
for module in self.modules():
|
||||
if not (
|
||||
isinstance(module, PreTrainedModel)
|
||||
and hasattr(module, "get_input_embeddings")
|
||||
):
|
||||
continue
|
||||
|
||||
try:
|
||||
input_embeddings = module.get_input_embeddings()
|
||||
except NotImplementedError:
|
||||
# Vision models may not implement get_input_embeddings - skip them
|
||||
# For GLM V4.6 for example, this skips only `self.visual`
|
||||
continue
|
||||
|
||||
if input_embeddings is None:
|
||||
continue
|
||||
|
||||
embedding_id = id(input_embeddings)
|
||||
if embedding_id in seen_modules:
|
||||
continue
|
||||
|
||||
seen_modules.add(embedding_id)
|
||||
hooks.append(
|
||||
input_embeddings.register_forward_hook(make_inputs_require_grads)
|
||||
)
|
||||
|
||||
self._require_grads_hooks = hooks
|
||||
if hooks:
|
||||
self._require_grads_hook = hooks[0]
|
||||
|
||||
PreTrainedModel.enable_input_require_grads = _patched_enable_input_require_grads
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
logger.info(
|
||||
"Unsloth: Patched enable_input_require_grads for vision model compatibility"
|
||||
)
|
||||
|
||||
|
||||
def torchvision_compatibility_check():
|
||||
if importlib.util.find_spec("torch") is None:
|
||||
raise ImportError("Unsloth: torch not found. Please install torch first.")
|
||||
if importlib.util.find_spec("torchvision") is None:
|
||||
return
|
||||
torch_version = importlib_version("torch")
|
||||
torchvision_version = importlib_version("torchvision")
|
||||
|
||||
# Torch version -> minimum required torchvision version
|
||||
# See https://pytorch.org/get-started/previous-versions/
|
||||
TORCH_TORCHVISION_COMPAT = [
|
||||
("2.9.0", "0.24.0"),
|
||||
("2.8.0", "0.23.0"),
|
||||
("2.7.0", "0.22.0"),
|
||||
("2.6.0", "0.21.0"),
|
||||
("2.5.0", "0.20.0"),
|
||||
("2.4.0", "0.19.0"),
|
||||
]
|
||||
|
||||
required_torchvision = None
|
||||
for min_torch, min_torchvision in TORCH_TORCHVISION_COMPAT:
|
||||
if Version(torch_version) >= Version(min_torch):
|
||||
required_torchvision = min_torchvision
|
||||
break
|
||||
|
||||
if required_torchvision is None:
|
||||
# Torch version not in compatibility table, skip check
|
||||
return
|
||||
|
||||
if Version(torchvision_version) < Version(required_torchvision):
|
||||
raise ImportError(
|
||||
f"Unsloth: torch=={torch_version} requires torchvision>={required_torchvision}, "
|
||||
f"but found torchvision=={torchvision_version}. "
|
||||
f"Please refer to https://pytorch.org/get-started/previous-versions/ for more information."
|
||||
)
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
logger.info(
|
||||
f"Unsloth: torch=={torch_version} and torchvision=={torchvision_version} are compatible."
|
||||
)
|
||||
|
||||
|
||||
# Fix TRL OpenEnv 0.26 NameError: name 'SamplingParams' is not defined
|
||||
def fix_openenv_no_vllm():
|
||||
if importlib.util.find_spec("trl") is None:
|
||||
return
|
||||
trl_location = importlib.util.find_spec("trl").origin
|
||||
trl_location = os.path.split(trl_location)[0]
|
||||
openenv = Path(trl_location) / "experimental" / "openenv" / "utils.py"
|
||||
if not openenv.exists():
|
||||
return
|
||||
from unsloth_zoo.log import logger
|
||||
|
||||
try:
|
||||
with open(openenv, "r+", encoding = "utf-8") as f:
|
||||
text = f.read()
|
||||
bad = (
|
||||
"if is_vllm_available():\n"
|
||||
" from vllm import SamplingParams\n"
|
||||
" from vllm.sampling_params import GuidedDecodingParams\n"
|
||||
)
|
||||
if bad + "\n" + "\n" in text:
|
||||
text = text.replace(
|
||||
bad + "\n" + "\n",
|
||||
bad
|
||||
+ (
|
||||
"else:\n"
|
||||
" from typing import Any\n"
|
||||
" SamplingParams = Any\n"
|
||||
" GuidedDecodingParams = Any\n"
|
||||
"\n"
|
||||
),
|
||||
)
|
||||
f.seek(0)
|
||||
f.write(text)
|
||||
f.truncate()
|
||||
logger.info(
|
||||
"Unsloth: Patching TRL OpenEnv to fix SamplingParams not defined"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.info(f"Unsloth: Failed patching TRL OpenEnv with error = {str(e)}")
|
||||
|
|
|
|||
|
|
@ -160,7 +160,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
|
|||
def forward(ctx, X: torch.Tensor, W: torch.Tensor, eps: float, gemma: bool = False):
|
||||
shape = X.shape
|
||||
dim: int = shape[-1]
|
||||
X = X.view(-1, dim)
|
||||
X = X.reshape(-1, dim)
|
||||
n_rows: int
|
||||
n_cols: int
|
||||
n_rows, n_cols = X.shape
|
||||
|
|
@ -199,7 +199,7 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
|
|||
def backward(ctx, dY: torch.Tensor):
|
||||
shape = dY.shape
|
||||
dim: int = shape[-1]
|
||||
dY = dY.view(-1, dim)
|
||||
dY = dY.reshape(-1, dim)
|
||||
X, W, r = ctx.saved_tensors
|
||||
n_rows: int
|
||||
n_cols: int
|
||||
|
|
|
|||
|
|
@ -1,22 +1,111 @@
|
|||
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
||||
# Copyright 2023-present Daniel Han-Chen, Michael 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
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Lesser General Public License as published by
|
||||
# the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU General Public License for more details.
|
||||
#
|
||||
# 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.
|
||||
# You should have received a copy of the GNU Lesser General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import triton
|
||||
import triton.language as tl
|
||||
import torch
|
||||
from ..device_type import DEVICE_COUNT
|
||||
from .utils import calculate_settings, torch_gpu_device, torch_device_stream
|
||||
|
||||
|
||||
def _rope_embedding_QK(
|
||||
Q,
|
||||
Q_batch_stride,
|
||||
Q_head_stride,
|
||||
Q_seq_stride,
|
||||
K,
|
||||
K_batch_stride,
|
||||
K_head_stride,
|
||||
K_seq_stride,
|
||||
cos,
|
||||
cos_row_stride,
|
||||
sin,
|
||||
sin_row_stride,
|
||||
rope_embedding_indices,
|
||||
seqlen,
|
||||
head_dim: tl.constexpr,
|
||||
n_heads_K: tl.constexpr,
|
||||
BACKWARD_PASS: tl.constexpr,
|
||||
HAS_ROPE_INDICES: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
row_position = tl.program_id(0)
|
||||
head_position = tl.program_id(1)
|
||||
col_offsets = tl.arange(0, BLOCK_SIZE)
|
||||
half_head_dim = head_dim // 2
|
||||
mask = col_offsets < half_head_dim
|
||||
|
||||
if HAS_ROPE_INDICES:
|
||||
rot_position = tl.load(
|
||||
rope_embedding_indices + row_position,
|
||||
eviction_policy = "evict_first",
|
||||
).to(tl.int32)
|
||||
else:
|
||||
rot_position = row_position % seqlen
|
||||
|
||||
cos_ptr = cos + rot_position * cos_row_stride
|
||||
sin_ptr = sin + rot_position * sin_row_stride
|
||||
sin1 = tl.load(
|
||||
sin_ptr + col_offsets,
|
||||
mask = mask,
|
||||
other = 0,
|
||||
)
|
||||
cos1 = tl.load(
|
||||
cos_ptr + col_offsets,
|
||||
mask = mask,
|
||||
other = 0,
|
||||
)
|
||||
if BACKWARD_PASS:
|
||||
sin1 = -sin1
|
||||
|
||||
batch_id = row_position // seqlen
|
||||
seq_index = row_position - batch_id * seqlen
|
||||
|
||||
q_ptr = (
|
||||
Q
|
||||
+ batch_id * Q_batch_stride
|
||||
+ head_position * Q_head_stride
|
||||
+ seq_index * Q_seq_stride
|
||||
)
|
||||
q0 = tl.load(q_ptr + col_offsets, mask = mask, other = 0)
|
||||
q1 = tl.load(q_ptr + half_head_dim + col_offsets, mask = mask, other = 0)
|
||||
tl.store(q_ptr + col_offsets, q0 * cos1 - q1 * sin1, mask = mask)
|
||||
tl.store(q_ptr + half_head_dim + col_offsets, q1 * cos1 + q0 * sin1, mask = mask)
|
||||
|
||||
if head_position < n_heads_K:
|
||||
k_ptr = (
|
||||
K
|
||||
+ batch_id * K_batch_stride
|
||||
+ head_position * K_head_stride
|
||||
+ seq_index * K_seq_stride
|
||||
)
|
||||
k0 = tl.load(k_ptr + col_offsets, mask = mask, other = 0)
|
||||
k1 = tl.load(k_ptr + half_head_dim + col_offsets, mask = mask, other = 0)
|
||||
tl.store(k_ptr + col_offsets, k0 * cos1 - k1 * sin1, mask = mask)
|
||||
tl.store(k_ptr + half_head_dim + col_offsets, k1 * cos1 + k0 * sin1, mask = mask)
|
||||
|
||||
|
||||
_rope_embedding_QK = triton.jit(_rope_embedding_QK)
|
||||
_rope_embedding_QK = triton.heuristics(
|
||||
{
|
||||
"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),
|
||||
"HAS_ROPE_INDICES": lambda args: bool(args["HAS_ROPE_INDICES"]),
|
||||
}
|
||||
)(_rope_embedding_QK)
|
||||
|
||||
|
||||
ROPE_GROUP_SIZE: int = 4
|
||||
|
||||
|
||||
|
|
@ -102,7 +191,7 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
n_heads: int
|
||||
head_dim: int
|
||||
batch, seq_len, n_heads, head_dim = Q.shape
|
||||
Q = Q.view(batch * seq_len, n_heads * head_dim)
|
||||
Q = Q.reshape(batch * seq_len, n_heads * head_dim)
|
||||
n_rows: int
|
||||
n_cols: int
|
||||
n_rows, n_cols = Q.shape
|
||||
|
|
@ -143,7 +232,7 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
ctx.n_groups = n_groups
|
||||
ctx.cos = cos
|
||||
ctx.sin = sin
|
||||
return Q.view(batch, seq_len, n_heads, head_dim)
|
||||
return Q.reshape(batch, seq_len, n_heads, head_dim)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dY):
|
||||
|
|
@ -153,7 +242,6 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
head_dim: int
|
||||
batch, seq_len, n_heads, head_dim = dY.shape
|
||||
dY = dY.reshape(batch * seq_len, n_heads * head_dim)
|
||||
# Must be reshape not view
|
||||
n_rows: int
|
||||
n_cols: int
|
||||
n_rows, n_cols = dY.shape
|
||||
|
|
@ -181,7 +269,7 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
BLOCK_SIZE = ctx.BLOCK_SIZE,
|
||||
num_warps = ctx.num_warps,
|
||||
)
|
||||
dY = dY.view(batch, seq_len, n_heads, head_dim)
|
||||
dY = dY.reshape(batch, seq_len, n_heads, head_dim)
|
||||
return (
|
||||
dY,
|
||||
None,
|
||||
|
|
@ -191,12 +279,150 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
|
||||
# [TODO] Unsure why RoPE Embedding is not torch.compiling properly
|
||||
@torch.compiler.disable
|
||||
def fast_rope_embedding(Q, K, cos, sin):
|
||||
Q = Fast_RoPE_Embedding.apply(Q.transpose(1, 2), cos, sin).transpose(1, 2)
|
||||
K = Fast_RoPE_Embedding.apply(K.transpose(1, 2), cos, sin).transpose(1, 2)
|
||||
# synchronize before cat to avoid race condition
|
||||
torch_device_stream(Q.device).synchronize()
|
||||
return Q, K
|
||||
def fast_rope_embedding(
|
||||
Q,
|
||||
K,
|
||||
cos,
|
||||
sin,
|
||||
rope_embedding_indices = None,
|
||||
):
|
||||
if rope_embedding_indices is not None:
|
||||
Q_out, K_out = Fast_RoPE_Embedding_QK.apply(
|
||||
Q, K, cos, sin, rope_embedding_indices
|
||||
)
|
||||
else:
|
||||
Q_out = Fast_RoPE_Embedding.apply(
|
||||
Q.transpose(1, 2).contiguous(), cos, sin
|
||||
).transpose(1, 2)
|
||||
K_out = Fast_RoPE_Embedding.apply(
|
||||
K.transpose(1, 2).contiguous(), cos, sin
|
||||
).transpose(1, 2)
|
||||
if DEVICE_COUNT > 1:
|
||||
torch_device_stream(Q.device).synchronize()
|
||||
return Q_out, K_out
|
||||
|
||||
|
||||
class Fast_RoPE_Embedding_QK(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, Q, K, cos, sin, rope_indices):
|
||||
has_indices = rope_indices is not None
|
||||
cos, sin = cos.squeeze(), sin.squeeze()
|
||||
|
||||
batch, n_heads_Q, seq_len, head_dim = Q.shape
|
||||
_, 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
|
||||
|
||||
if has_indices:
|
||||
# TRL's rotary indices are always in int32, so casting is just for safety
|
||||
rope_ptr = rope_indices.reshape(-1).to(dtype = torch.int32, device = Q.device)
|
||||
else:
|
||||
rope_ptr = cos.new_empty(1, dtype = torch.int32)
|
||||
|
||||
BLOCK_SIZE, num_warps = calculate_settings(head_dim)
|
||||
|
||||
Q_batch_stride, Q_head_stride, Q_seq_stride = (
|
||||
Q_out.stride(0),
|
||||
Q_out.stride(1),
|
||||
Q_out.stride(2),
|
||||
)
|
||||
K_batch_stride, K_head_stride, K_seq_stride = (
|
||||
K_out.stride(0),
|
||||
K_out.stride(1),
|
||||
K_out.stride(2),
|
||||
)
|
||||
|
||||
with torch_gpu_device(Q.device):
|
||||
_rope_embedding_QK[(batch * seq_len, n_heads_Q)](
|
||||
Q_out,
|
||||
Q_batch_stride,
|
||||
Q_head_stride,
|
||||
Q_seq_stride,
|
||||
K_out,
|
||||
K_batch_stride,
|
||||
K_head_stride,
|
||||
K_seq_stride,
|
||||
cos,
|
||||
cos.stride(0),
|
||||
sin,
|
||||
sin.stride(0),
|
||||
rope_ptr,
|
||||
seq_len,
|
||||
head_dim = head_dim,
|
||||
n_heads_K = n_heads_K,
|
||||
BACKWARD_PASS = False,
|
||||
HAS_ROPE_INDICES = has_indices,
|
||||
BLOCK_SIZE = BLOCK_SIZE,
|
||||
num_warps = num_warps,
|
||||
)
|
||||
|
||||
ctx.block_size = BLOCK_SIZE
|
||||
ctx.num_warps = num_warps
|
||||
ctx.has_indices = has_indices
|
||||
ctx.cos = cos
|
||||
ctx.sin = sin
|
||||
ctx.rope_indices = rope_ptr if has_indices else None
|
||||
ctx.seq_len = seq_len
|
||||
ctx.n_heads_Q = n_heads_Q
|
||||
ctx.n_heads_K = n_heads_K
|
||||
|
||||
return (
|
||||
Q_out,
|
||||
K_out,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dQ, dK):
|
||||
batch, _, _, head_dim = dQ.shape
|
||||
|
||||
rope_ptr = (
|
||||
ctx.rope_indices
|
||||
if ctx.has_indices
|
||||
else ctx.cos.new_empty(1, dtype = torch.int32)
|
||||
)
|
||||
|
||||
Q_batch_stride, Q_head_stride, Q_seq_stride = (
|
||||
dQ.stride(0),
|
||||
dQ.stride(1),
|
||||
dQ.stride(2),
|
||||
)
|
||||
K_batch_stride, K_head_stride, K_seq_stride = (
|
||||
dK.stride(0),
|
||||
dK.stride(1),
|
||||
dK.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,
|
||||
Q_batch_stride,
|
||||
Q_head_stride,
|
||||
Q_seq_stride,
|
||||
dK_out,
|
||||
K_batch_stride,
|
||||
K_head_stride,
|
||||
K_seq_stride,
|
||||
ctx.cos,
|
||||
ctx.cos.stride(0),
|
||||
ctx.sin,
|
||||
ctx.sin.stride(0),
|
||||
rope_ptr,
|
||||
ctx.seq_len,
|
||||
head_dim = head_dim,
|
||||
n_heads_K = ctx.n_heads_K,
|
||||
BACKWARD_PASS = True,
|
||||
HAS_ROPE_INDICES = ctx.has_indices,
|
||||
BLOCK_SIZE = ctx.block_size,
|
||||
num_warps = ctx.num_warps,
|
||||
)
|
||||
|
||||
return (dQ_out, dK_out, None, None, None)
|
||||
|
||||
|
||||
class Slow_RoPE_Embedding(torch.autograd.Function):
|
||||
|
|
@ -206,8 +432,8 @@ class Slow_RoPE_Embedding(torch.autograd.Function):
|
|||
# The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
|
||||
cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
|
||||
sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
|
||||
cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
|
||||
sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
|
||||
cos = cos[position_ids].unsqueeze(2) # [bs, seq_len, 1, dim]
|
||||
sin = sin[position_ids].unsqueeze(2) # [bs, seq_len, 1, dim]
|
||||
|
||||
# Q * cos + rotate_half(Q) * sin
|
||||
half = Q.shape[-1] // 2
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
__version__ = "2025.11.4"
|
||||
__version__ = "2025.12.5"
|
||||
|
||||
__all__ = [
|
||||
"SUPPORTS_BFLOAT16",
|
||||
|
|
@ -85,6 +85,7 @@ import re
|
|||
from dataclasses import dataclass, field
|
||||
import functools
|
||||
import textwrap
|
||||
import logging
|
||||
import warnings, subprocess, inspect, psutil, os, math
|
||||
from unsloth_zoo.utils import Version, get_quant_type
|
||||
from importlib.metadata import version as importlib_version
|
||||
|
|
@ -167,9 +168,9 @@ warnings.filterwarnings(
|
|||
)
|
||||
warnings.filterwarnings(action = "ignore", category = RuntimeWarning, module = "multiprocess")
|
||||
warnings.filterwarnings(action = "ignore", category = UserWarning, module = "triton")
|
||||
# Stop "Special tokens have been added in the vocabulary, ..."
|
||||
import logging
|
||||
warnings.filterwarnings(action = "ignore", category = UserWarning, module = "bitsandbytes")
|
||||
|
||||
# Stop "Special tokens have been added in the vocabulary, ..."
|
||||
logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITICAL + 1)
|
||||
|
||||
|
||||
|
|
@ -412,6 +413,16 @@ try:
|
|||
except:
|
||||
pass
|
||||
|
||||
# Flax classes are deprecated and will be removed in Diffusers v1.0.0.
|
||||
try:
|
||||
from diffusers.utils import logger as diffusers_logger
|
||||
|
||||
diffusers_logger.addFilter(HideLoggingMessage("are deprecated"))
|
||||
del diffusers_logger
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
# Errors out on
|
||||
# Some weights of Gemma3nForConditionalGeneration were not initialized from the model checkpoint
|
||||
from transformers.modeling_utils import logger as transformers_logger
|
||||
|
|
@ -761,6 +772,13 @@ elif DEVICE_TYPE == "xpu":
|
|||
|
||||
# =============================================
|
||||
# Get Xformers
|
||||
# Silence xformers CUDA mismatch warnings before import
|
||||
try:
|
||||
_xformers_logger = logging.getLogger("xformers")
|
||||
_xformers_logger.setLevel(logging.ERROR)
|
||||
del _xformers_logger
|
||||
except:
|
||||
pass
|
||||
try:
|
||||
from xformers import __version__ as xformers_version
|
||||
|
||||
|
|
@ -835,10 +853,11 @@ except ModuleNotFoundError:
|
|||
xformers_attention = None
|
||||
xformers_version = None
|
||||
except Exception as e:
|
||||
print(
|
||||
"========\nSwitching to PyTorch attention since your Xformers is broken.\n========\n"
|
||||
)
|
||||
print(str(e))
|
||||
if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") != "0":
|
||||
print(
|
||||
"========\nSwitching to PyTorch attention since your Xformers is broken.\n========\n"
|
||||
)
|
||||
print(str(e))
|
||||
xformers = None
|
||||
xformers_attention = None
|
||||
xformers_version = None
|
||||
|
|
@ -2302,17 +2321,20 @@ def verify_fp8_support_if_applicable(model_config):
|
|||
raise ValueError(
|
||||
f"Unsloth: FP8 quantization is only supported on CUDA GPUs. You are using {DEVICE_TYPE}."
|
||||
)
|
||||
major_version, minor_version = torch.cuda.get_device_capability()
|
||||
if quant_method == "fbgemm_fp8" and major_version < 9:
|
||||
# While L4 does support FP8 as data type, it doesn't have fbgemm (package) support yet. So we restrict it.
|
||||
raise ValueError(
|
||||
f"Unsloth: FBGEMM FP8 quantization is only supported on H100 and higher GPUs. L4 is not supported. You are using {torch.cuda.get_device_name()}. Refer to https://developer.nvidia.com/cuda-gpus for more details."
|
||||
)
|
||||
if quant_method == "fp8" and major_version * 10 + minor_version < 89:
|
||||
# In case of block quantized, we allow L4 because we fall back to torchao kernels.
|
||||
raise ValueError(
|
||||
f"Unsloth: FP8 quantization is only supported on L4 and higher GPUs with compute capability 8.9 or higher. You are using {torch.cuda.get_device_name()}. Refer to https://developer.nvidia.com/cuda-gpus for more details."
|
||||
)
|
||||
|
||||
# [TODO] Need to add FP8 support for Intel XPUs
|
||||
if DEVICE_TYPE == "cuda":
|
||||
major_version, minor_version = torch.cuda.get_device_capability()
|
||||
if quant_method == "fbgemm_fp8" and major_version < 9:
|
||||
# While L4 does support FP8 as data type, it doesn't have fbgemm (package) support yet. So we restrict it.
|
||||
raise ValueError(
|
||||
f"Unsloth: FBGEMM FP8 quantization is only supported on H100 and higher GPUs. L4 is not supported. You are using {torch.cuda.get_device_name()}. Refer to https://developer.nvidia.com/cuda-gpus for more details."
|
||||
)
|
||||
if quant_method == "fp8" and major_version * 10 + minor_version < 89:
|
||||
# In case of block quantized, we allow L4 because we fall back to torchao kernels.
|
||||
raise ValueError(
|
||||
f"Unsloth: FP8 quantization is only supported on L4 and higher GPUs with compute capability 8.9 or higher. You are using {torch.cuda.get_device_name()}. Refer to https://developer.nvidia.com/cuda-gpus for more details."
|
||||
)
|
||||
|
||||
|
||||
def _get_inference_mode_context_manager(model: torch.nn.Module):
|
||||
|
|
|
|||
|
|
@ -16,6 +16,13 @@ from .llama import *
|
|||
from ._utils import __version__
|
||||
from unsloth_zoo.hf_utils import dtype_from_config
|
||||
from unsloth_zoo.utils import _get_dtype
|
||||
from ..utils.packing import get_packed_info_from_kwargs
|
||||
from ..utils.attention_dispatch import (
|
||||
AttentionConfig,
|
||||
AttentionContext,
|
||||
run_attention,
|
||||
select_attention_backend,
|
||||
)
|
||||
|
||||
try:
|
||||
from transformers.models.cohere.modeling_cohere import (
|
||||
|
|
@ -104,6 +111,7 @@ def CohereAttention_fast_forward(
|
|||
Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
|
||||
K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
seq_info = get_packed_info_from_kwargs(kwargs, Q.device)
|
||||
if self.use_qk_norm:
|
||||
Q = fast_layernorm_compiled(self.q_norm, Q)
|
||||
K = fast_layernorm_compiled(self.k_norm, K)
|
||||
|
|
@ -112,12 +120,17 @@ def CohereAttention_fast_forward(
|
|||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
|
||||
cos, sin = position_embeddings
|
||||
if position_ids is None:
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
# Extend RoPE dynamically to fit in VRAM
|
||||
if position_embeddings:
|
||||
cos, sin = position_embeddings
|
||||
else:
|
||||
cos, sin = cos[position_ids], sin[position_ids]
|
||||
Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids)
|
||||
cos, sin = self.rotary_emb.get_cached(kv_seq_len, Q.device.index)
|
||||
|
||||
rope_position_ids = (
|
||||
position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
)
|
||||
# Useful for LongRoPE
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
|
||||
if past_key_value is not None:
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
|
|
@ -125,54 +138,33 @@ def CohereAttention_fast_forward(
|
|||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
if not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None:
|
||||
# Xformers memory efficient attention
|
||||
# Also has Flash Attention v2 dispatching
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
use_varlen = seq_info is not None and past_key_value is None
|
||||
backend = select_attention_backend(use_varlen)
|
||||
attention_config = AttentionConfig(
|
||||
backend = backend,
|
||||
n_kv_heads = n_kv_heads,
|
||||
n_groups = n_groups,
|
||||
flash_dense_kwargs = {"causal": True},
|
||||
flash_varlen_kwargs = {
|
||||
"dropout_p": 0.0,
|
||||
"causal": True,
|
||||
"softmax_scale": getattr(self, "softmax_scale", None),
|
||||
},
|
||||
)
|
||||
context = AttentionContext(
|
||||
bsz = bsz,
|
||||
q_len = q_len,
|
||||
kv_seq_len = kv_seq_len,
|
||||
n_heads = n_heads,
|
||||
head_dim = head_dim,
|
||||
requires_grad = hidden_states.requires_grad,
|
||||
seq_info = seq_info,
|
||||
attention_mask = attention_mask,
|
||||
causal_mask = causal_mask,
|
||||
)
|
||||
|
||||
# Group query attention
|
||||
if n_groups != 1:
|
||||
K = K.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
V = V.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
else:
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
A = xformers_attention(Q, K, V, attn_bias = causal_mask)
|
||||
A = A.view(bsz, q_len, n_heads, head_dim)
|
||||
A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)
|
||||
|
||||
elif HAS_FLASH_ATTENTION and attention_mask is None:
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
A = flash_attn_func(Q, K, V, causal = True)
|
||||
else:
|
||||
# Grouped query attention
|
||||
if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(
|
||||
bsz, n_kv_heads, n_groups, kv_seq_len, head_dim
|
||||
)
|
||||
V = V[:, :, None, :, :].expand(
|
||||
bsz, n_kv_heads, n_groups, kv_seq_len, head_dim
|
||||
)
|
||||
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
pass
|
||||
# Must be contiguous or else results are False!
|
||||
# https://github.com/pytorch/pytorch/issues/112577
|
||||
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
A = scaled_dot_product_attention(
|
||||
Q, K, V, attn_mask = attention_mask, is_causal = False
|
||||
)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2).contiguous()
|
||||
attn_output = A.reshape(bsz, q_len, n_heads * head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
attn_weights = None
|
||||
|
|
@ -215,6 +207,7 @@ def CohereDecoderLayer_fast_forward(
|
|||
output_attentions = output_attentions,
|
||||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Fully Connected
|
||||
|
|
@ -234,6 +227,7 @@ def CohereDecoderLayer_fast_forward(
|
|||
output_attentions = output_attentions,
|
||||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Fully Connected
|
||||
|
|
|
|||
|
|
@ -17,6 +17,14 @@ import os
|
|||
from ._utils import __version__
|
||||
from unsloth_zoo.utils import Version, _get_dtype
|
||||
from unsloth_zoo.hf_utils import dtype_from_config
|
||||
from ..utils.packing import get_packed_info_from_kwargs
|
||||
from ..utils.attention_dispatch import (
|
||||
AttentionConfig,
|
||||
AttentionContext,
|
||||
run_attention,
|
||||
select_attention_backend,
|
||||
SDPA,
|
||||
)
|
||||
from .llama import (
|
||||
LlamaRotaryEmbedding,
|
||||
LlamaLinearScalingRotaryEmbedding,
|
||||
|
|
@ -98,13 +106,10 @@ def FalconH1Attention_fast_forward(
|
|||
assert n_kv_heads * n_groups == n_heads
|
||||
|
||||
Q, K, V = self.apply_qkv(self, hidden_states)
|
||||
Q = Q.view(
|
||||
bsz, q_len, n_heads, head_dim
|
||||
) # .transpose(1, 2) # we will transpose after normalisation
|
||||
K = K.view(
|
||||
bsz, q_len, n_kv_heads, head_dim
|
||||
) # .transpose(1, 2) # we will transpose after normalisation
|
||||
Q = Q.view(bsz, q_len, n_heads, head_dim)
|
||||
K = K.view(bsz, q_len, n_kv_heads, head_dim)
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
seq_info = get_packed_info_from_kwargs(kwargs, hidden_states.device)
|
||||
|
||||
# Falcon H1 multiplies key states by a multiplier
|
||||
K = K * self.config.key_multiplier
|
||||
|
|
@ -116,20 +121,19 @@ def FalconH1Attention_fast_forward(
|
|||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
|
||||
# Extend RoPE dynamically to fit in VRAM
|
||||
if position_embeddings and kv_seq_len <= position_embeddings[0].shape[0]:
|
||||
cos, sin = position_embeddings
|
||||
else:
|
||||
# Extend RoPE dynamically to fit in VRA
|
||||
rotary_emb = self.rotary_emb
|
||||
rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
|
||||
device_index = Q.device.index
|
||||
cos, sin = rotary_emb.get_cached(kv_seq_len, Q.device.index)
|
||||
|
||||
if position_ids is None:
|
||||
# Useful for LongRoPE
|
||||
cos, sin = rotary_emb.get_cached(kv_seq_len, device_index)
|
||||
else:
|
||||
cos, sin = rotary_emb.get_cached(kv_seq_len, device_index)
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
rope_position_ids = (
|
||||
position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
)
|
||||
# Useful for LongRoPE
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
|
||||
if past_key_value is not None:
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
|
|
@ -137,54 +141,45 @@ def FalconH1Attention_fast_forward(
|
|||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
if not HAS_FLASH_ATTENTION and attention_mask is None:
|
||||
# Xformers memory efficient attention
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
K_M = V_M = bsz * kv_seq_len
|
||||
Q_M = bsz * q_len
|
||||
window = (-1, -1)
|
||||
use_varlen = (
|
||||
attention_mask is None
|
||||
and seq_info is not None
|
||||
and past_key_value is None
|
||||
and window == (-1, -1)
|
||||
)
|
||||
|
||||
# Group query attention
|
||||
K = K.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
V = V.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
else:
|
||||
# Xformers does support the forward pass though
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
backend = (
|
||||
SDPA if attention_mask is not None else select_attention_backend(use_varlen)
|
||||
)
|
||||
attention_config = AttentionConfig(
|
||||
backend = backend,
|
||||
n_kv_heads = n_kv_heads,
|
||||
n_groups = n_groups,
|
||||
flash_dense_kwargs = {
|
||||
"causal": True,
|
||||
"window_size": (kv_seq_len, kv_seq_len),
|
||||
},
|
||||
flash_varlen_kwargs = {
|
||||
"dropout_p": 0.0,
|
||||
"softmax_scale": None,
|
||||
"causal": True,
|
||||
},
|
||||
sdpa_kwargs = {} if attention_mask is None else {"attn_mask": attention_mask},
|
||||
)
|
||||
context = AttentionContext(
|
||||
bsz = bsz,
|
||||
q_len = q_len,
|
||||
kv_seq_len = kv_seq_len,
|
||||
n_heads = n_heads,
|
||||
head_dim = head_dim,
|
||||
requires_grad = hidden_states.requires_grad,
|
||||
seq_info = seq_info,
|
||||
attention_mask = attention_mask,
|
||||
causal_mask = causal_mask,
|
||||
)
|
||||
|
||||
A = xformers_attention(Q, K, V, attn_bias = causal_mask)
|
||||
A = A.view(bsz, q_len, n_heads, head_dim)
|
||||
|
||||
elif HAS_FLASH_ATTENTION and attention_mask is None:
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
sw = kv_seq_len
|
||||
window = (-1, -1) if (kv_seq_len <= sw) else (sw, sw)
|
||||
A = flash_attn_func(Q, K, V, causal = True, window_size = window)
|
||||
else:
|
||||
# Grouped query attention
|
||||
# if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
# pass
|
||||
# Must be contiguous or else results are False!
|
||||
# https://github.com/pytorch/pytorch/issues/112577
|
||||
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
A = scaled_dot_product_attention(
|
||||
Q, K, V, attn_mask = attention_mask, is_causal = False
|
||||
)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2).contiguous()
|
||||
A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)
|
||||
|
||||
attn_output = A.reshape(bsz, q_len, n_heads * head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
|
|
@ -442,6 +437,7 @@ def FalconH1DecoderLayer_fast_forward(
|
|||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
position_embeddings = position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
attention_hidden_states = attention_hidden_states * self.attn_out_multiplier
|
||||
|
||||
|
|
@ -486,6 +482,7 @@ def FalconH1DecoderLayer_fast_forward(
|
|||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
position_embeddings = position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
attention_hidden_states = attention_hidden_states * self.attn_out_multiplier
|
||||
|
||||
|
|
@ -726,7 +723,7 @@ class FastFalconH1Model(FastLlamaModel):
|
|||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# https://github.com/huggingface/transformers/pull/27931
|
||||
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
|
||||
import transformers.models.falcon_h1.modeling_falcon_h1
|
||||
|
|
|
|||
|
|
@ -16,6 +16,11 @@ from .llama import *
|
|||
from ._utils import __version__
|
||||
from unsloth_zoo.utils import _get_dtype
|
||||
from unsloth_zoo.hf_utils import dtype_from_config
|
||||
from ..utils.packing import (
|
||||
build_sdpa_packed_attention_mask,
|
||||
build_xformers_block_causal_mask,
|
||||
get_packed_info_from_kwargs,
|
||||
)
|
||||
import math
|
||||
|
||||
try:
|
||||
|
|
@ -110,6 +115,7 @@ def GemmaDecoderLayer_fast_forward(
|
|||
output_attentions = output_attentions,
|
||||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states += residual
|
||||
|
||||
|
|
@ -134,6 +140,7 @@ def GemmaDecoderLayer_fast_forward(
|
|||
output_attentions = output_attentions,
|
||||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
|
|
@ -164,6 +171,7 @@ def GemmaModel_fast_forward_inference(
|
|||
past_key_values,
|
||||
position_ids,
|
||||
attention_mask = None,
|
||||
**kwargs,
|
||||
):
|
||||
out_weights = tuple(
|
||||
torch.empty_like(
|
||||
|
|
@ -416,7 +424,7 @@ class FastGemmaModel(FastLlamaModel):
|
|||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# https://github.com/huggingface/transformers/pull/27931
|
||||
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
|
||||
import transformers.models.gemma.modeling_gemma
|
||||
|
|
|
|||
|
|
@ -16,6 +16,14 @@ from .llama import *
|
|||
from ._utils import __version__
|
||||
from unsloth_zoo.utils import _get_dtype
|
||||
from unsloth_zoo.hf_utils import dtype_from_config
|
||||
from ..utils.packing import get_packed_info_from_kwargs
|
||||
from ..utils.attention_dispatch import (
|
||||
AttentionConfig,
|
||||
AttentionContext,
|
||||
run_attention,
|
||||
select_attention_backend,
|
||||
SDPA,
|
||||
)
|
||||
from .gemma import (
|
||||
GemmaFixedRotaryEmbedding,
|
||||
GemmaFixedLinearScalingRotaryEmbedding,
|
||||
|
|
@ -98,19 +106,25 @@ def Gemma2Attention_fast_forward(
|
|||
Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
|
||||
K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
seq_info = get_packed_info_from_kwargs(kwargs, Q.device)
|
||||
|
||||
kv_seq_len = K.shape[-2]
|
||||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
|
||||
device_index = Q.device.index
|
||||
if position_ids is None:
|
||||
cos = self.rotary_emb.multi_gpu_cos_cached[device_index]
|
||||
sin = self.rotary_emb.multi_gpu_sin_cached[device_index]
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
cos = self.rotary_emb.multi_gpu_cos_cached[device_index]
|
||||
sin = self.rotary_emb.multi_gpu_sin_cached[device_index]
|
||||
|
||||
rope_position_ids = (
|
||||
position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
)
|
||||
if rope_position_ids is not None:
|
||||
# Useful for LongRoPE
|
||||
cos_var, sin_var = self.rotary_emb.get_cached(kv_seq_len, device_index)
|
||||
Q, K = fast_rope_embedding(Q, K, cos_var, sin_var, rope_position_ids)
|
||||
else:
|
||||
cos, sin = self.rotary_emb.get_cached(kv_seq_len, device_index)
|
||||
Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids)
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
|
||||
if past_key_value is not None:
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
|
|
@ -118,32 +132,68 @@ def Gemma2Attention_fast_forward(
|
|||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Only enable if the attention_mask is True
|
||||
has_sliding_window = type(causal_mask) is bool and causal_mask is True
|
||||
if HAS_FLASH_ATTENTION_SOFTCAPPING and attention_mask is None:
|
||||
window = (-1, -1)
|
||||
if has_sliding_window:
|
||||
sw = getattr(self.config, "sliding_window", None)
|
||||
sw = kv_seq_len if (sw is None or sw == "null") else sw
|
||||
window = (-1, -1) if (kv_seq_len <= sw) else (sw, sw)
|
||||
use_sliding_window = kwargs.get("use_sliding_window")
|
||||
has_sliding_window = (
|
||||
use_sliding_window
|
||||
if use_sliding_window is not None
|
||||
else isinstance(causal_mask, bool) and causal_mask is True
|
||||
)
|
||||
|
||||
use_flash = HAS_FLASH_ATTENTION_SOFTCAPPING and attention_mask is None
|
||||
|
||||
if use_flash:
|
||||
window = (-1, -1)
|
||||
sliding_window = getattr(self.config, "sliding_window", None)
|
||||
if has_sliding_window:
|
||||
sliding_window = (
|
||||
sliding_window if sliding_window is not None else kv_seq_len
|
||||
)
|
||||
window = (
|
||||
(-1, -1)
|
||||
if kv_seq_len <= sliding_window
|
||||
else (sliding_window, sliding_window)
|
||||
)
|
||||
|
||||
# FA uses 1 / sqrt for softmax_scale!
|
||||
if not hasattr(self, "_flash_attention_softmax_scale"):
|
||||
self._flash_attention_softmax_scale = 1.0 / (
|
||||
self.config.query_pre_attn_scalar**0.5
|
||||
)
|
||||
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
A = flash_attn_func(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
causal = True,
|
||||
softcap = self.config.attn_logit_softcapping,
|
||||
softmax_scale = self._flash_attention_softmax_scale,
|
||||
window_size = window,
|
||||
use_varlen = seq_info is not None and past_key_value is None
|
||||
|
||||
attention_config = AttentionConfig(
|
||||
backend = select_attention_backend(use_varlen),
|
||||
n_kv_heads = n_kv_heads,
|
||||
n_groups = n_groups,
|
||||
flash_dense_kwargs = {
|
||||
"causal": True,
|
||||
"softcap": self.config.attn_logit_softcapping,
|
||||
"softmax_scale": self._flash_attention_softmax_scale,
|
||||
"window_size": window,
|
||||
},
|
||||
flash_varlen_kwargs = {
|
||||
"dropout_p": 0.0,
|
||||
"softmax_scale": self._flash_attention_softmax_scale,
|
||||
"causal": True,
|
||||
"softcap": self.config.attn_logit_softcapping,
|
||||
"window_size": window,
|
||||
},
|
||||
)
|
||||
|
||||
context = AttentionContext(
|
||||
bsz = bsz,
|
||||
q_len = q_len,
|
||||
kv_seq_len = kv_seq_len,
|
||||
n_heads = n_heads,
|
||||
head_dim = head_dim,
|
||||
requires_grad = hidden_states.requires_grad,
|
||||
seq_info = seq_info,
|
||||
attention_mask = attention_mask,
|
||||
causal_mask = causal_mask,
|
||||
sliding_window = sliding_window,
|
||||
)
|
||||
|
||||
A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)
|
||||
A = A.reshape(bsz, q_len, n_heads * head_dim)
|
||||
else:
|
||||
fx = (
|
||||
|
|
@ -192,6 +242,7 @@ def Gemma2DecoderLayer_fast_forward(
|
|||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
_flag_for_generation = self._flag_for_generation,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = fast_rms_layernorm_inference_gemma(
|
||||
self.post_attention_layernorm, hidden_states, out_weight
|
||||
|
|
@ -222,6 +273,7 @@ def Gemma2DecoderLayer_fast_forward(
|
|||
output_attentions = output_attentions,
|
||||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = fast_rms_layernorm(
|
||||
self.post_attention_layernorm, hidden_states, gemma = True
|
||||
|
|
@ -263,6 +315,7 @@ def Gemma2Attention_fast_forward_inference(
|
|||
do_prefill = False,
|
||||
attention_mask = None,
|
||||
use_sliding_window = False,
|
||||
**kwargs,
|
||||
):
|
||||
Xn = hidden_states
|
||||
bsz, _, hd = hidden_states.size()
|
||||
|
|
@ -426,6 +479,7 @@ def Gemma2Model_fast_forward_inference(
|
|||
past_key_values,
|
||||
position_ids,
|
||||
attention_mask = None,
|
||||
**kwargs,
|
||||
):
|
||||
out_weights = tuple(
|
||||
torch.empty_like(
|
||||
|
|
@ -550,7 +604,7 @@ class FastGemma2Model(FastLlamaModel):
|
|||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# https://github.com/huggingface/transformers/pull/27931
|
||||
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
|
||||
import transformers.models.gemma2.modeling_gemma2
|
||||
|
|
|
|||
|
|
@ -17,6 +17,14 @@ import os
|
|||
from ._utils import __version__
|
||||
from unsloth_zoo.utils import _get_dtype
|
||||
from unsloth_zoo.hf_utils import dtype_from_config
|
||||
from ..utils.packing import get_packed_info_from_kwargs
|
||||
from ..utils.attention_dispatch import (
|
||||
AttentionConfig,
|
||||
AttentionContext,
|
||||
run_attention,
|
||||
select_attention_backend,
|
||||
SDPA,
|
||||
)
|
||||
from .llama import (
|
||||
LlamaRotaryEmbedding,
|
||||
LlamaLinearScalingRotaryEmbedding,
|
||||
|
|
@ -96,6 +104,7 @@ def GraniteAttention_fast_forward(
|
|||
Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
|
||||
K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
seq_info = get_packed_info_from_kwargs(kwargs, Q.device)
|
||||
|
||||
kv_seq_len = K.shape[-2]
|
||||
if past_key_value is not None:
|
||||
|
|
@ -103,10 +112,14 @@ def GraniteAttention_fast_forward(
|
|||
|
||||
assert position_embeddings is not None
|
||||
cos, sin = position_embeddings
|
||||
if position_ids is None:
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
rope_position_ids = (
|
||||
position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
)
|
||||
if rope_position_ids is not None:
|
||||
# Useful for LongRoPE
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
else:
|
||||
Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids)
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
|
||||
if past_key_value is not None:
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
|
|
@ -114,69 +127,59 @@ def GraniteAttention_fast_forward(
|
|||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
if not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None:
|
||||
# Xformers memory efficient attention
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
K_M = V_M = bsz * kv_seq_len
|
||||
Q_M = bsz * q_len
|
||||
use_varlen = (
|
||||
attention_mask is None and seq_info is not None and past_key_value is None
|
||||
)
|
||||
|
||||
# Group query attention
|
||||
K = K.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
V = V.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
else:
|
||||
# Xformers does support the forward pass though
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
backend = (
|
||||
SDPA if attention_mask is not None else select_attention_backend(use_varlen)
|
||||
)
|
||||
|
||||
A = xformers_attention(
|
||||
Q, K, V, attn_bias = causal_mask, scale = self.scaling, p = dropout_p
|
||||
)
|
||||
A = A.view(bsz, q_len, n_heads, head_dim)
|
||||
window = (kv_seq_len, kv_seq_len)
|
||||
softmax_scale = getattr(self, "scaling", None)
|
||||
attention_config = AttentionConfig(
|
||||
backend = backend,
|
||||
n_kv_heads = n_kv_heads,
|
||||
n_groups = n_groups,
|
||||
flash_dense_kwargs = {
|
||||
"causal": True,
|
||||
"softmax_scale": softmax_scale,
|
||||
"dropout_p": dropout_p,
|
||||
"window_size": window,
|
||||
},
|
||||
flash_varlen_kwargs = {
|
||||
"dropout_p": 0.0,
|
||||
"softmax_scale": softmax_scale,
|
||||
"causal": True,
|
||||
},
|
||||
sdpa_kwargs = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"attn_mask": attention_mask,
|
||||
"scale": softmax_scale,
|
||||
"dropout_p": dropout_p,
|
||||
}.items()
|
||||
if v is not None
|
||||
},
|
||||
xformers_kwargs = {
|
||||
"scale": softmax_scale,
|
||||
"p": dropout_p,
|
||||
},
|
||||
)
|
||||
|
||||
elif HAS_FLASH_ATTENTION and attention_mask is None:
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
window = (kv_seq_len, kv_seq_len)
|
||||
A = flash_attn_func(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
causal = True,
|
||||
window_size = window,
|
||||
softmax_scale = self.scaling,
|
||||
dropout_p = dropout_p,
|
||||
)
|
||||
else:
|
||||
# Grouped query attention
|
||||
# if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
# pass
|
||||
# Must be contiguous or else results are False!
|
||||
# https://github.com/pytorch/pytorch/issues/112577
|
||||
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
A = scaled_dot_product_attention(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
attn_mask = attention_mask,
|
||||
scale = self.scaling,
|
||||
is_causal = False,
|
||||
dropout_p = dropout_p,
|
||||
)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2).contiguous()
|
||||
context = AttentionContext(
|
||||
bsz = bsz,
|
||||
q_len = q_len,
|
||||
kv_seq_len = kv_seq_len,
|
||||
n_heads = n_heads,
|
||||
head_dim = head_dim,
|
||||
requires_grad = hidden_states.requires_grad,
|
||||
seq_info = seq_info,
|
||||
attention_mask = attention_mask,
|
||||
causal_mask = causal_mask,
|
||||
)
|
||||
|
||||
A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)
|
||||
|
||||
attn_output = A.reshape(bsz, q_len, n_heads * head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
|
|
@ -222,6 +225,7 @@ def GraniteDecoderLayer_fast_forward(
|
|||
padding_mask = padding_mask,
|
||||
position_embeddings = position_embeddings,
|
||||
_flag_for_generation = self._flag_for_generation,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = torch.add(residual, hidden_states, alpha = residual_multiplier)
|
||||
|
||||
|
|
@ -245,6 +249,7 @@ def GraniteDecoderLayer_fast_forward(
|
|||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
position_embeddings = position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = torch.add(residual, hidden_states, alpha = residual_multiplier)
|
||||
|
||||
|
|
|
|||
|
|
@ -16,7 +16,8 @@ import torch
|
|||
import gc
|
||||
import math
|
||||
import functools
|
||||
from typing import Any, Dict, Optional, Tuple, List, Union
|
||||
from typing import Optional, Tuple, List, Union
|
||||
|
||||
from ._utils import *
|
||||
from ._utils import patch_unsloth_smart_gradient_checkpointing
|
||||
from ._utils import __version__, importlib_version
|
||||
|
|
@ -25,6 +26,16 @@ from ._utils import (
|
|||
_get_inference_mode_context_manager,
|
||||
_prepare_model_for_qat,
|
||||
)
|
||||
from ..utils.packing import (
|
||||
get_packed_info_from_kwargs,
|
||||
mask_packed_sequence_boundaries,
|
||||
)
|
||||
from ..utils.attention_dispatch import (
|
||||
AttentionConfig,
|
||||
AttentionContext,
|
||||
run_attention,
|
||||
select_attention_backend,
|
||||
)
|
||||
from torch.nn.functional import scaled_dot_product_attention
|
||||
from transformers import __version__ as transformers_version
|
||||
from unsloth_zoo.utils import Version, _get_dtype
|
||||
|
|
@ -61,9 +72,6 @@ from transformers.modeling_attn_mask_utils import (
|
|||
)
|
||||
from ..kernels import *
|
||||
from ..tokenizer_utils import *
|
||||
|
||||
if HAS_FLASH_ATTENTION:
|
||||
from flash_attn import flash_attn_func
|
||||
from .vision import FastBaseModel
|
||||
|
||||
# Final patching code
|
||||
|
|
@ -551,7 +559,6 @@ def LlamaAttention_fast_forward(
|
|||
del self.temp_KV
|
||||
del self.RH_Q
|
||||
del self.attention
|
||||
|
||||
bsz, q_len, _ = hidden_states.size()
|
||||
|
||||
n_heads = self.config.num_attention_heads
|
||||
|
|
@ -564,6 +571,7 @@ def LlamaAttention_fast_forward(
|
|||
Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
|
||||
K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
seq_info = get_packed_info_from_kwargs(kwargs, Q.device)
|
||||
|
||||
kv_seq_len = K.shape[-2]
|
||||
if past_key_value is not None:
|
||||
|
|
@ -572,23 +580,20 @@ def LlamaAttention_fast_forward(
|
|||
if position_embeddings and kv_seq_len <= position_embeddings[0].shape[0]:
|
||||
cos, sin = position_embeddings
|
||||
else:
|
||||
# Extend RoPE dynamically to fit in VRA
|
||||
rotary_emb = self.rotary_emb
|
||||
rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
|
||||
|
||||
# if position_ids is None:
|
||||
# # Useful for LongRoPE
|
||||
# cos, sin = rotary_emb.get_cached(kv_seq_len, device = Q.device)
|
||||
# else:
|
||||
# cos, sin = rotary_emb.get_cached(seq_len = kv_seq_len, device = Q.device)
|
||||
cos, sin = rotary_emb.get_cached(kv_seq_len, Q.device.index)
|
||||
|
||||
rope_position_ids = position_ids
|
||||
if rope_position_ids is None and seq_info is not None:
|
||||
rope_position_ids = kwargs.get("position_ids")
|
||||
|
||||
# Q, K = (
|
||||
# fast_rope_embedding(Q, K, cos, sin)
|
||||
# if position_ids is None
|
||||
# else inplace_rope_embedding(Q, K, cos, sin, position_ids)
|
||||
# if rope_position_ids is None
|
||||
# else inplace_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
# )
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
|
||||
if past_key_value is not None:
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
|
|
@ -596,76 +601,28 @@ def LlamaAttention_fast_forward(
|
|||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
if not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None:
|
||||
# Xformers memory efficient attention
|
||||
# Also has Flash Attention v2 dispatching
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
use_varlen = seq_info is not None and past_key_value is None
|
||||
backend = select_attention_backend(use_varlen)
|
||||
config = AttentionConfig(
|
||||
backend = backend,
|
||||
n_kv_heads = n_kv_heads,
|
||||
n_groups = n_groups,
|
||||
flash_dense_kwargs = {"causal": True},
|
||||
flash_varlen_kwargs = {"dropout_p": 0.0, "causal": True},
|
||||
)
|
||||
context = AttentionContext(
|
||||
bsz = bsz,
|
||||
q_len = q_len,
|
||||
kv_seq_len = kv_seq_len,
|
||||
n_heads = n_heads,
|
||||
head_dim = head_dim,
|
||||
requires_grad = hidden_states.requires_grad,
|
||||
seq_info = seq_info,
|
||||
attention_mask = attention_mask,
|
||||
causal_mask = causal_mask,
|
||||
)
|
||||
|
||||
# Group query attention
|
||||
if n_groups != 1:
|
||||
K = K.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
V = V.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
else:
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
A = xformers_attention(Q, K, V, attn_bias = causal_mask)
|
||||
A = A.view(bsz, q_len, n_heads, head_dim)
|
||||
|
||||
elif HAS_FLASH_ATTENTION and attention_mask is None:
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
A = flash_attn_func(Q, K, V, causal = True)
|
||||
else:
|
||||
# when qlen==vlen and attn_mask is None, we should use causal attention
|
||||
Q_len = Q.shape[-2]
|
||||
K_len = K.shape[-2]
|
||||
if attention_mask is None and Q_len == K_len:
|
||||
is_causal = True
|
||||
else:
|
||||
is_causal = False
|
||||
# Grouped query attention
|
||||
if SDPA_HAS_GQA:
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
A = scaled_dot_product_attention(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
attn_mask = attention_mask,
|
||||
is_causal = is_causal,
|
||||
enable_gqa = n_groups != 1,
|
||||
)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2) # .contiguous()
|
||||
else:
|
||||
if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(
|
||||
bsz, n_kv_heads, n_groups, kv_seq_len, head_dim
|
||||
)
|
||||
V = V[:, :, None, :, :].expand(
|
||||
bsz, n_kv_heads, n_groups, kv_seq_len, head_dim
|
||||
)
|
||||
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
pass
|
||||
# Must be contiguous or else results are False!
|
||||
# https://github.com/pytorch/pytorch/issues/112577
|
||||
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
A = scaled_dot_product_attention(
|
||||
Q, K, V, attn_mask = attention_mask, is_causal = is_causal
|
||||
)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2).contiguous()
|
||||
pass
|
||||
A = run_attention(config = config, context = context, Q = Q, K = K, V = V)
|
||||
attn_output = A.reshape(bsz, q_len, n_heads * head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
attn_weights = None
|
||||
|
|
@ -715,6 +672,7 @@ def LlamaDecoderLayer_fast_forward(
|
|||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
position_embeddings = position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states += residual
|
||||
|
||||
|
|
@ -738,6 +696,7 @@ def LlamaDecoderLayer_fast_forward(
|
|||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
position_embeddings = position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
|
|
@ -815,8 +774,11 @@ def LlamaModel_fast_forward(
|
|||
|
||||
seq_length_with_past = seq_length
|
||||
|
||||
# Fix out of bounds tokenization
|
||||
if hasattr(self, "max_seq_length"):
|
||||
# Fix out of bounds tokenization unless we were given packed metadata
|
||||
allow_overlength = getattr(self, "_unsloth_allow_packed_overlength", False) or (
|
||||
"packed_seq_lengths" in kwargs
|
||||
)
|
||||
if hasattr(self, "max_seq_length") and not allow_overlength:
|
||||
if seq_length > self.max_seq_length:
|
||||
shape = input_ids.shape if input_ids is not None else inputs_embeds.shape
|
||||
logger.warning_once(
|
||||
|
|
@ -1072,10 +1034,12 @@ def LlamaModel_fast_forward(
|
|||
|
||||
mask = causal_mask
|
||||
if IS_GEMMA2:
|
||||
if idx % 2 == 0:
|
||||
use_sliding_window = idx % 2 == 0
|
||||
if use_sliding_window:
|
||||
mask = self.SWA_mask if use_static_mask else dynamic_SWA_mask
|
||||
else:
|
||||
mask = self.GA_mask if use_static_mask else dynamic_GA_mask
|
||||
kwargs["use_sliding_window"] = use_sliding_window
|
||||
|
||||
if gradient_checkpointing and not isinstance(
|
||||
decoder_layer, GradientCheckpointingLayer
|
||||
|
|
@ -1089,6 +1053,7 @@ def LlamaModel_fast_forward(
|
|||
output_attentions,
|
||||
padding_mask = padding_mask,
|
||||
position_embeddings = position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return custom_forward
|
||||
|
|
@ -1115,6 +1080,7 @@ def LlamaModel_fast_forward(
|
|||
use_cache = use_cache,
|
||||
padding_mask = padding_mask,
|
||||
position_embeddings = position_embeddings,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = layer_outputs[0]
|
||||
|
||||
|
|
@ -1175,6 +1141,7 @@ def _LlamaModel_fast_forward_inference(
|
|||
past_key_values,
|
||||
position_ids,
|
||||
attention_mask = None,
|
||||
**kwargs,
|
||||
):
|
||||
input_ids = input_ids[:, : self.max_seq_length]
|
||||
bsz, q_len = input_ids.shape
|
||||
|
|
@ -1306,6 +1273,7 @@ def CausalLM_fast_forward(fast_forward_inference):
|
|||
past_key_values,
|
||||
position_ids = position_ids,
|
||||
attention_mask = attention_mask,
|
||||
**kwargs,
|
||||
)
|
||||
else:
|
||||
causal_mask = (
|
||||
|
|
@ -1338,6 +1306,7 @@ def CausalLM_fast_forward(fast_forward_inference):
|
|||
output_attentions = output_attentions,
|
||||
output_hidden_states = output_hidden_states,
|
||||
return_dict = return_dict,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = outputs[0]
|
||||
|
||||
|
|
@ -1446,6 +1415,10 @@ def CausalLM_fast_forward(fast_forward_inference):
|
|||
shift_labels = torch.empty_like(labels)
|
||||
shift_labels[..., :-1] = labels[..., 1:]
|
||||
shift_labels[..., -1] = -100
|
||||
mask_packed_sequence_boundaries(
|
||||
shift_labels,
|
||||
kwargs.get("packed_seq_lengths"),
|
||||
)
|
||||
# shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
|
||||
n_items = kwargs.get("num_items_in_batch", None)
|
||||
if n_items is None:
|
||||
|
|
@ -1533,7 +1506,7 @@ def PeftModel_fast_forward(
|
|||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# https://github.com/huggingface/transformers/pull/27931
|
||||
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
|
||||
class LlamaRotaryEmbedding(torch.nn.Module):
|
||||
|
|
@ -2010,8 +1983,8 @@ def unsloth_fast_generate(
|
|||
> self.config.max_position_embeddings
|
||||
):
|
||||
raise ValueError(
|
||||
f'Unsloth: input length {kwargs["input_ids"].shape[-1]} + max_new_tokens {kwargs["max_new_tokens"]} exceeds the maximum sequence length of {self.config.max_position_embeddings}!\n'
|
||||
'You will need to do long context extension by increasing the `max_seq_length` in `FastLanguageModel.from_pretrained`.'
|
||||
f"Unsloth: input length {kwargs['input_ids'].shape[-1]} + max_new_tokens {kwargs['max_new_tokens']} exceeds the maximum sequence length of {self.config.max_position_embeddings}!\n"
|
||||
"You will need to do long context extension by increasing the `max_seq_length` in `FastLanguageModel.from_pretrained`."
|
||||
)
|
||||
|
||||
# Must patch accelerate for Xformers
|
||||
|
|
|
|||
|
|
@ -12,8 +12,10 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from ..device_type import DEVICE_TYPE_TORCH
|
||||
import importlib
|
||||
import os
|
||||
import torch
|
||||
import re
|
||||
import tempfile
|
||||
from typing import Union
|
||||
|
|
@ -30,12 +32,14 @@ from packaging.version import Version
|
|||
from transformers import __version__ as transformers_version
|
||||
from unsloth.models._utils import TorchAOConfig
|
||||
from unsloth_zoo.utils import Version
|
||||
import torch
|
||||
import gc
|
||||
|
||||
transformers_version = Version(transformers_version)
|
||||
SUPPORTS_FOURBIT = transformers_version >= Version("4.37")
|
||||
|
||||
LOCAL_RANK_KEYS = ("LOCAL_RANK", "RANK")
|
||||
WORLD_SIZE_KEYS = ("WORLD_SIZE",)
|
||||
|
||||
BAD_MAPPINGS = {
|
||||
"unsloth/Qwen3-32B-unsloth-bnb-4bit".lower(): "unsloth/Qwen3-32B-bnb-4bit".lower(), # 32B dynamic quant is way too big
|
||||
"unsloth/Qwen3-30B-A3B-unsloth-bnb-4bit".lower(): "unsloth/Qwen3-30B-A3B".lower(), # HF loads MoEs too slowly
|
||||
|
|
@ -45,6 +49,50 @@ BAD_MAPPINGS = {
|
|||
}
|
||||
|
||||
|
||||
def _get_env_int(keys):
|
||||
for key in keys:
|
||||
value = os.environ.get(key)
|
||||
if value is None:
|
||||
continue
|
||||
try:
|
||||
return int(value)
|
||||
except ValueError:
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
def _infer_distributed_ranks():
|
||||
if torch.distributed.is_available() and torch.distributed.is_initialized():
|
||||
try:
|
||||
return torch.distributed.get_rank(), torch.distributed.get_world_size()
|
||||
except Exception:
|
||||
pass
|
||||
return _get_env_int(LOCAL_RANK_KEYS), _get_env_int(WORLD_SIZE_KEYS)
|
||||
|
||||
|
||||
def is_distributed():
|
||||
rank, world_size = _infer_distributed_ranks()
|
||||
return (world_size or 1) > 1 or (rank is not None and rank > 0)
|
||||
|
||||
|
||||
def prepare_device_map():
|
||||
rank, world_size = _infer_distributed_ranks()
|
||||
distributed = (world_size or 1) > 1 or (rank is not None and rank > 0)
|
||||
if not distributed:
|
||||
return None, False
|
||||
|
||||
local_rank = 0 if rank is None else rank
|
||||
device_map = {"": f"{DEVICE_TYPE_TORCH}:{local_rank}"}
|
||||
try:
|
||||
if DEVICE_TYPE_TORCH == "cuda":
|
||||
torch.cuda.set_device(local_rank)
|
||||
elif DEVICE_TYPE_TORCH == "xpu" and hasattr(torch, "xpu"):
|
||||
torch.xpu.set_device(local_rank)
|
||||
except Exception:
|
||||
pass
|
||||
return device_map, True
|
||||
|
||||
|
||||
def __get_model_name(
|
||||
model_name,
|
||||
load_in_4bit = True,
|
||||
|
|
|
|||
|
|
@ -17,6 +17,16 @@ import os
|
|||
from ._utils import __version__
|
||||
from unsloth_zoo.utils import _get_dtype
|
||||
from unsloth_zoo.hf_utils import dtype_from_config
|
||||
from ..utils.packing import (
|
||||
get_packed_info_from_kwargs,
|
||||
mask_packed_sequence_boundaries,
|
||||
)
|
||||
from ..utils.attention_dispatch import (
|
||||
AttentionConfig,
|
||||
AttentionContext,
|
||||
run_attention,
|
||||
select_attention_backend,
|
||||
)
|
||||
from .llama import (
|
||||
LlamaRotaryEmbedding,
|
||||
LlamaLinearScalingRotaryEmbedding,
|
||||
|
|
@ -76,6 +86,7 @@ def MistralAttention_fast_forward(
|
|||
Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
|
||||
K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
seq_info = get_packed_info_from_kwargs(kwargs, Q.device)
|
||||
|
||||
kv_seq_len = K.shape[-2]
|
||||
if past_key_value is not None:
|
||||
|
|
@ -83,12 +94,13 @@ def MistralAttention_fast_forward(
|
|||
|
||||
# Extend RoPE dynamically to fit in VRAM
|
||||
self.rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
|
||||
|
||||
cos, sin = self.rotary_emb.get_cached(kv_seq_len, Q.device.index)
|
||||
if position_ids is None:
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
else:
|
||||
Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids)
|
||||
|
||||
rope_position_ids = (
|
||||
position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
)
|
||||
# Useful for LongRoPE
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
|
||||
if past_key_value is not None:
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
|
|
@ -96,68 +108,38 @@ def MistralAttention_fast_forward(
|
|||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
if not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None:
|
||||
# Xformers memory efficient attention
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
K_M = V_M = bsz * kv_seq_len
|
||||
Q_M = bsz * q_len
|
||||
sw_cfg = getattr(self.config, "sliding_window", None)
|
||||
sw = kv_seq_len if (sw_cfg is None or sw_cfg == "null") else sw_cfg
|
||||
window_size = (-1, -1) if (kv_seq_len <= sw) else (sw, sw)
|
||||
|
||||
has_swa = isinstance(causal_mask, xformers.attn_bias.BlockDiagonalCausalMask)
|
||||
|
||||
# Group query attention
|
||||
K = K.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
V = V.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
|
||||
if has_swa:
|
||||
Q = Q.view(1, Q_M, n_heads, head_dim)
|
||||
K = K.view(1, K_M, n_heads, head_dim)
|
||||
V = V.view(1, V_M, n_heads, head_dim)
|
||||
else:
|
||||
# Xformers does support the forward pass though
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
|
||||
if has_swa:
|
||||
Q = Q.view(1, Q_M, n_kv_heads, n_groups, head_dim)
|
||||
K = K.view(1, K_M, n_kv_heads, n_groups, head_dim)
|
||||
V = V.view(1, V_M, n_kv_heads, n_groups, head_dim)
|
||||
|
||||
A = xformers_attention(Q, K, V, attn_bias = causal_mask)
|
||||
A = A.view(bsz, q_len, n_heads, head_dim)
|
||||
|
||||
elif HAS_FLASH_ATTENTION and attention_mask is None:
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
sw = getattr(self.config, "sliding_window", None)
|
||||
sw = kv_seq_len if (sw is None or sw == "null") else sw
|
||||
window = (-1, -1) if (kv_seq_len <= sw) else (sw, sw)
|
||||
A = flash_attn_func(Q, K, V, causal = True, window_size = window)
|
||||
else:
|
||||
# Grouped query attention
|
||||
# if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
# pass
|
||||
# Must be contiguous or else results are False!
|
||||
# https://github.com/pytorch/pytorch/issues/112577
|
||||
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
A = scaled_dot_product_attention(
|
||||
Q, K, V, attn_mask = attention_mask, is_causal = False
|
||||
)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2).contiguous()
|
||||
use_varlen = (
|
||||
seq_info is not None and past_key_value is None and window_size == (-1, -1)
|
||||
)
|
||||
backend = select_attention_backend(use_varlen)
|
||||
attention_config = AttentionConfig(
|
||||
backend = backend,
|
||||
n_kv_heads = n_kv_heads,
|
||||
n_groups = n_groups,
|
||||
flash_dense_kwargs = {"causal": True, "window_size": window_size},
|
||||
flash_varlen_kwargs = {
|
||||
"dropout_p": 0.0,
|
||||
"causal": True,
|
||||
"softmax_scale": getattr(self, "softmax_scale", None),
|
||||
},
|
||||
)
|
||||
context = AttentionContext(
|
||||
bsz = bsz,
|
||||
q_len = q_len,
|
||||
kv_seq_len = kv_seq_len,
|
||||
n_heads = n_heads,
|
||||
head_dim = head_dim,
|
||||
requires_grad = hidden_states.requires_grad,
|
||||
seq_info = seq_info,
|
||||
attention_mask = attention_mask,
|
||||
causal_mask = causal_mask,
|
||||
)
|
||||
|
||||
A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)
|
||||
attn_output = A.reshape(bsz, q_len, n_heads * head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
attn_weights = None
|
||||
|
|
@ -283,6 +265,7 @@ def MistralForCausalLM_fast_forward(
|
|||
output_attentions = output_attentions,
|
||||
output_hidden_states = output_hidden_states,
|
||||
return_dict = return_dict,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
|
|
@ -376,6 +359,10 @@ def MistralForCausalLM_fast_forward(
|
|||
shift_labels = torch.empty_like(labels)
|
||||
shift_labels[..., :-1] = labels[..., 1:]
|
||||
shift_labels[..., -1] = -100
|
||||
mask_packed_sequence_boundaries(
|
||||
shift_labels,
|
||||
kwargs.get("packed_seq_lengths"),
|
||||
)
|
||||
loss = fast_cross_entropy_loss(
|
||||
logits = shift_logits,
|
||||
labels = shift_labels,
|
||||
|
|
@ -439,7 +426,7 @@ class FastMistralModel(FastLlamaModel):
|
|||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# https://github.com/huggingface/transformers/pull/27931
|
||||
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
|
||||
import transformers.models.mistral.modeling_mistral
|
||||
|
|
|
|||
|
|
@ -60,7 +60,7 @@ class FastQwen2Model(FastLlamaModel):
|
|||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# https://github.com/huggingface/transformers/pull/27931
|
||||
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
|
||||
import transformers.models.qwen2.modeling_qwen2
|
||||
|
|
|
|||
|
|
@ -16,6 +16,13 @@ from .llama import *
|
|||
import os
|
||||
from ._utils import __version__
|
||||
from unsloth_zoo.utils import Version, _get_dtype
|
||||
from ..utils.packing import get_packed_info_from_kwargs
|
||||
from ..utils.attention_dispatch import (
|
||||
AttentionConfig,
|
||||
AttentionContext,
|
||||
run_attention,
|
||||
select_attention_backend,
|
||||
)
|
||||
from .llama import (
|
||||
LlamaRotaryEmbedding,
|
||||
LlamaLinearScalingRotaryEmbedding,
|
||||
|
|
@ -95,6 +102,7 @@ def Qwen3Attention_fast_forward(
|
|||
bsz, q_len, n_kv_heads, head_dim
|
||||
) # .transpose(1, 2) # we will transpose after normalisation
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
seq_info = get_packed_info_from_kwargs(kwargs, hidden_states.device)
|
||||
|
||||
# Qwen3 has QKNorm. This seems to be the only difference from Qwen2.
|
||||
# Note that using fast_layernorm_compiled causes issues as the dimensions don't match up.
|
||||
|
|
@ -110,20 +118,19 @@ def Qwen3Attention_fast_forward(
|
|||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
|
||||
# Extend RoPE dynamically to fit in VRAM
|
||||
if position_embeddings and kv_seq_len <= position_embeddings[0].shape[0]:
|
||||
cos, sin = position_embeddings
|
||||
else:
|
||||
# Extend RoPE dynamically to fit in VRA
|
||||
rotary_emb = self.rotary_emb
|
||||
rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
|
||||
device_index = Q.device.index
|
||||
cos, sin = rotary_emb.get_cached(kv_seq_len, Q.device.index)
|
||||
|
||||
if position_ids is None:
|
||||
# Useful for LongRoPE
|
||||
cos, sin = rotary_emb.get_cached(kv_seq_len, device_index)
|
||||
else:
|
||||
cos, sin = rotary_emb.get_cached(kv_seq_len, device_index)
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
rope_position_ids = (
|
||||
position_ids if position_ids is not None else kwargs.get("position_ids")
|
||||
)
|
||||
# Useful for LongRoPE
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids)
|
||||
|
||||
if past_key_value is not None:
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
|
|
@ -131,74 +138,32 @@ def Qwen3Attention_fast_forward(
|
|||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
# Attention module
|
||||
if not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None:
|
||||
# Xformers memory efficient attention
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
K_M = V_M = bsz * kv_seq_len
|
||||
Q_M = bsz * q_len
|
||||
use_varlen = seq_info is not None and past_key_value is None
|
||||
backend = select_attention_backend(use_varlen)
|
||||
attention_config = AttentionConfig(
|
||||
backend = backend,
|
||||
n_kv_heads = n_kv_heads,
|
||||
n_groups = n_groups,
|
||||
flash_dense_kwargs = {"causal": True},
|
||||
flash_varlen_kwargs = {
|
||||
"dropout_p": 0.0,
|
||||
"causal": True,
|
||||
"softmax_scale": getattr(self, "softmax_scale", None),
|
||||
},
|
||||
)
|
||||
context = AttentionContext(
|
||||
bsz = bsz,
|
||||
q_len = q_len,
|
||||
kv_seq_len = kv_seq_len,
|
||||
n_heads = n_heads,
|
||||
head_dim = head_dim,
|
||||
requires_grad = hidden_states.requires_grad,
|
||||
seq_info = seq_info,
|
||||
attention_mask = attention_mask,
|
||||
causal_mask = causal_mask,
|
||||
)
|
||||
|
||||
has_swa = isinstance(causal_mask, xformers.attn_bias.BlockDiagonalCausalMask)
|
||||
|
||||
# Group query attention
|
||||
K = K.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
V = V.view(bsz, kv_seq_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, kv_seq_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
K = K.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
V = V.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
|
||||
if has_swa:
|
||||
Q = Q.view(1, Q_M, n_heads, head_dim)
|
||||
K = K.view(1, K_M, n_heads, head_dim)
|
||||
V = V.view(1, V_M, n_heads, head_dim)
|
||||
else:
|
||||
# Xformers does support the forward pass though
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
|
||||
if has_swa:
|
||||
Q = Q.view(1, Q_M, n_kv_heads, n_groups, head_dim)
|
||||
K = K.view(1, K_M, n_kv_heads, n_groups, head_dim)
|
||||
V = V.view(1, V_M, n_kv_heads, n_groups, head_dim)
|
||||
|
||||
A = xformers_attention(Q, K, V, attn_bias = causal_mask)
|
||||
A = A.view(bsz, q_len, n_heads, head_dim)
|
||||
|
||||
elif HAS_FLASH_ATTENTION and attention_mask is None:
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
sw = kv_seq_len
|
||||
window = (-1, -1) if (kv_seq_len <= sw) else (sw, sw)
|
||||
A = flash_attn_func(Q, K, V, causal = True, window_size = window)
|
||||
else:
|
||||
# Grouped query attention
|
||||
# if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)
|
||||
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
# pass
|
||||
# Must be contiguous or else results are False!
|
||||
# https://github.com/pytorch/pytorch/issues/112577
|
||||
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
# when qlen==vlen and attn_mask is None, we should use causal attention
|
||||
Q_len = Q.shape[-2]
|
||||
K_len = K.shape[-2]
|
||||
if attention_mask is None and Q_len == K_len:
|
||||
is_causal = True
|
||||
else:
|
||||
is_causal = False
|
||||
|
||||
A = scaled_dot_product_attention(
|
||||
Q, K, V, attn_mask = attention_mask, is_causal = is_causal
|
||||
)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2).contiguous()
|
||||
A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)
|
||||
|
||||
attn_output = A.reshape(bsz, q_len, n_heads * head_dim)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
|
|
@ -451,7 +416,7 @@ class FastQwen3Model(FastLlamaModel):
|
|||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# https://github.com/huggingface/transformers/pull/27931
|
||||
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
|
||||
import transformers.models.qwen3.modeling_qwen3
|
||||
|
|
|
|||
|
|
@ -59,12 +59,14 @@ def Qwen3MoeSparseMoeBlock_fast_forward(self, X, temp_gate = None, temp_up = Non
|
|||
self.gate_proj, X, out = temp_gate
|
||||
) # pretty much the only change from transformers implementation.
|
||||
|
||||
routing_weights = torch_nn_functional_softmax(router_logits, dim = -1)
|
||||
routing_weights = torch_nn_functional_softmax(
|
||||
router_logits, dim = -1, dtype = torch.float32
|
||||
)
|
||||
routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim = -1)
|
||||
routing_weights /= routing_weights.sum(dim = -1, keepdim = True)
|
||||
# we cast back to the input dtype
|
||||
routing_weights = routing_weights.to(X.dtype)
|
||||
final_X = torch.zeros((bsz * seq_len, hd), dtype = X.dtype, device = X.device)
|
||||
final_X = torch.zeros((bsz * seq_len, hd), dtype = torch.float32, device = X.device)
|
||||
|
||||
# One hot encode the selected experts to create an expert mask
|
||||
# this will be used to easily index which expert is going to be sollicitated
|
||||
|
|
@ -128,7 +130,7 @@ def Qwen3MoeDecoderLayer_fast_forward(
|
|||
position_embeddings = position_embeddings,
|
||||
_flag_for_generation = self._flag_for_generation,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
hidden_states += residual
|
||||
|
||||
# MoE Router MLP
|
||||
residual = hidden_states
|
||||
|
|
@ -138,7 +140,7 @@ def Qwen3MoeDecoderLayer_fast_forward(
|
|||
hidden_states, router_logits = Qwen3MoeSparseMoeBlock_fast_forward(
|
||||
self.mlp, hidden_states
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
hidden_states += residual
|
||||
else:
|
||||
residual = hidden_states
|
||||
hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states)
|
||||
|
|
@ -200,7 +202,7 @@ class FastQwen3MoeModel(FastQwen3Model):
|
|||
|
||||
# Solves https://github.com/unslothai/unsloth/issues/168
|
||||
# Static KV Cache was introduced in 4.38.0, causing training to be much slower.
|
||||
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# Inference can now be CUDAGraphed, but we shall retain the old rotary embeddings.
|
||||
# https://github.com/huggingface/transformers/pull/27931
|
||||
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py\
|
||||
import transformers.models.qwen3_moe.modeling_qwen3_moe
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ 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
|
||||
from unsloth_zoo.rl_replacements import RL_REPLACEMENTS
|
||||
from .rl_replacements import (
|
||||
|
|
@ -32,6 +33,7 @@ from .rl_replacements import (
|
|||
RL_PRE_ITEMS,
|
||||
RL_CONFIG_CHANGES,
|
||||
RL_METRICS_CHANGES,
|
||||
RL_ADDITIONAL_FUNCTIONS,
|
||||
)
|
||||
|
||||
torch_compile_options = {
|
||||
|
|
@ -312,7 +314,7 @@ class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}):
|
|||
if getattr(args, "_n_gpu", 1) != 1:
|
||||
args._n_gpu = 1
|
||||
if "model" in locals() and hasattr(model, "for_training"):
|
||||
model.for_training()
|
||||
model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True))
|
||||
super().__init__({RLTrainer_call_args}{RLTrainer_kwargs})
|
||||
if "model" in locals() and hasattr(model, "for_inference"):
|
||||
model.for_inference()
|
||||
|
|
@ -329,6 +331,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
try:
|
||||
trainer = eval(f"trl.trainer.{trainer_file}")
|
||||
except Exception as error:
|
||||
print(f"Unsloth: Could not import trl.trainer.{trainer_file}: {error}")
|
||||
return
|
||||
|
||||
# Get SFTTrainer and SFTConfig names
|
||||
|
|
@ -347,8 +350,14 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
and trainer_file.split("_")[0] in x.lower()
|
||||
]
|
||||
if len(name) != 1:
|
||||
logger.info(
|
||||
f"Unsloth: Could not find Trainer class in trl.trainer.{trainer_file}. Found: {name}"
|
||||
)
|
||||
return
|
||||
if len(config) != 1:
|
||||
logger.info(
|
||||
f"Unsloth: Could not find Config class in trl.trainer.{trainer_file}. Found: {config}"
|
||||
)
|
||||
return
|
||||
|
||||
# Get SFTTrainer, SFTConfig
|
||||
|
|
@ -356,17 +365,25 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
RLConfig_name = config[0]
|
||||
try:
|
||||
RLTrainer = eval(f"trl.trainer.{trainer_file}.{RLTrainer_name}")
|
||||
except:
|
||||
except Exception as e:
|
||||
logger.info(
|
||||
f"Unsloth: Could not load {RLTrainer_name} from trl.trainer.{trainer_file}: {e}"
|
||||
)
|
||||
return
|
||||
try:
|
||||
RLConfig = eval(f"trl.trainer.{trainer_file}.{RLConfig_name}")
|
||||
except:
|
||||
except Exception as e:
|
||||
logger.info(
|
||||
f"Unsloth: Could not load {RLConfig_name} from trl.trainer.{trainer_file}: {e}"
|
||||
)
|
||||
return
|
||||
|
||||
# Check name
|
||||
if RLTrainer.__name__.startswith("Unsloth"):
|
||||
print(f"Unsloth: {RLTrainer.__name__} is already patched.")
|
||||
return
|
||||
if RLConfig.__name__.startswith("Unsloth"):
|
||||
print(f"Unsloth: {RLConfig.__name__} is already patched.")
|
||||
return
|
||||
|
||||
# Get old source
|
||||
|
|
@ -442,7 +459,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
" force_float32 = True\n"
|
||||
"mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32')\n"
|
||||
"dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None)\n"
|
||||
"if dtype is None: dtype = model.get_input_embeddings().dtype\n"
|
||||
"if dtype is None: dtype = model.get_input_embeddings().weight.dtype\n"
|
||||
"from unsloth_zoo.utils import _get_dtype\n"
|
||||
"dtype = _get_dtype(dtype)\n"
|
||||
"float16 = dtype == torch.float16\n"
|
||||
|
|
@ -453,13 +470,24 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
" args.fp16 = False\n"
|
||||
" args.bf16 = False\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n"
|
||||
" # args.mixed_precision is a new argument which needs to be set now\n"
|
||||
"elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':\n"
|
||||
" # Mixed precision training\n"
|
||||
" args.fp16 = float16\n"
|
||||
" args.bf16 = not float16\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16'\n"
|
||||
" # args.mixed_precision is a new argument which needs to be set now\n"
|
||||
"elif mixed_precision_dtype == 'bfloat16':\n"
|
||||
" # Both False since bfloat16 full finetuning doesn't do any autocasting.\n"
|
||||
" args.fp16 = False\n"
|
||||
" args.bf16 = False\n"
|
||||
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n"
|
||||
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n"
|
||||
" # args.mixed_precision is a new argument which needs to be set now\n"
|
||||
"\n"
|
||||
)
|
||||
"elif mixed_precision_dtype == 'bfloat16':\n" " # Both False since bfloat16 full finetuning doesn't do any autocasting.\n" " args.fp16 = False\n" " args.bf16 = False\n" " os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n"
|
||||
extra_args += mixed_precision
|
||||
|
||||
# Check if per_device_eval_batch_size (default 8) bigger than bsz
|
||||
|
|
@ -565,7 +593,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
if "model" in call_args:
|
||||
training_check = (
|
||||
"if model is not None and hasattr(model, 'for_training'):\n"
|
||||
" model.for_training()\n"
|
||||
" model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True))\n"
|
||||
"if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right'\n"
|
||||
"if 'processing_class' in locals():\n"
|
||||
" if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right'\n"
|
||||
|
|
@ -977,7 +1005,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
|
|||
RLTrainer_source,
|
||||
f"trl.trainer.{trainer_file}",
|
||||
imports,
|
||||
overwrite = True,
|
||||
overwrite = False,
|
||||
)
|
||||
|
||||
# Patch Trainer
|
||||
|
|
@ -1242,6 +1270,11 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
|
|||
+ r", load_tensors = True))",
|
||||
source,
|
||||
)
|
||||
# All these are to fix multiple commas before lora_request (in case the original code ends with something like ",)")
|
||||
# https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L1388 for eg has such an ending
|
||||
source = re.sub(r"\,[\s]{1,}\,[\s]{0,}lora_request", ", lora_request", source)
|
||||
source = re.sub(r"[\s]{1,}\,[\s]{0,}lora_request", ", lora_request", source)
|
||||
source = re.sub(r"[\,]{1,}[\s]{0,}lora_request", ", lora_request", source)
|
||||
# Prefer using unsloth's sampling params and fallback to trl's if not found
|
||||
# We'll enable this later separately when combining both this and GRPOConfig params
|
||||
# source = re.sub(
|
||||
|
|
@ -1291,15 +1324,27 @@ def patch_trl_rl_trainers():
|
|||
import trl.trainer
|
||||
|
||||
all_trainers = dir(trl.trainer)
|
||||
all_trainers = [x for x in all_trainers if x.islower() and x.endswith("_trainer")]
|
||||
all_trainers = [
|
||||
x
|
||||
for x in all_trainers
|
||||
if x.islower() and x.endswith("_trainer") and x != "base_trainer"
|
||||
]
|
||||
for trainer in all_trainers:
|
||||
_patch_trl_rl_trainers(trainer)
|
||||
return
|
||||
|
||||
|
||||
def patch_trl_openenv():
|
||||
for function in RL_ADDITIONAL_FUNCTIONS["openenv"]:
|
||||
logger.info(f"Unsloth: Patching trl openenv with function: {function.__name__}")
|
||||
function() # Call the function to apply the patch
|
||||
return
|
||||
|
||||
|
||||
def PatchFastRL(algorithm = None, FastLanguageModel = None):
|
||||
if FastLanguageModel is not None:
|
||||
PatchRL(FastLanguageModel)
|
||||
patch_trl_rl_trainers()
|
||||
patch_trl_openenv()
|
||||
if type(algorithm) is str and algorithm.islower():
|
||||
PatchRLStatistics(algorithm)
|
||||
|
|
|
|||
|
|
@ -26,6 +26,10 @@ import torch
|
|||
import inspect
|
||||
from collections import defaultdict
|
||||
from unsloth_zoo.rl_replacements import RL_REPLACEMENTS, left_pack_padding
|
||||
from unsloth_zoo.utils import Version
|
||||
from importlib.metadata import version as importlib_version
|
||||
from unsloth_zoo.log import logger
|
||||
import importlib.util
|
||||
from ..device_type import (
|
||||
is_hip,
|
||||
get_device_type,
|
||||
|
|
@ -42,6 +46,7 @@ RL_FUNCTIONS = defaultdict(list)
|
|||
RL_PRE_ITEMS = defaultdict(list)
|
||||
RL_CONFIG_CHANGES = defaultdict(list)
|
||||
RL_METRICS_CHANGES = defaultdict(list)
|
||||
RL_ADDITIONAL_FUNCTIONS = defaultdict(list)
|
||||
|
||||
torch_compile_options = {
|
||||
"epilogue_fusion": True,
|
||||
|
|
@ -216,6 +221,27 @@ def grpo_trainer__prepare_inputs(function_name, function):
|
|||
RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__prepare_inputs)
|
||||
|
||||
|
||||
# Remove collective RPC of reload weights from generate
|
||||
# trl added reload weights (potentially for quantized models), we don't need it for our use case (LoRA primarily)
|
||||
# https://github.com/huggingface/trl/commit/7856d3b1f6518601732f489883b341bb6dd36434#diff-964e6fd373aa93037604064cb2b822d7f8e2735e33f791065acf2c4c3552d393R1168-R1169
|
||||
def grpo_trainer__generate_single_turn(function_name, function):
|
||||
if function_name != "_generate_single_turn":
|
||||
return function
|
||||
|
||||
# Remove the reload_weights collective RPC call from the generate function's source
|
||||
# function = function.replace('self.llm.collective_rpc("reload_weights")', "")
|
||||
# The regex below does the same thing but is more flexible and can handle single or double quotes
|
||||
function = re.sub(
|
||||
r"self\.llm\.collective_rpc\(\s*(['\"])reload_weights\1\s*\)",
|
||||
"",
|
||||
function,
|
||||
)
|
||||
return function
|
||||
|
||||
|
||||
RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__generate_single_turn)
|
||||
|
||||
|
||||
# Fix incorrect special tokens handling and truncation in older TRL versions
|
||||
def grpo_trainer__generate_and_score_completions(function_name, function):
|
||||
if function_name != "_generate_and_score_completions":
|
||||
|
|
@ -906,3 +932,54 @@ def grpo_trainer_metrics(RLTrainer_source, RLConfig_source):
|
|||
|
||||
|
||||
RL_METRICS_CHANGES["grpo_trainer"].append(grpo_trainer_metrics)
|
||||
|
||||
|
||||
def openenv_vllm_reload_weights():
|
||||
# This function patches the trl openenv generate_rollout_completions function to:
|
||||
# 1. Remove the reload_weights call (unsloth handles weight reloading)
|
||||
# 2. Fix wake_up call to be compatible with unsloth (remove tags to wake everything)
|
||||
#
|
||||
# The issue: TRL's wake_up(tags=["kv_cache"]) only wakes kv_cache, leaving is_sleeping=True
|
||||
# at the executor level. This causes unsloth's patched generate to try waking up again,
|
||||
# resulting in double create_and_map on already-mapped handles.
|
||||
#
|
||||
# The fix: Use wake_up() with no tags, which wakes everything. Unsloth's patched
|
||||
# CuMemAllocator.wake_up skips weights anyway, so this is safe.
|
||||
if importlib.util.find_spec("trl") is None:
|
||||
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}")
|
||||
return
|
||||
|
||||
src = inspect.getsource(openenv_utils.generate_rollout_completions)
|
||||
src = textwrap.dedent(src)
|
||||
original_src = src
|
||||
|
||||
# Remove the reload_weights call - unsloth handles this differently
|
||||
src = re.sub(r'.*\.collective_rpc\("reload_weights"\).*\n?', "", src)
|
||||
|
||||
# Change wake_up(tags=["kv_cache"]) to wake_up() - wake everything to set is_sleeping=False
|
||||
# This prevents double wake_up issues. Unsloth's allocator skips weights anyway.
|
||||
src = re.sub(r"\.wake_up\(tags=\[.*?\]\)", ".wake_up()", src)
|
||||
|
||||
if original_src == src:
|
||||
logger.warning("Unsloth: Warning - regex did not match, patch may have failed")
|
||||
return
|
||||
|
||||
# Execute and explicitly assign to module
|
||||
local_ns = {}
|
||||
exec(compile(src, "<unsloth>", "exec"), openenv_utils.__dict__, local_ns)
|
||||
patched_func = local_ns["generate_rollout_completions"]
|
||||
|
||||
# Patch both the utils module and the parent openenv module
|
||||
openenv_utils.generate_rollout_completions = patched_func
|
||||
openenv.generate_rollout_completions = patched_func
|
||||
logger.info("Unsloth: Patched trl openenv generate_rollout_completions")
|
||||
|
||||
|
||||
RL_ADDITIONAL_FUNCTIONS["openenv"].append(openenv_vllm_reload_weights)
|
||||
|
|
|
|||
|
|
@ -430,8 +430,10 @@ class FastBaseModel:
|
|||
|
||||
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
|
||||
|
||||
arch_name = model_type_arch.title()
|
||||
arch_name = arch_name.replace("_Vl_", "_VL_").replace("_Moe", "_MoE")
|
||||
statistics = (
|
||||
f"==((====))== Unsloth {__version__}: Fast {model_type_arch.title()} patching. Transformers: {transformers_version}.{vllm_version}\n"
|
||||
f"==((====))== Unsloth {__version__}: Fast {arch_name} patching. Transformers: {transformers_version}.{vllm_version}\n"
|
||||
f" {chr(92)}{chr(92)} /| {gpu_stats_name}Num GPUs = {DEVICE_COUNT}. Max memory: {max_memory} GB. Platform: {platform_system}.\n"
|
||||
f"O^O/ {chr(92)}_/ {chr(92)} Torch: {torch.__version__}. {gpu_stats_snippet} Triton: {triton_version}\n"
|
||||
f"{chr(92)} / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. FA [Xformers = {xformers_version}. FA2 = {HAS_FLASH_ATTENTION}]\n"
|
||||
|
|
@ -670,25 +672,34 @@ class FastBaseModel:
|
|||
model.fast_generate = model.generate
|
||||
model.fast_generate_batches = error_out_no_vllm
|
||||
if offload_embedding:
|
||||
embed_tokens = model.get_input_embeddings()
|
||||
nbytes = embed_tokens.weight.numel() * embed_tokens.weight.itemsize
|
||||
ngb = round(nbytes / 1024 / 1024 / 1024, 2)
|
||||
print(f"Unsloth: Offloading embeddings to RAM to save {ngb} GB.")
|
||||
embed_tokens.to("cpu")
|
||||
if bool(
|
||||
os.environ.get("WSL_DISTRO_NAME") or os.environ.get("WSL_INTEROP")
|
||||
):
|
||||
# WSL doesn't work with offloaded embeddings
|
||||
pass
|
||||
elif os.name == "nt":
|
||||
# Windows doesn't work with offloaded embeddings
|
||||
pass
|
||||
else:
|
||||
embed_tokens = model.get_input_embeddings()
|
||||
nbytes = embed_tokens.weight.numel() * embed_tokens.weight.itemsize
|
||||
ngb = round(nbytes / 1024 / 1024 / 1024, 2)
|
||||
print(f"Unsloth: Offloading embeddings to RAM to save {ngb} GB.")
|
||||
embed_tokens.to("cpu")
|
||||
|
||||
# Add hooks to move inputs to CPU and back to CUDA
|
||||
# [TODO] Doesn't seem to work!
|
||||
# def pre_hook(module, args):
|
||||
# args[0]._old_device = args[0].device
|
||||
# return (args[0].to("cpu", non_blocking = True))
|
||||
# def post_hook(module, args, output):
|
||||
# old_device = getattr(args[0], "_old_device", "cuda")
|
||||
# return output.to(old_device, non_blocking = True)
|
||||
# embed_tokens.register_forward_pre_hook(pre_hook, prepend = True)
|
||||
# embed_tokens.register_forward_hook (post_hook, prepend = True)
|
||||
# Must free GPU memory otherwise will not free!
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
# Add hooks to move inputs to CPU and back to CUDA
|
||||
# [TODO] Doesn't seem to work!
|
||||
# def pre_hook(module, args):
|
||||
# args[0]._old_device = args[0].device
|
||||
# return (args[0].to("cpu", non_blocking = True))
|
||||
# def post_hook(module, args, output):
|
||||
# old_device = getattr(args[0], "_old_device", "cuda")
|
||||
# return output.to(old_device, non_blocking = True)
|
||||
# embed_tokens.register_forward_pre_hook(pre_hook, prepend = True)
|
||||
# embed_tokens.register_forward_hook (post_hook, prepend = True)
|
||||
# Must free GPU memory otherwise will not free!
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
else:
|
||||
from unsloth_zoo.vllm_utils import (
|
||||
load_vllm,
|
||||
|
|
|
|||
152
unsloth/save.py
152
unsloth/save.py
|
|
@ -42,6 +42,7 @@ import re
|
|||
from transformers.models.llama.modeling_llama import logger
|
||||
from .tokenizer_utils import fix_sentencepiece_gguf
|
||||
from .models.loader_utils import get_model_name
|
||||
from .models._utils import _convert_torchao_model
|
||||
from .ollama_template_mappers import OLLAMA_TEMPLATES, MODEL_TO_OLLAMA_TEMPLATE_MAPPER
|
||||
from transformers import ProcessorMixin
|
||||
from huggingface_hub import HfApi
|
||||
|
|
@ -2734,11 +2735,35 @@ def unsloth_generic_push_to_hub_merged(
|
|||
gc.collect()
|
||||
|
||||
|
||||
def unsloth_save_pretrained_torchao(
|
||||
self,
|
||||
def _unsloth_save_torchao_with_attached_config(
|
||||
model,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
tokenizer = None,
|
||||
torchao_config = None,
|
||||
tokenizer,
|
||||
push_to_hub: bool = False,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
):
|
||||
"""Save a QAT-trained model by converting fake-quantized weights to real quantized weights."""
|
||||
# Convert QAT fake-quantized weights to real quantized weights
|
||||
_convert_torchao_model(model)
|
||||
|
||||
# TorchAO does not support safe_serialization reliably
|
||||
safe_serialization = False
|
||||
|
||||
if push_to_hub:
|
||||
model.push_to_hub(
|
||||
save_directory, safe_serialization = safe_serialization, token = token
|
||||
)
|
||||
tokenizer.push_to_hub(save_directory, token = token)
|
||||
else:
|
||||
model.save_pretrained(save_directory, safe_serialization = safe_serialization)
|
||||
tokenizer.save_pretrained(save_directory)
|
||||
|
||||
|
||||
def _unsloth_save_torchao_with_given_config(
|
||||
model,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
tokenizer,
|
||||
torchao_config,
|
||||
push_to_hub: bool = False,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
):
|
||||
|
|
@ -2749,23 +2774,26 @@ def unsloth_save_pretrained_torchao(
|
|||
`torchao_config` (TorchAOBaseConfig): configuration for torchao quantization, full list: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
|
||||
`push_to_hub` (bool): whether to push the checkpoint to huggingface hub or save locally
|
||||
"""
|
||||
|
||||
if push_to_hub:
|
||||
assert token is not None, "Unsloth: Please specify a token for uploading!"
|
||||
|
||||
assert (
|
||||
torchao_config is not None
|
||||
), "Unsloth: Please specify a torchao_config for post-training quantization!"
|
||||
|
||||
# first merge the lora weights
|
||||
arguments = dict(locals())
|
||||
arguments["model"] = self
|
||||
arguments["tokenizer"] = tokenizer
|
||||
arguments["push_to_hub"] = False # We save ourselves
|
||||
arguments["save_method"] = "merged_16bit" # Must be 16bit
|
||||
del arguments["self"]
|
||||
del arguments["torchao_config"]
|
||||
|
||||
if token is None and push_to_hub:
|
||||
token = get_token()
|
||||
|
||||
if not isinstance(self, PeftModelForCausalLM) and not isinstance(self, PeftModel):
|
||||
self.save_pretrained(save_directory)
|
||||
if not isinstance(model, PeftModelForCausalLM) and not isinstance(model, PeftModel):
|
||||
model.save_pretrained(save_directory)
|
||||
tokenizer.save_pretrained(save_directory)
|
||||
else:
|
||||
unsloth_generic_save(**arguments)
|
||||
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
|
||||
|
|
@ -2778,26 +2806,20 @@ def unsloth_save_pretrained_torchao(
|
|||
)
|
||||
from torchao import quantize_
|
||||
|
||||
if torchao_config is None:
|
||||
from torchao.quantization import Int8DynamicActivationInt8WeightConfig
|
||||
|
||||
print(
|
||||
"Unsloth: You did not specify a `torchao_config`, so defaulting to `Int8DynamicActivationInt8WeightConfig`"
|
||||
)
|
||||
torchao_config = Int8DynamicActivationInt8WeightConfig()
|
||||
quantization_config = TorchAoConfig(quant_type = torchao_config)
|
||||
|
||||
# Determine if this is a VLM
|
||||
is_vlm = False
|
||||
if hasattr(self, "config") and hasattr(self.config, "architectures"):
|
||||
if hasattr(model, "config") and hasattr(model.config, "architectures"):
|
||||
is_vlm = any(
|
||||
x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
|
||||
for x in self.config.architectures
|
||||
for x in model.config.architectures
|
||||
)
|
||||
is_vlm = is_vlm or hasattr(self.config, "vision_config")
|
||||
is_vlm = is_vlm or hasattr(model.config, "vision_config")
|
||||
auto_model = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM
|
||||
auto_processor = AutoProcessor if is_vlm else AutoTokenizer
|
||||
|
||||
tokenizer = auto_processor.from_pretrained(arguments["save_directory"])
|
||||
tokenizer = auto_processor.from_pretrained(save_directory)
|
||||
|
||||
# TorchAO must only use bfloat16 for loading (float16 fails)
|
||||
if HAS_TORCH_DTYPE:
|
||||
|
|
@ -2805,8 +2827,9 @@ def unsloth_save_pretrained_torchao(
|
|||
else:
|
||||
kwargs = {"dtype": torch.bfloat16}
|
||||
|
||||
model = auto_model.from_pretrained(
|
||||
arguments["save_directory"],
|
||||
# Reload with quantization applied
|
||||
quantized_model = auto_model.from_pretrained(
|
||||
save_directory,
|
||||
device_map = "auto",
|
||||
quantization_config = quantization_config,
|
||||
**kwargs,
|
||||
|
|
@ -2817,25 +2840,92 @@ def unsloth_save_pretrained_torchao(
|
|||
# TorchAO does not support safe_serialization right now 0.14.0 seems broken!
|
||||
safe_serialization = Version(importlib_version("torchao")) > Version("0.14.0")
|
||||
safe_serialization = False
|
||||
|
||||
if push_to_hub:
|
||||
if token is None and push_to_hub:
|
||||
token = get_token()
|
||||
model.push_to_hub(
|
||||
quantized_model.push_to_hub(
|
||||
torchao_save_directory, safe_serialization = safe_serialization, token = token
|
||||
)
|
||||
tokenizer.push_to_hub(torchao_save_directory, token = token)
|
||||
else:
|
||||
model.save_pretrained(
|
||||
quantized_model.save_pretrained(
|
||||
torchao_save_directory, safe_serialization = safe_serialization
|
||||
)
|
||||
tokenizer.save_pretrained(torchao_save_directory)
|
||||
|
||||
# Clean up the intermediate unquantized model
|
||||
if os.path.exists(save_directory):
|
||||
try:
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(save_directory)
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
def unsloth_save_pretrained_torchao(
|
||||
self,
|
||||
save_directory: Union[str, os.PathLike],
|
||||
tokenizer = None,
|
||||
torchao_config = None,
|
||||
push_to_hub: bool = False,
|
||||
token: Optional[Union[str, bool]] = None,
|
||||
):
|
||||
"""Saves a torchao quantized model checkpoint.
|
||||
|
||||
This function handles two mutually exclusive workflows:
|
||||
|
||||
1. **QAT (Quantization-Aware Training)**: If the model was trained with `qat_scheme`
|
||||
parameter, do NOT pass `torchao_config`. The function will convert the QAT
|
||||
fake-quantized weights to real quantized weights and save directly.
|
||||
|
||||
2. **PTQ (Post-Training Quantization)**: If you want to apply quantization to a
|
||||
regular model, pass a `torchao_config`. The model must NOT have been trained
|
||||
with `qat_scheme`.
|
||||
|
||||
Args:
|
||||
`save_directory`: local folder path or huggingface hub ID when `push_to_hub` is True
|
||||
`tokenizer`: the tokenizer to save alongside the model
|
||||
`torchao_config` (TorchAOBaseConfig): configuration for torchao quantization.
|
||||
Required for PTQ, must be None for QAT models.
|
||||
Options: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize
|
||||
`push_to_hub` (bool): whether to push to huggingface hub or save locally
|
||||
`token`: HuggingFace token for pushing to hub
|
||||
"""
|
||||
if token is None and push_to_hub:
|
||||
token = get_token()
|
||||
|
||||
has_qat_config = (
|
||||
hasattr(self, "_torchao_config") and self._torchao_config is not None
|
||||
)
|
||||
|
||||
if torchao_config is not None:
|
||||
# PTQ path: user provided a config, model must NOT have QAT config
|
||||
assert not has_qat_config, (
|
||||
"Unsloth: You passed `torchao_config` but this model was trained with `qat_scheme`. "
|
||||
"For QAT models, do not pass `torchao_config` - the quantization config is already "
|
||||
"attached to the model from training."
|
||||
)
|
||||
_unsloth_save_torchao_with_given_config(
|
||||
model = self,
|
||||
save_directory = save_directory,
|
||||
tokenizer = tokenizer,
|
||||
torchao_config = torchao_config,
|
||||
push_to_hub = push_to_hub,
|
||||
token = token,
|
||||
)
|
||||
else:
|
||||
# QAT path: no config provided, model must have QAT config
|
||||
assert has_qat_config, (
|
||||
"Unsloth: No `torchao_config` provided and model was not trained with `qat_scheme`. "
|
||||
"Either train with `qat_scheme` parameter, or provide a `torchao_config` for "
|
||||
"post-training quantization."
|
||||
)
|
||||
_unsloth_save_torchao_with_attached_config(
|
||||
model = self,
|
||||
save_directory = save_directory,
|
||||
tokenizer = tokenizer,
|
||||
push_to_hub = push_to_hub,
|
||||
token = token,
|
||||
)
|
||||
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,8 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
|
@ -21,13 +23,20 @@ import trl
|
|||
import inspect
|
||||
from trl import SFTTrainer
|
||||
from . import is_bfloat16_supported
|
||||
from unsloth.utils import (
|
||||
configure_padding_free,
|
||||
configure_sample_packing,
|
||||
enable_padding_free_metadata,
|
||||
enable_sample_packing,
|
||||
)
|
||||
from unsloth_zoo.training_utils import (
|
||||
unsloth_train as _unsloth_train,
|
||||
)
|
||||
from unsloth_zoo.vision_utils import (
|
||||
UnslothVisionDataCollator,
|
||||
)
|
||||
from packaging.version import Version
|
||||
from unsloth_zoo.hf_utils import get_transformers_model_type
|
||||
from unsloth_zoo.utils import Version
|
||||
import dataclasses
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -38,8 +47,59 @@ __all__ = [
|
|||
"UnslothVisionDataCollator",
|
||||
]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_AUTO_PADDING_FREE_ENV_DISABLED = os.environ.get(
|
||||
"UNSLOTH_DISABLE_AUTO_PADDING_FREE", ""
|
||||
).strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
PADDING_FREE_BLOCKLIST = {
|
||||
"gemma2", # - gemma2: Uses slow_attention_softcapping which has torch.compile issues
|
||||
"gpt_oss", # - gpt_oss: Uses Flex Attention which doesn't handle padding_free correctly
|
||||
}
|
||||
|
||||
|
||||
def _should_pack(config) -> bool:
|
||||
if config is None or not getattr(config, "packing", False):
|
||||
return False
|
||||
return not getattr(config, "_unsloth_disable_auto_packing", False)
|
||||
|
||||
|
||||
def _should_auto_padding_free(config) -> bool:
|
||||
if (
|
||||
config is None
|
||||
or _AUTO_PADDING_FREE_ENV_DISABLED
|
||||
or getattr(config, "packing", False)
|
||||
):
|
||||
return False
|
||||
return not getattr(config, "padding_free", False)
|
||||
|
||||
|
||||
def _disable_sample_packing(config):
|
||||
if config is None:
|
||||
return
|
||||
for attr, value in (("packing", False), ("padding_free", False)):
|
||||
if hasattr(config, attr):
|
||||
setattr(config, attr, value)
|
||||
if hasattr(config, "remove_unused_columns"):
|
||||
setattr(config, "remove_unused_columns", True)
|
||||
setattr(config, "_unsloth_disable_auto_packing", True)
|
||||
|
||||
|
||||
_AUTO_PACK_SKIP_MESSAGES = (
|
||||
"packing is not supported",
|
||||
"padding-free training",
|
||||
"passing a custom data collator",
|
||||
)
|
||||
|
||||
|
||||
def _should_skip_auto_packing_error(exc: Exception) -> bool:
|
||||
message = str(exc).lower()
|
||||
return any(msg in message for msg in _AUTO_PACK_SKIP_MESSAGES)
|
||||
|
||||
|
||||
# Unsloth gradient accumulation fix:
|
||||
from transformers import __version__ as transformers_version
|
||||
from transformers import __version__ as transformers_version, ProcessorMixin
|
||||
|
||||
if Version(transformers_version) > Version("4.45.2"):
|
||||
|
||||
|
|
@ -211,6 +271,141 @@ def _backwards_compatible_trainer(trainer_class, config_class):
|
|||
return new_init
|
||||
|
||||
|
||||
def _patch_sft_trainer_auto_packing(trl_module):
|
||||
sft_trainer = getattr(trl_module, "SFTTrainer", None)
|
||||
if sft_trainer is None:
|
||||
return
|
||||
if getattr(sft_trainer, "_unsloth_auto_packing_wrapped", False):
|
||||
return
|
||||
|
||||
original_init = sft_trainer.__init__
|
||||
|
||||
@wraps(original_init)
|
||||
def new_init(self, *args, **kwargs):
|
||||
config_arg = None
|
||||
if len(args) >= 2:
|
||||
config_arg = args[1]
|
||||
else:
|
||||
config_arg = kwargs.get("args")
|
||||
|
||||
# Check if model type is unsupported for padding_free
|
||||
model = kwargs.get("model")
|
||||
is_unsupported_model = False
|
||||
is_vlm = False
|
||||
if model is not None:
|
||||
model_config = getattr(model, "config", None)
|
||||
if model_config is not None:
|
||||
model_types = get_transformers_model_type(model_config)
|
||||
# Blocklist: models that don't work correctly with padding_free
|
||||
is_unsupported_model = any(
|
||||
x in PADDING_FREE_BLOCKLIST for x in model_types
|
||||
)
|
||||
|
||||
# Check if VLM
|
||||
architectures = getattr(model_config, "architectures", None)
|
||||
if architectures is None:
|
||||
architectures = []
|
||||
is_vlm = any(
|
||||
x.endswith("ForConditionalGeneration") for x in architectures
|
||||
)
|
||||
is_vlm = is_vlm or hasattr(model_config, "vision_config")
|
||||
|
||||
processing_class = kwargs.get("processing_class") or kwargs.get("tokenizer")
|
||||
data_collator = kwargs.get("data_collator")
|
||||
|
||||
# We also disable vision language models for padding free collators
|
||||
blocked = (
|
||||
(data_collator is not None)
|
||||
or isinstance(processing_class, ProcessorMixin)
|
||||
or is_vlm
|
||||
or is_unsupported_model
|
||||
or (
|
||||
os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1"
|
||||
) # Disable padding free on forced logits
|
||||
)
|
||||
requested_pack = bool(getattr(config_arg, "packing", False))
|
||||
if blocked:
|
||||
if hasattr(config_arg, "packing"):
|
||||
setattr(config_arg, "packing", False)
|
||||
if hasattr(config_arg, "padding_free"):
|
||||
setattr(config_arg, "padding_free", False)
|
||||
|
||||
if blocked and requested_pack:
|
||||
reason = "custom data collator"
|
||||
if data_collator is None and isinstance(processing_class, ProcessorMixin):
|
||||
reason = "processor-based model"
|
||||
elif is_vlm:
|
||||
reason = "vision-language model"
|
||||
elif is_unsupported_model:
|
||||
reason = f"unsupported model type(s): {', '.join(model_types)}"
|
||||
message = "Unsloth: Sample packing skipped " f"({reason} detected)."
|
||||
print(message)
|
||||
|
||||
packing_active = False
|
||||
if _should_pack(config_arg) and not blocked:
|
||||
configure_sample_packing(config_arg)
|
||||
packing_active = True
|
||||
logger.info("Unsloth: Sample packing enabled for SFTTrainer instance.")
|
||||
|
||||
auto_padding_free_active = False
|
||||
padding_free_requested = getattr(config_arg, "padding_free", None) is True
|
||||
if not blocked:
|
||||
if padding_free_requested:
|
||||
configure_padding_free(config_arg)
|
||||
elif _should_auto_padding_free(config_arg):
|
||||
configure_padding_free(config_arg)
|
||||
auto_padding_free_active = True
|
||||
logger.info(
|
||||
"Unsloth: Padding-free batching auto-enabled for SFTTrainer instance."
|
||||
)
|
||||
|
||||
try:
|
||||
original_init(self, *args, **kwargs)
|
||||
except ValueError as exc:
|
||||
if packing_active and _should_skip_auto_packing_error(exc):
|
||||
logger.info(
|
||||
"Unsloth: Auto sample packing failed because trainer reported an incompatible setup (%s).",
|
||||
exc,
|
||||
)
|
||||
_disable_sample_packing(config_arg)
|
||||
packing_active = False
|
||||
original_init(self, *args, **kwargs)
|
||||
else:
|
||||
raise
|
||||
|
||||
trainer_args = getattr(self, "args", None)
|
||||
trainer_packing = bool(trainer_args and getattr(trainer_args, "packing", False))
|
||||
trainer_padding_free = bool(
|
||||
trainer_args and getattr(trainer_args, "padding_free", False)
|
||||
)
|
||||
|
||||
if blocked and trainer_args is not None:
|
||||
# Mirror the block on the trainer args to avoid re-enabling later
|
||||
setattr(trainer_args, "packing", False)
|
||||
setattr(trainer_args, "padding_free", False)
|
||||
|
||||
if (
|
||||
not blocked
|
||||
and trainer_packing
|
||||
and (packing_active or _should_pack(trainer_args))
|
||||
):
|
||||
enable_sample_packing(self.model, self)
|
||||
print(
|
||||
"🦥 Unsloth: Packing enabled - training is >2x faster and uses less VRAM!"
|
||||
)
|
||||
elif not blocked and trainer_padding_free:
|
||||
enable_padding_free_metadata(self.model, self)
|
||||
message = (
|
||||
"🦥 Unsloth: Padding-free auto-enabled, enabling faster training."
|
||||
if auto_padding_free_active
|
||||
else "🦥 Unsloth: Padding-free enabled, enabling faster training."
|
||||
)
|
||||
print(message)
|
||||
|
||||
sft_trainer.__init__ = new_init
|
||||
sft_trainer._unsloth_auto_packing_wrapped = True
|
||||
|
||||
|
||||
def _patch_trl_trainer():
|
||||
import trl
|
||||
|
||||
|
|
@ -237,4 +432,6 @@ def _patch_trl_trainer():
|
|||
except:
|
||||
continue
|
||||
|
||||
_patch_sft_trainer_auto_packing(trl)
|
||||
|
||||
trl.__UNSLOTH_BACKWARDS_COMPATIBLE__ = True
|
||||
|
|
|
|||
|
|
@ -0,0 +1,48 @@
|
|||
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Lesser General Public License as published by
|
||||
# the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Lesser General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
from .packing import (
|
||||
configure_padding_free,
|
||||
configure_sample_packing,
|
||||
enable_padding_free_metadata,
|
||||
enable_sample_packing,
|
||||
mark_allow_overlength,
|
||||
)
|
||||
from .attention_dispatch import (
|
||||
AttentionConfig,
|
||||
AttentionContext,
|
||||
FLASH_DENSE,
|
||||
FLASH_VARLEN,
|
||||
SDPA,
|
||||
XFORMERS,
|
||||
run_attention,
|
||||
select_attention_backend,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"configure_sample_packing",
|
||||
"configure_padding_free",
|
||||
"enable_sample_packing",
|
||||
"enable_padding_free_metadata",
|
||||
"mark_allow_overlength",
|
||||
"AttentionConfig",
|
||||
"AttentionContext",
|
||||
"FLASH_VARLEN",
|
||||
"FLASH_DENSE",
|
||||
"XFORMERS",
|
||||
"SDPA",
|
||||
"run_attention",
|
||||
"select_attention_backend",
|
||||
]
|
||||
283
unsloth/utils/attention_dispatch.py
Normal file
283
unsloth/utils/attention_dispatch.py
Normal file
|
|
@ -0,0 +1,283 @@
|
|||
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Lesser General Public License as published by
|
||||
# the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Lesser General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""Shared helpers for attention backend selection and execution."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional, Tuple
|
||||
|
||||
from torch import Tensor
|
||||
from torch.nn.functional import scaled_dot_product_attention
|
||||
|
||||
from ..models._utils import *
|
||||
from ..utils.packing import (
|
||||
build_sdpa_packed_attention_mask,
|
||||
build_xformers_block_causal_mask,
|
||||
)
|
||||
|
||||
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"
|
||||
FLASH_DENSE = "flash_dense"
|
||||
XFORMERS = "xformers"
|
||||
SDPA = "sdpa"
|
||||
|
||||
|
||||
XFORMERS_BLOCK_DIAG_CLS = (
|
||||
xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttentionConfig:
|
||||
"""
|
||||
Per-layer attention metadata.
|
||||
|
||||
NOTE(djsaunde): I had originally intended this to be populated once per layer, but
|
||||
we're currently constructing it on every forward pass since it can possibly be
|
||||
invalid from one forward pass to the next (e.g., switching from training to
|
||||
inference). For now, I'm keeping separate from AttentionContext for the sake of
|
||||
better grouping of params.
|
||||
"""
|
||||
|
||||
backend: str
|
||||
n_kv_heads: int
|
||||
n_groups: int
|
||||
flash_dense_kwargs: Optional[dict[str, Any]] = None
|
||||
flash_varlen_kwargs: Optional[dict[str, Any]] = None
|
||||
sdpa_kwargs: Optional[dict[str, Any]] = None
|
||||
xformers_kwargs: Optional[dict[str, Any]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttentionContext:
|
||||
"""Per-call info required to run attention."""
|
||||
|
||||
bsz: int
|
||||
q_len: int
|
||||
kv_seq_len: int
|
||||
n_heads: int
|
||||
head_dim: int
|
||||
requires_grad: bool
|
||||
seq_info: Optional[Tuple[Tensor, Tensor, int]]
|
||||
attention_mask: Optional[Tensor]
|
||||
causal_mask: Optional[Any]
|
||||
sliding_window: Optional[int] = None
|
||||
|
||||
|
||||
def select_attention_backend(use_varlen: bool = False) -> str:
|
||||
"""Return attention backend based on availability / priority order."""
|
||||
|
||||
if HAS_FLASH_ATTENTION:
|
||||
if use_varlen:
|
||||
return FLASH_VARLEN
|
||||
else:
|
||||
return FLASH_DENSE
|
||||
if HAS_XFORMERS:
|
||||
return XFORMERS
|
||||
return SDPA
|
||||
|
||||
|
||||
def run_attention(
|
||||
*,
|
||||
config: AttentionConfig,
|
||||
context: AttentionContext,
|
||||
Q: Tensor,
|
||||
K: Tensor,
|
||||
V: Tensor,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Run attention using config / context info.
|
||||
|
||||
Backend choice is prioritized for speed: FlashAttention when installed
|
||||
(`flash_varlen` for packed/variable-length inputs with `seq_info`, otherwise dense
|
||||
flash), then xFormers if flash is unavailable, with PyTorch SDPA as the final
|
||||
fallback (e.g., CPU or no fused kernels).
|
||||
|
||||
Varlen flash is preferred when packing metadata is present because it avoids padding
|
||||
and keeps peak memory low. xFormers and SDPA can also handle packed batches (we
|
||||
pass a block-diagonal mask into each).
|
||||
"""
|
||||
|
||||
backend = config.backend
|
||||
if backend == FLASH_VARLEN and context.seq_info is None:
|
||||
backend = FLASH_DENSE if HAS_FLASH_ATTENTION else SDPA
|
||||
flash_dense_kwargs = config.flash_dense_kwargs or {}
|
||||
flash_varlen_kwargs = config.flash_varlen_kwargs or {}
|
||||
sdpa_kwargs = config.sdpa_kwargs or {}
|
||||
xformers_kwargs = config.xformers_kwargs or {}
|
||||
|
||||
bsz = context.bsz
|
||||
n_heads = context.n_heads
|
||||
q_len = context.q_len
|
||||
head_dim = context.head_dim
|
||||
kv_seq_len = context.kv_seq_len
|
||||
requires_grad = context.requires_grad
|
||||
sliding_window = context.sliding_window
|
||||
|
||||
if backend == FLASH_VARLEN:
|
||||
Q_f = Q.transpose(1, 2).reshape(bsz * q_len, n_heads, head_dim)
|
||||
K_f = K.transpose(1, 2).reshape(bsz * q_len, config.n_kv_heads, head_dim)
|
||||
V_f = V.transpose(1, 2).reshape(bsz * q_len, config.n_kv_heads, head_dim)
|
||||
_, cu_seqlens, max_seqlen = context.seq_info
|
||||
return flash_attn_varlen_func(
|
||||
Q_f,
|
||||
K_f,
|
||||
V_f,
|
||||
cu_seqlens,
|
||||
cu_seqlens,
|
||||
max_seqlen,
|
||||
max_seqlen,
|
||||
**flash_varlen_kwargs,
|
||||
).view(bsz, q_len, n_heads, head_dim)
|
||||
elif backend == FLASH_DENSE:
|
||||
Q_t = Q.transpose(1, 2)
|
||||
K_t = K.transpose(1, 2)
|
||||
V_t = V.transpose(1, 2)
|
||||
return flash_attn_func(Q_t, K_t, V_t, **flash_dense_kwargs).reshape(
|
||||
bsz, q_len, n_heads, head_dim
|
||||
)
|
||||
elif backend == XFORMERS:
|
||||
attn_bias = build_xformers_block_causal_mask(
|
||||
context.seq_info,
|
||||
sliding_window = sliding_window,
|
||||
base_mask = context.causal_mask,
|
||||
)
|
||||
|
||||
Q_t = Q.transpose(1, 2)
|
||||
K_t = K.transpose(1, 2)
|
||||
V_t = V.transpose(1, 2)
|
||||
|
||||
K_mod = K_t
|
||||
V_mod = V_t
|
||||
Q_mod = Q_t
|
||||
|
||||
if config.n_groups != 1:
|
||||
K_mod = K_t.view(bsz, kv_seq_len, config.n_kv_heads, 1, head_dim)
|
||||
V_mod = V_t.view(bsz, kv_seq_len, config.n_kv_heads, 1, head_dim)
|
||||
K_mod = K_mod.expand(
|
||||
bsz, kv_seq_len, config.n_kv_heads, config.n_groups, head_dim
|
||||
)
|
||||
V_mod = V_mod.expand(
|
||||
bsz, kv_seq_len, config.n_kv_heads, config.n_groups, head_dim
|
||||
)
|
||||
|
||||
if requires_grad:
|
||||
K_mod = K_mod.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
V_mod = V_mod.reshape(bsz, kv_seq_len, n_heads, head_dim)
|
||||
else:
|
||||
Q_mod = Q_t.view(
|
||||
bsz, q_len, config.n_kv_heads, config.n_groups, head_dim
|
||||
)
|
||||
|
||||
has_block = XFORMERS_BLOCK_DIAG_CLS is not None and isinstance(
|
||||
attn_bias, XFORMERS_BLOCK_DIAG_CLS
|
||||
)
|
||||
|
||||
if config.n_groups != 1 and has_block:
|
||||
if not requires_grad:
|
||||
Q_mod = Q_mod.view(
|
||||
1, bsz * q_len, config.n_kv_heads, config.n_groups, head_dim
|
||||
)
|
||||
K_mod = K_mod.view(
|
||||
1, bsz * kv_seq_len, config.n_kv_heads, config.n_groups, head_dim
|
||||
)
|
||||
V_mod = V_mod.view(
|
||||
1, bsz * kv_seq_len, config.n_kv_heads, config.n_groups, head_dim
|
||||
)
|
||||
else:
|
||||
Q_mod = Q_mod.view(1, bsz * q_len, n_heads, head_dim)
|
||||
K_mod = K_mod.view(1, bsz * kv_seq_len, n_heads, head_dim)
|
||||
V_mod = V_mod.view(1, bsz * kv_seq_len, n_heads, head_dim)
|
||||
|
||||
out = xformers_attention(
|
||||
Q_mod,
|
||||
K_mod,
|
||||
V_mod,
|
||||
attn_bias = attn_bias,
|
||||
**xformers_kwargs,
|
||||
)
|
||||
|
||||
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.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)
|
||||
return out
|
||||
else:
|
||||
local_mask = context.attention_mask
|
||||
is_causal_local = False
|
||||
if context.seq_info is not None and local_mask is None:
|
||||
local_mask = build_sdpa_packed_attention_mask(
|
||||
context.seq_info,
|
||||
dtype = Q.dtype,
|
||||
device = Q.device,
|
||||
sliding_window = sliding_window,
|
||||
)
|
||||
else:
|
||||
q_len_local = Q.shape[-2]
|
||||
k_len_local = K.shape[-2]
|
||||
is_causal_local = local_mask is None and q_len_local == k_len_local
|
||||
|
||||
kwargs = dict(sdpa_kwargs)
|
||||
kwargs.setdefault("attn_mask", local_mask)
|
||||
kwargs.setdefault("is_causal", is_causal_local)
|
||||
|
||||
if SDPA_HAS_GQA:
|
||||
kwargs.setdefault("enable_gqa", config.n_groups != 1)
|
||||
out = scaled_dot_product_attention(Q, K, V, **kwargs)
|
||||
return out.transpose(1, 2)
|
||||
|
||||
K_mod = K
|
||||
V_mod = V
|
||||
if config.n_groups != 1:
|
||||
K_mod = K[:, :, None, :, :].expand(
|
||||
bsz, config.n_kv_heads, config.n_groups, kv_seq_len, head_dim
|
||||
)
|
||||
V_mod = V[:, :, None, :, :].expand(
|
||||
bsz, config.n_kv_heads, config.n_groups, kv_seq_len, head_dim
|
||||
)
|
||||
K_mod = K_mod.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
V_mod = V_mod.reshape(bsz, n_heads, kv_seq_len, head_dim)
|
||||
|
||||
out = scaled_dot_product_attention(
|
||||
Q.contiguous(),
|
||||
K_mod.contiguous(),
|
||||
V_mod.contiguous(),
|
||||
**kwargs,
|
||||
)
|
||||
return out.transpose(1, 2).contiguous()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AttentionConfig",
|
||||
"AttentionContext",
|
||||
"select_attention_backend",
|
||||
"run_attention",
|
||||
]
|
||||
344
unsloth/utils/packing.py
Normal file
344
unsloth/utils/packing.py
Normal file
|
|
@ -0,0 +1,344 @@
|
|||
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Lesser General Public License as published by
|
||||
# the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Lesser General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
"""Utilities for enabling packed (padding-free) batches across Unsloth."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Iterable, Optional, Sequence, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
from xformers.ops.fmha.attn_bias import (
|
||||
BlockDiagonalCausalMask as _XFormersBlockMask,
|
||||
)
|
||||
except Exception:
|
||||
try:
|
||||
from xformers.attn_bias import BlockDiagonalCausalMask as _XFormersBlockMask
|
||||
except Exception:
|
||||
_XFormersBlockMask = None
|
||||
|
||||
_XFORMERS_MASK_CACHE_MAXSIZE = 32
|
||||
_XFORMERS_MASK_CACHE: OrderedDict[Tuple[Tuple[int, ...], int], Any] = OrderedDict()
|
||||
|
||||
|
||||
def _window_cache_key(sliding_window: Optional[int]) -> int:
|
||||
if sliding_window is None or sliding_window <= 0:
|
||||
return 0
|
||||
return int(sliding_window)
|
||||
|
||||
|
||||
def _get_cached_block_mask(
|
||||
lengths: Tuple[int, ...],
|
||||
sliding_window: Optional[int],
|
||||
):
|
||||
if _XFormersBlockMask is None:
|
||||
return None
|
||||
|
||||
window_key = _window_cache_key(sliding_window)
|
||||
cache_key = (lengths, window_key)
|
||||
cached = _XFORMERS_MASK_CACHE.get(cache_key)
|
||||
if cached is not None:
|
||||
_XFORMERS_MASK_CACHE.move_to_end(cache_key)
|
||||
return cached
|
||||
|
||||
mask = _XFormersBlockMask.from_seqlens(list(lengths))
|
||||
if window_key and mask is not None and hasattr(mask, "make_local_attention"):
|
||||
mask = mask.make_local_attention(window_size = window_key)
|
||||
|
||||
_XFORMERS_MASK_CACHE[cache_key] = mask
|
||||
if len(_XFORMERS_MASK_CACHE) > _XFORMERS_MASK_CACHE_MAXSIZE:
|
||||
_XFORMERS_MASK_CACHE.popitem(last = False)
|
||||
return mask
|
||||
|
||||
|
||||
class _TrlPackingWarningFilter(logging.Filter):
|
||||
to_filter = (
|
||||
"attention implementation is not",
|
||||
"kernels-community",
|
||||
)
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
message = record.getMessage()
|
||||
return not any(substring in message for substring in self.to_filter)
|
||||
|
||||
|
||||
_TRL_FILTER_INSTALLED = False
|
||||
|
||||
|
||||
def _ensure_trl_warning_filter():
|
||||
global _TRL_FILTER_INSTALLED
|
||||
if _TRL_FILTER_INSTALLED:
|
||||
return
|
||||
logging.getLogger("trl.trainer.sft_trainer").addFilter(_TrlPackingWarningFilter())
|
||||
_TRL_FILTER_INSTALLED = True
|
||||
|
||||
|
||||
def mark_allow_overlength(module):
|
||||
"""Mark a module hierarchy so padding-free batches can exceed max_seq_length."""
|
||||
if module is None:
|
||||
return
|
||||
if hasattr(module, "max_seq_length"):
|
||||
setattr(module, "_unsloth_allow_packed_overlength", True)
|
||||
children = getattr(module, "children", None)
|
||||
if children is None:
|
||||
return
|
||||
for child in children():
|
||||
mark_allow_overlength(child)
|
||||
|
||||
|
||||
def configure_sample_packing(config):
|
||||
"""Mutate an ``SFTConfig`` so TRL prepares packed batches."""
|
||||
_ensure_trl_warning_filter()
|
||||
setattr(config, "packing", True)
|
||||
setattr(config, "padding_free", True)
|
||||
|
||||
|
||||
def configure_padding_free(config):
|
||||
"""Mutate an ``SFTConfig`` so TRL enables padding-free batching without packing."""
|
||||
_ensure_trl_warning_filter()
|
||||
setattr(config, "padding_free", True)
|
||||
|
||||
|
||||
def enable_sample_packing(
|
||||
model,
|
||||
trainer,
|
||||
*,
|
||||
sequence_lengths_key: str = "seq_lengths",
|
||||
) -> None:
|
||||
"""Enable runtime support for packed batches on an existing trainer."""
|
||||
if model is None or trainer is None:
|
||||
raise ValueError("model and trainer must not be None")
|
||||
|
||||
mark_allow_overlength(model)
|
||||
|
||||
if hasattr(trainer, "args") and hasattr(trainer.args, "remove_unused_columns"):
|
||||
trainer.args.remove_unused_columns = False
|
||||
|
||||
collator = getattr(trainer, "data_collator", None)
|
||||
if collator is None or not hasattr(collator, "torch_call"):
|
||||
return
|
||||
if getattr(collator, "_unsloth_packing_wrapped", False):
|
||||
return
|
||||
|
||||
if hasattr(collator, "padding_free"):
|
||||
collator.padding_free = True
|
||||
if hasattr(collator, "return_position_ids"):
|
||||
collator.return_position_ids = True
|
||||
|
||||
original_torch_call = collator.torch_call
|
||||
|
||||
def torch_call_with_lengths(examples: Sequence[dict]):
|
||||
batch = original_torch_call(examples)
|
||||
if examples and isinstance(examples[0], dict):
|
||||
seq_lengths: list[int] = []
|
||||
for example in examples:
|
||||
lengths = example.get(sequence_lengths_key)
|
||||
if isinstance(lengths, Iterable):
|
||||
seq_lengths.extend(int(length) for length in lengths)
|
||||
if seq_lengths:
|
||||
batch["packed_seq_lengths"] = torch.tensor(
|
||||
seq_lengths, dtype = torch.int32
|
||||
)
|
||||
if "attention_mask" in batch:
|
||||
batch.pop("attention_mask")
|
||||
return batch
|
||||
|
||||
collator.torch_call = torch_call_with_lengths
|
||||
collator._unsloth_packing_wrapped = True
|
||||
|
||||
|
||||
def enable_padding_free_metadata(model, trainer):
|
||||
"""Inject seq-length metadata when padding-free batching is enabled without packing."""
|
||||
collator = getattr(trainer, "data_collator", None)
|
||||
if (
|
||||
collator is None
|
||||
or getattr(collator, "_unsloth_padding_free_lengths_wrapped", False)
|
||||
or not getattr(collator, "padding_free", False)
|
||||
):
|
||||
return
|
||||
|
||||
mark_allow_overlength(model)
|
||||
if hasattr(collator, "return_position_ids"):
|
||||
collator.return_position_ids = True
|
||||
|
||||
original_torch_call = collator.torch_call
|
||||
|
||||
def torch_call_with_padding_free_metadata(examples: Sequence[dict]):
|
||||
seq_lengths: list[int] = []
|
||||
if examples and isinstance(examples[0], dict):
|
||||
for example in examples:
|
||||
lengths = example.get("seq_lengths")
|
||||
if lengths is None:
|
||||
ids = example.get("input_ids")
|
||||
if ids is None:
|
||||
continue
|
||||
lengths = [len(ids)]
|
||||
example["seq_lengths"] = lengths
|
||||
seq_lengths.extend(lengths)
|
||||
|
||||
batch = original_torch_call(examples)
|
||||
if seq_lengths:
|
||||
batch["packed_seq_lengths"] = torch.tensor(
|
||||
seq_lengths,
|
||||
dtype = torch.int32,
|
||||
)
|
||||
return batch
|
||||
|
||||
collator.torch_call = torch_call_with_padding_free_metadata
|
||||
collator._unsloth_padding_free_lengths_wrapped = True
|
||||
|
||||
|
||||
def get_packed_info_from_kwargs(
|
||||
kwargs: dict,
|
||||
device: torch.device,
|
||||
) -> Optional[Tuple[torch.Tensor, torch.Tensor, int]]:
|
||||
"""Return packed sequence metadata expected by the attention kernels."""
|
||||
|
||||
seq_lengths = kwargs.get("packed_seq_lengths")
|
||||
if seq_lengths is None:
|
||||
return None
|
||||
|
||||
lengths = seq_lengths.to(device = device, dtype = torch.int32, non_blocking = True)
|
||||
cu_seqlens = torch.empty(lengths.numel() + 1, dtype = torch.int32, device = device)
|
||||
cu_seqlens[0] = 0
|
||||
torch.cumsum(lengths, dim = 0, dtype = torch.int32, out = cu_seqlens[1:])
|
||||
|
||||
max_seqlen = int(lengths.max().item())
|
||||
return lengths, cu_seqlens, max_seqlen
|
||||
|
||||
|
||||
def build_xformers_block_causal_mask(
|
||||
seq_info: Optional[Tuple[torch.Tensor, torch.Tensor, int]],
|
||||
*,
|
||||
sliding_window: Optional[int] = None,
|
||||
base_mask: Optional[Any] = None,
|
||||
):
|
||||
if _XFormersBlockMask is None:
|
||||
return None
|
||||
if seq_info is not None:
|
||||
seq_lengths, _, _ = seq_info
|
||||
lengths_tensor = seq_lengths.to("cpu", torch.int32)
|
||||
if lengths_tensor.numel() == 0:
|
||||
return None
|
||||
lengths = tuple(int(x) for x in lengths_tensor.tolist())
|
||||
mask = _get_cached_block_mask(lengths, sliding_window)
|
||||
else:
|
||||
mask = base_mask
|
||||
|
||||
if (
|
||||
sliding_window is not None
|
||||
and sliding_window > 0
|
||||
and mask is not None
|
||||
and hasattr(mask, "make_local_attention")
|
||||
):
|
||||
mask = mask.make_local_attention(window_size = sliding_window)
|
||||
return mask
|
||||
|
||||
|
||||
def build_sdpa_packed_attention_mask(
|
||||
seq_info: Tuple[torch.Tensor, torch.Tensor, int],
|
||||
*,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
sliding_window: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
seq_lengths, _, _ = seq_info
|
||||
total_tokens = int(seq_lengths.sum().item())
|
||||
mask = torch.full(
|
||||
(total_tokens, total_tokens),
|
||||
float("-inf"),
|
||||
dtype = dtype,
|
||||
device = device,
|
||||
)
|
||||
offset = 0
|
||||
for length in seq_lengths.tolist():
|
||||
length = int(length)
|
||||
if length <= 0:
|
||||
continue
|
||||
block = torch.zeros((length, length), dtype = dtype, device = device)
|
||||
upper = torch.triu(
|
||||
torch.ones((length, length), device = device), diagonal = 1
|
||||
).bool()
|
||||
block = block.masked_fill(upper, float("-inf"))
|
||||
if (
|
||||
sliding_window is not None
|
||||
and sliding_window > 0
|
||||
and length > sliding_window
|
||||
):
|
||||
idx = torch.arange(length, device = device)
|
||||
dist = idx.unsqueeze(1) - idx.unsqueeze(0)
|
||||
window_mask = dist >= sliding_window
|
||||
block = block.masked_fill(window_mask, float("-inf"))
|
||||
mask[offset : offset + length, offset : offset + length] = block
|
||||
offset += length
|
||||
return mask.unsqueeze(0).unsqueeze(0)
|
||||
|
||||
|
||||
def _normalize_packed_lengths(
|
||||
seq_lengths: Any,
|
||||
*,
|
||||
device: torch.device,
|
||||
) -> Optional[torch.Tensor]:
|
||||
if seq_lengths is None:
|
||||
return None
|
||||
if isinstance(seq_lengths, torch.Tensor):
|
||||
lengths = seq_lengths.to(device = device, dtype = torch.int64)
|
||||
else:
|
||||
lengths = torch.tensor(seq_lengths, device = device, dtype = torch.int64)
|
||||
if lengths.ndim != 1:
|
||||
lengths = lengths.reshape(-1)
|
||||
if lengths.numel() == 0:
|
||||
return None
|
||||
return lengths
|
||||
|
||||
|
||||
def mask_packed_sequence_boundaries(
|
||||
shift_labels: torch.Tensor,
|
||||
seq_lengths: Any,
|
||||
*,
|
||||
ignore_index: int = -100,
|
||||
) -> bool:
|
||||
"""Mark final token of every packed sample so CE ignores boundary predictions."""
|
||||
lengths = _normalize_packed_lengths(seq_lengths, device = shift_labels.device)
|
||||
if lengths is None:
|
||||
return False
|
||||
|
||||
flat = shift_labels.reshape(-1)
|
||||
total_tokens = flat.shape[0]
|
||||
boundary_positions = torch.cumsum(lengths, dim = 0) - 1
|
||||
valid = boundary_positions < total_tokens
|
||||
if not torch.all(valid):
|
||||
boundary_positions = boundary_positions[valid]
|
||||
if boundary_positions.numel() == 0:
|
||||
return False
|
||||
flat[boundary_positions] = ignore_index
|
||||
return True
|
||||
|
||||
|
||||
__all__ = [
|
||||
"configure_sample_packing",
|
||||
"configure_padding_free",
|
||||
"enable_sample_packing",
|
||||
"enable_padding_free_metadata",
|
||||
"mark_allow_overlength",
|
||||
"get_packed_info_from_kwargs",
|
||||
"build_xformers_block_causal_mask",
|
||||
"build_sdpa_packed_attention_mask",
|
||||
"mask_packed_sequence_boundaries",
|
||||
]
|
||||
Loading…
Add table
Add a link
Reference in a new issue