Float8 GRPO, RL (#3640)

* Enable FP8 + RL training for bf16 models (#3440)

* Enable FP8 + RL training for bf16 models

**Summary:** Enable FP8 + RL training using TorchAO for 1.33x faster training and 42% less model memory usage:
- We quantize the frozen LoRA weights into fp8 and keep the LoRA adapters in bf16
- We leverage TorchAO's `Float8Tensor`, which calls into fbgemm's fp8 x fp8 rowwise matmul kernel
- For now, we need to do an offline quantization first, because vllm doesn't support on-the-fly quantization for torchao yet  (this is in progress: https://github.com/vllm-project/vllm/pull/26327)

**Example usage:**
```
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Qwen3-8B-Base",
    max_seq_length = 2048,
    load_in_4bit = False,
    fast_inference = True,
    max_lora_rank = 32,
    load_in_fp8 = True,  # set this to True
)

\# the rest is the same as before
model = FastLanguageModel.get_peft_model(...)
```

**Initial results:**
```
\# fp8
{'train_runtime': 1725.4337, 'train_samples_per_second': 0.232, 'train_steps_per_second': 0.058, 'train_loss': 0.00015715716748673002, 'epoch': 0.01}

\# bf16
{'train_runtime': 2297.8145, 'train_samples_per_second': 0.174, 'train_steps_per_second': 0.044, 'train_loss': 0.00016081033063528594, 'epoch': 0.01}
```

<img width="1199" height="448" alt="Screenshot 2025-11-11 at 4 10 50 PM" src="https://github.com/user-attachments/assets/b6304afd-89e9-42b1-8064-775807e17b23" />

Test script: https://gist.github.com/andrewor14/5b85119fae46845d07b608d420907423

**Requires:**
- https://github.com/pytorch/ao/pull/3158 (torchao nightly or 0.15.0+)
- https://github.com/unslothai/unsloth-zoo/pull/351

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update utils.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* _get_inference_mode_context_manager

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update utils.py

* Update utils.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Update __init__.py

* Fix/save torchao model loading logic (#3621)

* make loading gpt-oss-BF16 faster. Linked to unsloth-zoo PR #314

* fix model loading and clean merged model directory

* revert default quant

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* revert mapper.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Update loader_utils.py

* Update loader_utils.py

* Add 128x128 PerBlock FP8 + RL (#3629)

* Add 128x128 PerBlock FP8 + RL

**Summary:** Following https://github.com/unslothai/unsloth/pull/3440,
this PR extends torchao FP8 + RL support to also handle 128x128
PerBlock granularity (in addition to PerRow).

**Example usage:**

```
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Qwen3-8B-Base",
    max_seq_length = 2048,
    load_in_4bit = False,
    fast_inference = True,
    max_lora_rank = 32,
    load_in_fp8 = "block",  # or "row" or True
)
```

**Initial results:** TBD

**Note:**
- Requires https://github.com/pytorch/ao/pull/3370

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Version

* Update vision.py

* Update rl.py

* Add torch 2.9.1

* Fix auto installer

* Update fp8.py

* Float8

* Update fp8.py

* Update mapper.py

* Update mapper.py

* Update loader_utils.py

* Update loader.py

* Update fp8.py

* Versioning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: andrewor14 <andrewor14@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
This commit is contained in:
Daniel Han 2025-11-25 07:23:26 -08:00 committed by GitHub
commit 38aa148aba
15 changed files with 818 additions and 156 deletions

View file

@ -197,9 +197,9 @@ pip install unsloth
</details>
### Advanced Pip Installation
`⚠Do **NOT** use this if you have Conda.` Pip is a bit more complex since there are dependency issues. The pip command is different for `torch 2.2,2.3,2.4,2.5` and CUDA versions.
`⚠Do **NOT** use this if you have Conda.` Pip is a bit more complex since there are dependency issues. The pip command is different for `torch 2.2,2.3,2.4,2.5,2.6,2.7,2.8,2.9` and CUDA versions.
For other torch versions, we support `torch211`, `torch212`, `torch220`, `torch230`, `torch240` and for CUDA versions, we support `cu118` and `cu121` and `cu124`. For Ampere devices (A100, H100, RTX3090) and above, use `cu118-ampere` or `cu121-ampere` or `cu124-ampere`.
For other torch versions, we support `torch211`, `torch212`, `torch220`, `torch230`, `torch240`, `torch250`, `torch260`, `torch270`, `torch280`, `torch290` and for CUDA versions, we support `cu118` and `cu121` and `cu124`. For Ampere devices (A100, H100, RTX3090) and above, use `cu118-ampere` or `cu121-ampere` or `cu124-ampere`.
For example, if you have `torch 2.4` and `CUDA 12.1`, use:
```bash
@ -207,10 +207,10 @@ pip install --upgrade pip
pip install "unsloth[cu121-torch240] @ git+https://github.com/unslothai/unsloth.git"
```
Another example, if you have `torch 2.5` and `CUDA 12.4`, use:
Another example, if you have `torch 2.9` and `CUDA 13.0`, use:
```bash
pip install --upgrade pip
pip install "unsloth[cu124-torch250] @ git+https://github.com/unslothai/unsloth.git"
pip install "unsloth[cu130-torch290] @ git+https://github.com/unslothai/unsloth.git"
```
And other examples:
@ -242,7 +242,7 @@ v = V(re.match(r"[0-9\.]{3,}", torch.__version__).group(0))
cuda = str(torch.version.cuda)
is_ampere = torch.cuda.get_device_capability()[0] >= 8
USE_ABI = torch._C._GLIBCXX_USE_CXX11_ABI
if cuda not in ("11.8", "12.1", "12.4", "12.6", "12.8"): raise RuntimeError(f"CUDA = {cuda} not supported!")
if cuda not in ("11.8", "12.1", "12.4", "12.6", "12.8", "13.0"): raise RuntimeError(f"CUDA = {cuda} not supported!")
if v <= V('2.1.0'): raise RuntimeError(f"Torch = {v} too old!")
elif v <= V('2.1.1'): x = 'cu{}{}-torch211'
elif v <= V('2.1.2'): x = 'cu{}{}-torch212'
@ -255,10 +255,12 @@ elif v < V('2.7.0'): x = 'cu{}{}-torch260'
elif v < V('2.7.9'): x = 'cu{}{}-torch270'
elif v < V('2.8.0'): x = 'cu{}{}-torch271'
elif v < V('2.8.9'): x = 'cu{}{}-torch280'
elif v < V('2.9.1'): x = 'cu{}{}-torch290'
elif v < V('2.9.2'): x = 'cu{}{}-torch291'
else: raise RuntimeError(f"Torch = {v} too new!")
if v > V('2.6.9') and cuda not in ("11.8", "12.6", "12.8"): raise RuntimeError(f"CUDA = {cuda} not supported!")
x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "")
print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"')
if v > V('2.6.9') and cuda not in ("11.8", "12.6", "12.8", "13.0"): raise RuntimeError(f"CUDA = {cuda} not supported!")
x = x.format(cuda.replace(".", ""), "-ampere" if False else "") # is_ampere is broken due to flash-attn
print(f'pip install --upgrade pip && pip install --no-deps git+https://github.com/unslothai/unsloth-zoo.git && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git" --no-build-isolation')
```
### Docker Installation
You can use our pre-built Docker container with all dependencies to use Unsloth instantly with no setup required.

View file

@ -60,7 +60,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2025.11.4",
"unsloth_zoo>=2025.11.5",
"torchvision",
"unsloth[triton]",
]
@ -273,10 +273,22 @@ cu128onlytorch280 = [
cu130onlytorch280 = [
]
cu126onlytorch290 = [
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.33.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.33.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')",
]
cu128onlytorch290 = [
"xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.33.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu128/xformers-0.0.33.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')",
]
cu130onlytorch290 = [
"xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu130/xformers-0.0.33.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')",
]
cu126onlytorch291 = [
]
cu128onlytorch291 = [
]
cu130onlytorch291 = [
]
cu118 = [
"unsloth[huggingface]",
@ -453,6 +465,21 @@ cu130-torch290 = [
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu130onlytorch290]",
]
cu126-torch291 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu126onlytorch291]",
]
cu128-torch291 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu128onlytorch291]",
]
cu130-torch291 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu130onlytorch291]",
]
kaggle = [
"unsloth[huggingface]",
]
@ -490,7 +517,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2025.11.4",
"unsloth_zoo>=2025.11.5",
"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",
@ -721,6 +748,21 @@ cu130-ampere-torch290 = [
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu130onlytorch290]",
]
cu126-ampere-torch291 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu126onlytorch291]",
]
cu128-ampere-torch291 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu128onlytorch291]",
]
cu130-ampere-torch291 = [
"unsloth[huggingface]",
"bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0",
"unsloth[cu130onlytorch291]",
]
flashattentiontorch260abiFALSEcu12x = [
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'",

View file

@ -177,6 +177,7 @@ if DEVICE_TYPE == "cuda":
print(
"Unsloth: `bitsandbytes` is not installed - 4bit QLoRA unallowed, but 16bit and full finetuning works!"
)
bnb = None
try:
cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
libcuda_dirs()
@ -209,7 +210,8 @@ if DEVICE_TYPE == "cuda":
del find_number, latest_cuda
del possible_cudas, find_cuda
importlib.reload(bnb)
if bnb is not None:
importlib.reload(bnb)
importlib.reload(triton)
try:
libcuda_dirs = lambda: None

View file

@ -34,7 +34,8 @@ elif v < V('2.7.9'): x = 'cu{}{}-torch270'
elif v < V('2.8.0'): x = 'cu{}{}-torch271'
elif v < V('2.8.9'): x = 'cu{}{}-torch280'
elif v < V('2.9.1'): x = 'cu{}{}-torch290'
elif v < V('2.9.2'): x = 'cu{}{}-torch291'
else: raise RuntimeError(f"Torch = {v} too new!")
if v > V('2.6.9') and cuda not in ("11.8", "12.6", "12.8", "13.0"): raise RuntimeError(f"CUDA = {cuda} not supported!")
x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "")
print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"')
x = x.format(cuda.replace(".", ""), "-ampere" if False else "") # is_ampere is broken due to flash-attn
print(f'pip install --upgrade pip && pip install --no-deps git+https://github.com/unslothai/unsloth-zoo.git && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git" --no-build-isolation')

View file

@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import torch
import torch.nn as nn
import triton
@ -519,7 +520,41 @@ def fp8_fbgemm_block_linear(X, weight, weight_scale, bias = None):
return FP8_fbgemm_block_linear.apply(X, weight, weight_scale, bias)
def test_has_fbgemm():
# We must manually check if the faster FBGEMM works on the specific GPU
# For example RTX 5090 and RTX 4090 does not work
# [TODO] Investigate with TorchAO why FBGEMM fails on consumer GPUs
M, N, K = 128, 128, 128
xq = torch.ones(M, K, dtype = torch.float8_e4m3fn, device = "cuda")
wq = xq
M, K = xq.shape
N, _ = wq.shape
block_scale = torch.ones(M // 128, K // 128, dtype = torch.float32, device = "cuda")
has_fbgemm = False
try:
out = torch.ops.fbgemm.f8f8bf16_blockwise(xq, wq, block_scale, block_scale)
assert torch.unique(out).item() == 128
has_fbgemm = True
del out
except Exception as e:
e = str(e)
if "cutlass cannot initialize" in e.lower():
print(
f"Unsloth: FBGEMM on the current GPU cannot load - will switch to Triton kernels"
)
else:
print(
f"Unsloth: FBGEMM on the current GPU cannot load with error = {e} - will switch to Triton kernels"
)
has_fbgemm = False
del block_scale, xq
torch.cuda.empty_cache()
return has_fbgemm
fp8_block_quant_linear = fp8_torch_block_quant_forward
if "UNSLOTH_HAS_FBGEMM" not in os.environ:
os.environ["UNSLOTH_HAS_FBGEMM"] = "0"
try:
import fbgemm_gpu
@ -527,8 +562,14 @@ try:
# This is both fast and accurate hence preferred.
# This makes it 15% faster than the torchao implementation.
if Version(fbgemm_gpu.__version__) >= Version("1.4.0"):
logger.info(f"Using fbgemm_gpu block quantized FP8 matmul")
fp8_block_quant_linear = fp8_fbgemm_block_linear
# We must manually confirm if blockwise FBGEMM works!
# This check is a must for consumer grade GPUs which fail
if test_has_fbgemm():
os.environ["UNSLOTH_HAS_FBGEMM"] = "1"
logger.info(f"Using fbgemm_gpu block quantized FP8 matmul")
fp8_block_quant_linear = fp8_fbgemm_block_linear
else:
os.environ["UNSLOTH_HAS_FBGEMM"] = "0"
except:
pass

View file

@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import importlib
import triton
import ctypes
@ -35,7 +36,7 @@ import functools
import torch
torch_Tensor = torch.Tensor
from packaging.version import Version
from unsloth_zoo.utils import Version
if DEVICE_TYPE == "xpu" and Version(torch.__version__) < Version("2.6.0"):
raise RuntimeError(
@ -55,7 +56,6 @@ if DEVICE_TYPE == "xpu":
# tl.math.tanh now is libdevice.tanh
from packaging.version import Version
import triton
import triton.language as tl
@ -211,6 +211,22 @@ torch_float16 = torch.float16
torch_bfloat16 = torch.bfloat16
# Check whether torchao can be imported to get Float8Tensor
if importlib.util.find_spec("torchao") is not None:
try:
from torchao.quantization import Float8Tensor
except:
import torchao
if Version(torchao.__version__) >= Version("0.15.0"):
print(
f"Unsloth: `from torchao.quantization import Float8Tensor` failed on version={torchao.__version__}"
)
Float8Tensor = type(None)
else:
Float8Tensor = type(None)
def QUANT_STATE(W):
return getattr(W, "quant_state", None)
@ -335,6 +351,8 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM:
@torch.inference_mode
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
# TODO: After adding XPU BNB support, check this function
if isinstance(W, Float8Tensor):
return W.dequantize()
if quant_state is None:
return W
if W.dtype == torch.float8_e4m3fn:
@ -441,6 +459,8 @@ elif DEVICE_TYPE in ("cuda", "hip") and HAS_CUDA_STREAM:
@torch.inference_mode
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
if isinstance(W, Float8Tensor):
return W.dequantize()
if quant_state is None:
return W
if W.dtype == torch.float8_e4m3fn:
@ -551,6 +571,8 @@ else:
@torch.inference_mode
def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False):
if isinstance(W, Float8Tensor):
return W.dequantize()
if quant_state is None:
return W
if W.dtype == torch.float8_e4m3fn:
@ -984,11 +1006,21 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None):
else:
reshape = False
if W.dtype == torch.float8_e4m3fn:
if isinstance(W, Float8Tensor):
assert W.ndim == 2
if W.block_size[0] == W.shape[0] and W.block_size[1] == 1:
# In the backward pass, rowwise scaled becomes colwise scaled after we
# transpose the weight tensor. Use this case to detect backward.
# TODO: would be simpler if we simply don't call `matmul_lora` in backward
W = W.dequantize()
else:
W = W.contiguous()
out = torch_matmul(X, W.t(), out = out)
elif W.dtype == torch.float8_e4m3fn:
out = fp8_linear(X, W, W_quant)
else:
W = fast_dequantize(W.t(), W_quant, use_global_buffer = True)
out = torch_matmul(X, W, out = out)
W = fast_dequantize(W, W_quant, use_global_buffer = True)
out = torch_matmul(X, W.t(), out = out)
if W_quant is not None:
del W

View file

@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
__version__ = "2025.11.3"
__version__ = "2025.11.4"
__all__ = [
"SUPPORTS_BFLOAT16",
@ -71,6 +71,7 @@ __all__ = [
"dequantize_module_weight",
"patch_hf_quantizer",
"verify_fp8_support_if_applicable",
"_get_inference_mode_context_manager",
]
import torch
@ -2056,7 +2057,7 @@ except:
@dataclass
class TorchAOConfig:
qat_scheme: str = "int4"
qat_scheme: Optional[str] = "int4"
# Each (config, filter_fn) pair defines a quantization rule
base_config_and_filter_fns: List[
@ -2306,3 +2307,22 @@ def verify_fp8_support_if_applicable(model_config):
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):
"""
If the state dict was quantized using torchao, we will run into
the following error when calling ops like aten.t() in inference mode.
This is a bug in PyTorch that affects all tensor subclasses.
Cannot set version_counter for inference tensor
For now, we work around this issue by using `torch.no_grad()` in this case.
See https://github.com/pytorch/pytorch/issues/164872 for more details.
Otherwise, just return `torch.inference_mode()`.
"""
torchao_config = getattr(model, "torchao_config", None)
if torchao_config is not None and torchao_config.qat_scheme is None:
return torch.no_grad()
else:
return torch.inference_mode()

View file

@ -21,7 +21,10 @@ from ._utils import *
from ._utils import patch_unsloth_smart_gradient_checkpointing
from ._utils import __version__, importlib_version
from ._utils import move_to_device
from ._utils import _prepare_model_for_qat
from ._utils import (
_get_inference_mode_context_manager,
_prepare_model_for_qat,
)
from torch.nn.functional import scaled_dot_product_attention
from transformers import __version__ as transformers_version
from unsloth_zoo.utils import Version, _get_dtype
@ -2030,7 +2033,7 @@ def unsloth_fast_generate(
# Mixed precision autocast
with (
torch.inference_mode(),
_get_inference_mode_context_manager(self),
torch.autocast(device_type = DEVICE_TYPE_TORCH, dtype = dtype),
):
output = self._old_generate(*args, **kwargs)

View file

@ -31,7 +31,12 @@ from .cohere import FastCohereModel
from transformers import AutoConfig
from transformers import __version__ as transformers_version
from peft import PeftConfig, PeftModel
from .loader_utils import get_model_name
from .loader_utils import (
_get_fp8_mode_and_check_settings,
_offline_quantize_to_fp8,
_tag_model_with_fp8_torchao_config,
get_model_name,
)
import os, contextlib, sys
try:
@ -140,6 +145,7 @@ class FastLanguageModel(FastLlamaModel):
max_lora_rank = 64,
disable_log_stats = True,
qat_scheme = None,
load_in_fp8 = False, # fp8 LoRA (True, False, 'block')
*args,
**kwargs,
):
@ -183,6 +189,7 @@ class FastLanguageModel(FastLlamaModel):
max_lora_rank = max_lora_rank,
disable_log_stats = disable_log_stats,
qat_scheme = qat_scheme,
load_in_fp8 = load_in_fp8,
*args,
**kwargs,
)
@ -212,9 +219,28 @@ class FastLanguageModel(FastLlamaModel):
)
load_in_4bit = False
# Find FP8, BnB 4bit, other mapped names
old_model_name = model_name
fp8_mode = None
if not use_exact_model_name:
model_name = get_model_name(model_name, load_in_4bit)
new_model_name = get_model_name(
model_name, load_in_4bit = load_in_4bit, load_in_fp8 = load_in_fp8
)
if new_model_name is None and load_in_fp8 != False:
fp8_mode = _get_fp8_mode_and_check_settings(
load_in_fp8,
fast_inference,
full_finetuning,
load_in_4bit,
load_in_8bit,
load_in_16bit,
use_exact_model_name,
)
model_name = _offline_quantize_to_fp8(model_name, fp8_mode)
else:
assert new_model_name is not None
model_name = new_model_name
# Check if pre-quantized models are allowed
# For eg AMD GPUs need blocksize = 128, but our pre-quants are blocksize = 64
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
@ -354,6 +380,7 @@ class FastLanguageModel(FastLlamaModel):
if model_name.lower().endswith("-bf16"):
load_in_4bit = False
load_in_8bit = False
load_in_fp8 = False
load_in_16bit = True
model_config = AutoConfig.from_pretrained(
@ -476,6 +503,8 @@ class FastLanguageModel(FastLlamaModel):
random_state = random_state,
max_lora_rank = max_lora_rank,
disable_log_stats = disable_log_stats,
qat_scheme = qat_scheme,
load_in_fp8 = load_in_fp8,
*args,
**kwargs,
)
@ -554,6 +583,9 @@ class FastLanguageModel(FastLlamaModel):
}
model.config.update({"quantization_config": quantization_config})
if load_in_fp8 != False:
_tag_model_with_fp8_torchao_config(model, fp8_mode)
if is_peft:
# From https://github.com/huggingface/peft/issues/184
# Now add PEFT adapters
@ -634,6 +666,7 @@ class FastModel(FastBaseModel):
max_lora_rank = 64,
disable_log_stats = True,
qat_scheme = None,
load_in_fp8 = False, # fp8 LoRA (True, False, 'block')
*args,
**kwargs,
):
@ -660,6 +693,7 @@ class FastModel(FastBaseModel):
)
dtype = torch.float16
assert dtype in (torch.float16, torch.bfloat16, torch.float32)
assert load_in_fp8 in (True, False, "block")
patch_compiled_autograd()
patch_compiling_bitsandbytes()
@ -694,9 +728,28 @@ class FastModel(FastBaseModel):
)
load_in_4bit = False
# Find FP8, BnB 4bit, other mapped names
old_model_name = model_name
fp8_mode = None
if not use_exact_model_name:
model_name = get_model_name(model_name, load_in_4bit)
new_model_name = get_model_name(
model_name, load_in_4bit = load_in_4bit, load_in_fp8 = load_in_fp8
)
if new_model_name is None and load_in_fp8 != False:
fp8_mode = _get_fp8_mode_and_check_settings(
load_in_fp8,
fast_inference,
full_finetuning,
load_in_4bit,
load_in_8bit,
load_in_16bit,
use_exact_model_name,
)
model_name = _offline_quantize_to_fp8(model_name, fp8_mode)
else:
assert new_model_name is not None
model_name = new_model_name
# Check if pre-quantized models are allowed
# For eg AMD GPUs need blocksize = 128, but our pre-quants are blocksize = 64
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
@ -708,6 +761,7 @@ class FastModel(FastBaseModel):
if model_name.lower().endswith("-bf16"):
load_in_4bit = False
load_in_8bit = False
load_in_fp8 = False
load_in_16bit = True
# Check modelscope
@ -786,6 +840,8 @@ class FastModel(FastBaseModel):
string += "_load_in_8bit_"
if load_in_16bit:
string += "_load_in_16bit_"
if load_in_fp8:
string += "load_in_fp8"
os.environ["UNSLOTH_MODEL_NAME"] = string
# Check versions
@ -966,6 +1022,7 @@ class FastModel(FastBaseModel):
if model_name.lower().endswith("-bf16"):
load_in_4bit = False
load_in_8bit = False
load_in_fp8 = False
load_in_16bit = True
model_config = AutoConfig.from_pretrained(
@ -1130,6 +1187,9 @@ class FastModel(FastBaseModel):
}
model.config.update({"quantization_config": quantization_config})
if load_in_fp8 != False:
_tag_model_with_fp8_torchao_config(model, fp8_mode)
if is_peft:
# From https://github.com/huggingface/peft/issues/184
# Now add PEFT adapters

View file

@ -12,11 +12,26 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from .mapper import INT_TO_FLOAT_MAPPER, FLOAT_TO_INT_MAPPER, MAP_TO_UNSLOTH_16bit
import importlib
import os
import re
import tempfile
from typing import Union
from .mapper import (
INT_TO_FLOAT_MAPPER,
FLOAT_TO_INT_MAPPER,
MAP_TO_UNSLOTH_16bit,
FLOAT_TO_FP8_BLOCK_MAPPER,
FLOAT_TO_FP8_ROW_MAPPER,
)
# https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading!
from packaging.version import Version
from transformers import __version__ as transformers_version
from unsloth.models._utils import TorchAOConfig
from unsloth_zoo.utils import Version
import torch
import gc
transformers_version = Version(transformers_version)
SUPPORTS_FOURBIT = transformers_version >= Version("4.37")
@ -36,11 +51,28 @@ def __get_model_name(
INT_TO_FLOAT_MAPPER = None,
FLOAT_TO_INT_MAPPER = None,
MAP_TO_UNSLOTH_16bit = None,
load_in_fp8 = False,
FLOAT_TO_FP8_BLOCK_MAPPER = None,
FLOAT_TO_FP8_ROW_MAPPER = None,
):
model_name = str(model_name)
lower_model_name = model_name.lower()
if not SUPPORTS_FOURBIT and lower_model_name in INT_TO_FLOAT_MAPPER:
assert load_in_fp8 in (True, False, "block")
if load_in_fp8 != False:
if load_in_fp8 == True and (os.environ.get("UNSLOTH_HAS_FBGEMM", "0") == "1"):
if lower_model_name in FLOAT_TO_FP8_ROW_MAPPER:
# Faster row scaling only works if FBGEMM works!
return FLOAT_TO_FP8_ROW_MAPPER[lower_model_name]
elif lower_model_name in FLOAT_TO_FP8_BLOCK_MAPPER:
# Otherwise we use the slower blockwise type
return FLOAT_TO_FP8_BLOCK_MAPPER[lower_model_name]
else:
if lower_model_name in FLOAT_TO_FP8_BLOCK_MAPPER:
return FLOAT_TO_FP8_BLOCK_MAPPER[lower_model_name]
return None
elif not SUPPORTS_FOURBIT and lower_model_name in INT_TO_FLOAT_MAPPER:
model_name = INT_TO_FLOAT_MAPPER[lower_model_name]
print(
f"Unsloth: Your transformers version of {transformers_version} does not support native "
@ -103,13 +135,17 @@ def _get_new_mapper():
return {}, {}, {}
def get_model_name(model_name, load_in_4bit = True):
def get_model_name(model_name, load_in_4bit = True, load_in_fp8 = False):
assert load_in_fp8 in (True, False, "block")
new_model_name = __get_model_name(
model_name = model_name,
load_in_4bit = load_in_4bit,
INT_TO_FLOAT_MAPPER = INT_TO_FLOAT_MAPPER,
FLOAT_TO_INT_MAPPER = FLOAT_TO_INT_MAPPER,
MAP_TO_UNSLOTH_16bit = MAP_TO_UNSLOTH_16bit,
load_in_fp8 = load_in_fp8,
FLOAT_TO_FP8_BLOCK_MAPPER = FLOAT_TO_FP8_BLOCK_MAPPER,
FLOAT_TO_FP8_ROW_MAPPER = FLOAT_TO_FP8_ROW_MAPPER,
)
# In the rare case, we convert bad model names to other names
# For eg too large dynamic quants or MoEs
@ -135,6 +171,9 @@ def get_model_name(model_name, load_in_4bit = True):
INT_TO_FLOAT_MAPPER = NEW_INT_TO_FLOAT_MAPPER,
FLOAT_TO_INT_MAPPER = NEW_FLOAT_TO_INT_MAPPER,
MAP_TO_UNSLOTH_16bit = NEW_MAP_TO_UNSLOTH_16bit,
load_in_fp8 = load_in_fp8,
FLOAT_TO_FP8_BLOCK_MAPPER = FLOAT_TO_FP8_BLOCK_MAPPER,
FLOAT_TO_FP8_ROW_MAPPER = FLOAT_TO_FP8_ROW_MAPPER,
)
if upgraded_model_name is not None:
raise NotImplementedError(
@ -143,4 +182,194 @@ def get_model_name(model_name, load_in_4bit = True):
'pip install --upgrade --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"\n'
'pip install --upgrade --no-cache-dir "git+https://github.com/unslothai/unsloth-zoo.git"\n'
)
if load_in_fp8 != False:
# Handle on the fly TorchAO FP8 quantization
return new_model_name
return new_model_name if new_model_name is not None else model_name
def _get_torchao_fp8_config(fp8_mode: str):
"""
Return a `torchao.quantization.Float8DynamicActivationFloat8WeightConfig`
to be used for `load_in_fp8=True`.
"""
from torchao.quantization import (
Float8DynamicActivationFloat8WeightConfig,
PerBlock,
PerRow,
)
if fp8_mode == "row":
granularity = PerRow()
elif fp8_mode == "block":
granularity = (PerBlock([1, 128]), PerBlock([128, 128]))
else:
raise ValueError("Unsloth: `load_in_fp8` supports only 'row' or 'block'")
return Float8DynamicActivationFloat8WeightConfig(
granularity = granularity,
activation_value_lb = 1e-12,
)
def _offline_quantize_to_fp8(model_name: str, fp8_mode: str) -> str:
"""
Quantizes the model to fp8 using torchao and saving the quantized model to a
temporary location. Return the path to the quantized model.
Note: Once on-the-fly quantization is added in vllm in
https://github.com/vllm-project/vllm/pull/26327, we should
dynamically quantize the model there instead:
llm = LLM(
...
hf_overrides={"quantization_config_file": "torchao_config.json"},
)
"""
temp_dir = tempfile.gettempdir()
new_model_name = model_name.split("/")[-1] + "-fp8-" + fp8_mode
new_model_name = os.path.join(temp_dir, new_model_name)
print(
f"Unsloth: Quantizing '{model_name}' to fp8, using model_name='{new_model_name}' instead"
)
if not os.path.isdir(new_model_name):
from transformers import (
AutoModelForCausalLM,
AutoModelForImageTextToText,
AutoTokenizer,
AutoProcessor,
TorchAoConfig,
AutoConfig,
)
qconfig = _get_torchao_fp8_config(fp8_mode)
qconfig = TorchAoConfig(qconfig)
config = AutoConfig.from_pretrained(model_name)
is_vlm = any(
x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
for x in config.architectures
)
is_vlm = is_vlm or hasattr(config, "vision_config")
auto_model = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM
auto_processor = AutoProcessor if is_vlm else AutoTokenizer
model = auto_model.from_pretrained(
model_name,
torch_dtype = "auto",
device_map = "auto",
quantization_config = qconfig,
)
tokenizer = auto_processor.from_pretrained(model_name)
model.save_pretrained(new_model_name, safe_serialization = False)
del model
for _ in range(2):
torch.cuda.empty_cache()
gc.collect()
tokenizer.save_pretrained(new_model_name)
return new_model_name
def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
"""
Tag a model with a `TorchAOConfig` so downstream callers will know what to do with it.
"""
try:
base_config = _get_torchao_fp8_config(fp8_mode)
model.torchao_config = TorchAOConfig(
qat_scheme = None,
base_config_and_filter_fns = [(base_config, None)],
)
except:
pass
def _get_fp8_mode_and_check_settings(
load_in_fp8: Union[bool, str],
fast_inference: bool,
full_finetuning: bool,
load_in_4bit: bool,
load_in_8bit: bool,
load_in_16bit: bool,
use_exact_model_name: bool,
) -> str:
"""
Assuming `load_in_fp8` is enabled, raise appropriate errors on incompatible settings
and environment. Currently this feature requires:
1. H100 GPUs or after
2. torchao 0.15.0+ (or nightly)
3. torch 2.9.0+
4. If fbgemm_gpu_genai is installed, require 1.4.1+
Returns the fp8 mode, one of "row" or "block".
"""
assert load_in_fp8 is not False
if load_in_fp8 is True:
fp8_mode = "row" # default
else:
fp8_mode = load_in_fp8
# Check user settings
if fp8_mode not in ["row", "block"]:
raise ValueError(
f"Unsloth: `load_in_fp8` can only be 'row' or 'block', got '{fp8_mode}'"
)
if not fast_inference:
raise ValueError(
"Unsloth: `load_in_fp8` is only supported for `fast_inference` for now"
)
if full_finetuning:
raise ValueError(
"Unsloth: `load_in_fp8` is not compatible with full finetuning"
)
if load_in_4bit or load_in_8bit or load_in_16bit:
raise ValueError(
"Unsloth: `load_in_fp8` is not compatible with `load_in_4bit`, `load_in_8bit` or `load_in_16bit`",
)
if use_exact_model_name:
raise ValueError("Unsloth: `load_in_fp8` requires `use_exact_model_name=False`")
# Check if this is Hopper or above
if not (
torch.cuda.is_available()
and torch.version.cuda
and torch.cuda.get_device_capability() >= (9, 0)
):
raise ValueError(
"Unsloth: On the fly `load_in_fp8` requires H100 GPUs or after. Try `unsloth/Qwen3-8B` instead."
)
# Check if torch >= 2.9.0
if Version(torch.__version__) < Version("2.9.0"):
raise ValueError(
"Unsloth: On the fly `load_in_fp8` requires torch 2.9.0+. Try `unsloth/Qwen3-8B` instead."
)
# Check if torchao has this PR: https://github.com/pytorch/ao/pull/3158,
# which will be released in 0.15.0.
if importlib.util.find_spec("torchao") is None:
raise ValueError(
"Unsloth: Please install torchao for on the fly float8 to work! Try `unsloth/Qwen3-8B` instead."
)
import torchao
error_message = (
"Unsloth: `load_in_fp8` requires torchao 0.15.0+ (or nightly).\n"
f"You have torchao version={torchao.__version__}\n"
"Use `pip install --upgrade --force-reinstall torchao`"
)
if Version(torchao.__version__) < Version("0.15.0"):
raise ValueError(error_message)
# If fbgemm_gpu_genai is installed, check if it's >= 1.4.1
if (
importlib.util.find_spec("fbgemm_gpu") is not None
and importlib.util.find_spec("fbgemm_gpu.experimental") is not None
):
import fbgemm_gpu.experimental.gen_ai
if Version(fbgemm_gpu.__version__) < Version("1.4.1"):
raise ValueError(
"Unsloth: On the fly `load_in_fp8` is only compatible with fbgemm_gpu_genai 1.4.1+. Try `unsloth/Qwen3-8B` instead."
)
return fp8_mode

View file

@ -15,6 +15,9 @@
__all__ = [
"INT_TO_FLOAT_MAPPER",
"FLOAT_TO_INT_MAPPER",
"MAP_TO_UNSLOTH_16bit",
"FLOAT_TO_FP8_BLOCK_MAPPER",
"FLOAT_TO_FP8_ROW_MAPPER",
]
__INT_TO_FLOAT_MAPPER = \
@ -233,21 +236,35 @@ __INT_TO_FLOAT_MAPPER = \
"meta-llama/Meta-Llama-3.1-8B",
"unsloth/Meta-Llama-3.1-8B-bnb-4bit",
),
"unsloth/Meta-Llama-3.1-8B-Instruct-unsloth-bnb-4bit" : (
"unsloth/Meta-Llama-3.1-8B-Instruct",
"meta-llama/Meta-Llama-3.1-8B-Instruct",
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
),
"unsloth/Meta-Llama-3.1-8B-Instruct-unsloth-bnb-4bit" : {
"8" : (
"RedHatAI/Llama-3.1-8B-Instruct-FP8",
"unsloth/Llama-3.1-8B-Instruct-FP8-Block",
"unsloth/Llama-3.1-8B-Instruct-FP8-Dynamic",
),
"16" : (
"unsloth/Meta-Llama-3.1-8B-Instruct",
"meta-llama/Meta-Llama-3.1-8B-Instruct",
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
),
},
"unsloth/Llama-3.1-8B-unsloth-bnb-4bit" : (
"unsloth/Llama-3.1-8B",
"meta-llama/Llama-3.1-8B",
"unsloth/Llama-3.1-8B-bnb-4bit",
),
"unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit" : (
"unsloth/Llama-3.1-8B-Instruct",
"meta-llama/Llama-3.1-8B-Instruct",
"unsloth/Llama-3.1-8B-Instruct-bnb-4bit",
),
"unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit" : {
"8" : (
"RedHatAI/Llama-3.1-8B-Instruct-FP8",
"unsloth/Llama-3.1-8B-Instruct-FP8-Block",
"unsloth/Llama-3.1-8B-Instruct-FP8-Dynamic",
),
"16" : (
"unsloth/Llama-3.1-8B-Instruct",
"meta-llama/Llama-3.1-8B-Instruct",
"unsloth/Llama-3.1-8B-Instruct-bnb-4bit",
),
},
"unsloth/Meta-Llama-3.1-70B-bnb-4bit" : (
"unsloth/Meta-Llama-3.1-70B",
"meta-llama/Meta-Llama-3.1-70B",
@ -474,16 +491,30 @@ __INT_TO_FLOAT_MAPPER = \
"meta-llama/Llama-3.2-3B",
"unsloth/Llama-3.2-3B-bnb-4bit",
),
"unsloth/Llama-3.2-1B-Instruct-unsloth-bnb-4bit" : (
"unsloth/Llama-3.2-1B-Instruct",
"meta-llama/Llama-3.2-1B-Instruct",
"unsloth/Llama-3.2-1B-Instruct-bnb-4bit",
),
"unsloth/Llama-3.2-3B-Instruct-unsloth-bnb-4bit" : (
"unsloth/Llama-3.2-3B-Instruct",
"meta-llama/Llama-3.2-3B-Instruct",
"unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
),
"unsloth/Llama-3.2-1B-Instruct-unsloth-bnb-4bit" : {
"8": (
"RedHatAI/Llama-3.2-1B-Instruct-FP8",
"unsloth/Llama-3.2-1B-Instruct-FP8-Block",
"unsloth/Llama-3.2-1B-Instruct-FP8-Dynamic",
),
"16" : (
"unsloth/Llama-3.2-1B-Instruct",
"meta-llama/Llama-3.2-1B-Instruct",
"unsloth/Llama-3.2-1B-Instruct-bnb-4bit",
),
},
"unsloth/Llama-3.2-3B-Instruct-unsloth-bnb-4bit" : {
"8": (
"RedHatAI/Llama-3.2-3B-Instruct-FP8",
"unsloth/Llama-3.2-3B-Instruct-FP8-Block",
"unsloth/Llama-3.2-3B-Instruct-FP8-Dynamic",
),
"16" : (
"unsloth/Llama-3.2-3B-Instruct",
"meta-llama/Llama-3.2-3B-Instruct",
"unsloth/Llama-3.2-3B-Instruct-bnb-4bit",
),
},
"unsloth/Llama-3.1-Nemotron-70B-Instruct-bnb-4bit" : (
"unsloth/Llama-3.1-Nemotron-70B-Instruct",
"nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
@ -561,10 +592,18 @@ __INT_TO_FLOAT_MAPPER = \
"unsloth/QwQ-32B-Preview",
"Qwen/QwQ-32B-Preview",
),
"unsloth/Llama-3.3-70B-Instruct-bnb-4bit" : (
"unsloth/Llama-3.3-70B-Instruct",
"meta-llama/Llama-3.3-70B-Instruct",
),
"unsloth/Llama-3.3-70B-Instruct-unsloth-bnb-4bit" : {
"8" : (
"RedHatAI/Llama-3.3-70B-Instruct-FP8",
"unsloth/Llama-3.3-70B-Instruct-FP8-Block",
"unsloth/Llama-3.3-70B-Instruct-FP8-Dynamic",
),
"16" : (
"unsloth/Llama-3.3-70B-Instruct",
"meta-llama/Llama-3.3-70B-Instruct",
"unsloth/Llama-3.3-70B-Instruct-bnb-4bit",
),
},
"unsloth/phi-4-unsloth-bnb-4bit" : (
"unsloth/phi-4",
"microsoft/phi-4",
@ -733,36 +772,78 @@ __INT_TO_FLOAT_MAPPER = \
"mistralai/Mistral-Small-3.1-24B-Base-2503",
"unsloth/Mistral-Small-3.1-24B-Base-2503-bnb-4bit",
),
"unsloth/Qwen3-0.6B-unsloth-bnb-4bit" : (
"unsloth/Qwen3-0.6B",
"Qwen/Qwen3-0.6B",
"unsloth/Qwen3-0.6B-bnb-4bit",
),
"unsloth/Qwen3-1.7B-unsloth-bnb-4bit" : (
"unsloth/Qwen3-1.7B",
"Qwen/Qwen3-1.7B",
"unsloth/Qwen3-1.7B-bnb-4bit",
),
"unsloth/Qwen3-4B-unsloth-bnb-4bit" : (
"unsloth/Qwen3-4B",
"Qwen/Qwen3-4B",
"unsloth/Qwen3-4B-bnb-4bit",
),
"unsloth/Qwen3-8B-unsloth-bnb-4bit" : (
"unsloth/Qwen3-8B",
"Qwen/Qwen3-8B",
"unsloth/Qwen3-8B-bnb-4bit",
),
"unsloth/Qwen3-14B-unsloth-bnb-4bit" : (
"unsloth/Qwen3-14B",
"Qwen/Qwen3-14B",
"unsloth/Qwen3-14B-bnb-4bit",
),
"unsloth/Qwen3-32B-unsloth-bnb-4bit" : (
"unsloth/Qwen3-32B",
"Qwen/Qwen3-32B",
"unsloth/Qwen3-32B-bnb-4bit",
),
"unsloth/Qwen3-0.6B-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-0.6B-FP8",
"unsloth/Qwen3-0.6B-FP8",
"unsloth/Qwen3-0.6B-FP8",
),
"16" : (
"unsloth/Qwen3-0.6B",
"Qwen/Qwen3-0.6B",
"unsloth/Qwen3-0.6B-bnb-4bit",
),
},
"unsloth/Qwen3-1.7B-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-1.7B-FP8",
"unsloth/Qwen3-1.7B-FP8",
"unsloth/Qwen3-1.7B-FP8",
),
"16" : (
"unsloth/Qwen3-1.7B",
"Qwen/Qwen3-1.7B",
"unsloth/Qwen3-1.7B-bnb-4bit",
),
},
"unsloth/Qwen3-4B-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-4B-FP8",
"unsloth/Qwen3-4B-FP8",
"unsloth/Qwen3-4B-FP8",
),
"16" : (
"unsloth/Qwen3-4B",
"Qwen/Qwen3-4B",
"unsloth/Qwen3-4B-bnb-4bit",
),
},
"unsloth/Qwen3-8B-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-8B-FP8",
"unsloth/Qwen3-8B-FP8",
"unsloth/Qwen3-8B-FP8",
),
"16" : (
"unsloth/Qwen3-8B",
"Qwen/Qwen3-8B",
"unsloth/Qwen3-8B-bnb-4bit",
),
},
"unsloth/Qwen3-14B-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-14B-FP8",
"unsloth/Qwen3-14B-FP8",
"unsloth/Qwen3-14B-FP8",
),
"16" : (
"unsloth/Qwen3-14B",
"Qwen/Qwen3-14B",
"unsloth/Qwen3-14B-bnb-4bit",
),
},
"unsloth/Qwen3-32B-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-32B-FP8",
"unsloth/Qwen3-32B-FP8",
"unsloth/Qwen3-32B-FP8",
),
"16" : (
"unsloth/Qwen3-32B",
"Qwen/Qwen3-32B",
"unsloth/Qwen3-32B-bnb-4bit",
),
},
"unsloth/Qwen3-30B-A3B-unsloth-bnb-4bit" : (
"unsloth/Qwen3-30B-A3B",
"Qwen/Qwen3-30B-A3B",
@ -936,16 +1017,30 @@ __INT_TO_FLOAT_MAPPER = \
"openai/gpt-oss-120b",
"unsloth/gpt-oss-120b-unsloth-bnb-4bit",
),
"unsloth/Qwen3-4B-Instruct-2507-unsloth-bnb-4bit" : (
"unsloth/Qwen3-4B-Instruct-2507",
"Qwen/Qwen3-4B-Instruct-2507",
"unsloth/Qwen3-4B-Instruct-2507-bnb-4bit",
),
"unsloth/Qwen3-4B-Thinking-2507-unsloth-bnb-4bit" : (
"unsloth/Qwen3-4B-Thinking-2507",
"Qwen/Qwen3-4B-Thinking-2507",
"unsloth/Qwen3-4B-Thinking-2507-bnb-4bit",
),
"unsloth/Qwen3-4B-Instruct-2507-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-4B-Instruct-2507-FP8",
"unsloth/Qwen3-4B-Instruct-2507-FP8",
"unsloth/Qwen3-4B-Instruct-2507-FP8",
),
"16" : (
"unsloth/Qwen3-4B-Instruct-2507",
"Qwen/Qwen3-4B-Instruct-2507",
"unsloth/Qwen3-4B-Instruct-2507-bnb-4bit",
),
},
"unsloth/Qwen3-4B-Thinking-2507-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-4B-Thinking-2507-FP8",
"unsloth/Qwen3-4B-Thinking-2507-FP8",
"unsloth/Qwen3-4B-Thinking-2507-FP8",
),
"16" : (
"unsloth/Qwen3-4B-Thinking-2507",
"Qwen/Qwen3-4B-Thinking-2507",
"unsloth/Qwen3-4B-Thinking-2507-bnb-4bit",
),
},
"unsloth/gemma-3-270m-it-unsloth-bnb-4bit" : (
"unsloth/gemma-3-270m-it",
"google/gemma-3-270m-it",
@ -1012,83 +1107,139 @@ __INT_TO_FLOAT_MAPPER = \
"unsloth/granite-4.0-h-small-base",
"ibm-granite/granite-4.0-h-small-base",
),
"unsloth/Qwen3-VL-4B-Thinking-unsloth-bnb-4bit" : (
"unsloth/Qwen3-VL-4B-Thinking",
"Qwen/Qwen3-VL-4B-Thinking",
"unsloth/Qwen3-VL-4B-Thinking-bnb-4bit",
),
"unsloth/Qwen3-VL-8B-Thinking-unsloth-bnb-4bit" : (
"unsloth/Qwen3-VL-8B-Thinking",
"Qwen/Qwen3-VL-8B-Thinking",
"unsloth/Qwen3-VL-8B-Thinking-bnb-4bit",
),
"unsloth/Qwen3-VL-4B-Instruct-unsloth-bnb-4bit" : (
"unsloth/Qwen3-VL-4B-Instruct",
"Qwen/Qwen3-VL-4B-Instruct",
"unsloth/Qwen3-VL-4B-Instruct-bnb-4bit",
),
"unsloth/Qwen3-VL-8B-Instruct-unsloth-bnb-4bit" : (
"unsloth/Qwen3-VL-8B-Instruct",
"Qwen/Qwen3-VL-8B-Instruct",
"unsloth/Qwen3-VL-8B-Instruct-bnb-4bit",
),
"unsloth/Qwen3-VL-2B-Thinking-unsloth-bnb-4bit" : (
"unsloth/Qwen3-VL-2B-Thinking",
"Qwen/Qwen3-VL-2B-Thinking",
"unsloth/Qwen3-VL-2B-Thinking-bnb-4bit",
),
"unsloth/Qwen3-VL-32B-Thinking-unsloth-bnb-4bit" : (
"unsloth/Qwen3-VL-32B-Thinking",
"Qwen/Qwen3-VL-32B-Thinking",
"unsloth/Qwen3-VL-32B-Thinking-bnb-4bit",
),
"unsloth/Qwen3-VL-2B-Instruct-unsloth-bnb-4bit" : (
"unsloth/Qwen3-VL-2B-Instruct",
"Qwen/Qwen3-VL-2B-Instruct",
"unsloth/Qwen3-VL-2B-Instruct-bnb-4bit",
),
"unsloth/Qwen3-VL-32B-Instruct-unsloth-bnb-4bit" : (
"unsloth/Qwen3-VL-32B-Instruct",
"Qwen/Qwen3-VL-32B-Instruct",
"unsloth/Qwen3-VL-32B-Instruct-bnb-4bit",
),
"unsloth/Qwen3-VL-4B-Thinking-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-VL-4B-Thinking-FP8",
"unsloth/Qwen3-VL-4B-Thinking-FP8",
"unsloth/Qwen3-VL-4B-Thinking-FP8",
),
"16" : (
"unsloth/Qwen3-VL-4B-Thinking",
"Qwen/Qwen3-VL-4B-Thinking",
"unsloth/Qwen3-VL-4B-Thinking-bnb-4bit",
),
},
"unsloth/Qwen3-VL-8B-Thinking-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-VL-8B-Thinking-FP8",
"unsloth/Qwen3-VL-8B-Thinking-FP8",
"unsloth/Qwen3-VL-8B-Thinking-FP8",
),
"16" : (
"unsloth/Qwen3-VL-8B-Thinking",
"Qwen/Qwen3-VL-8B-Thinking",
"unsloth/Qwen3-VL-8B-Thinking-bnb-4bit",
),
},
"unsloth/Qwen3-VL-4B-Instruct-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-VL-4B-Instruct-FP8",
"unsloth/Qwen3-VL-4B-Instruct-FP8",
"unsloth/Qwen3-VL-4B-Instruct-FP8",
),
"16" : (
"unsloth/Qwen3-VL-4B-Instruct",
"Qwen/Qwen3-VL-4B-Instruct",
"unsloth/Qwen3-VL-4B-Instruct-bnb-4bit",
),
},
"unsloth/Qwen3-VL-8B-Instruct-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-VL-8B-Instruct-FP8",
"unsloth/Qwen3-VL-8B-Instruct-FP8",
"unsloth/Qwen3-VL-8B-Instruct-FP8",
),
"16" : (
"unsloth/Qwen3-VL-8B-Instruct",
"Qwen/Qwen3-VL-8B-Instruct",
"unsloth/Qwen3-VL-8B-Instruct-bnb-4bit",
),
},
"unsloth/Qwen3-VL-2B-Thinking-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-VL-2B-Thinking-FP8",
"unsloth/Qwen3-VL-2B-Thinking-FP8",
"unsloth/Qwen3-VL-2B-Thinking-FP8",
),
"16" : (
"unsloth/Qwen3-VL-2B-Thinking",
"Qwen/Qwen3-VL-2B-Thinking",
"unsloth/Qwen3-VL-2B-Thinking-bnb-4bit",
),
},
"unsloth/Qwen3-VL-32B-Thinking-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-VL-32B-Thinking-FP8",
"unsloth/Qwen3-VL-32B-Thinking-FP8",
"unsloth/Qwen3-VL-32B-Thinking-FP8",
),
"16" : (
"unsloth/Qwen3-VL-32B-Thinking",
"Qwen/Qwen3-VL-32B-Thinking",
"unsloth/Qwen3-VL-32B-Thinking-bnb-4bit",
),
},
"unsloth/Qwen3-VL-2B-Instruct-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-VL-2B-Instruct-FP8",
"unsloth/Qwen3-VL-2B-Instruct-FP8",
"unsloth/Qwen3-VL-2B-Instruct-FP8",
),
"16" : (
"unsloth/Qwen3-VL-2B-Instruct",
"Qwen/Qwen3-VL-2B-Instruct",
"unsloth/Qwen3-VL-2B-Instruct-bnb-4bit",
),
},
"unsloth/Qwen3-VL-32B-Instruct-unsloth-bnb-4bit" : {
"8" : (
"Qwen/Qwen3-VL-32B-Instruct-FP8",
"unsloth/Qwen3-VL-32B-Instruct-FP8",
"unsloth/Qwen3-VL-32B-Instruct-FP8",
),
"16" : (
"unsloth/Qwen3-VL-32B-Instruct",
"Qwen/Qwen3-VL-32B-Instruct",
"unsloth/Qwen3-VL-32B-Instruct-bnb-4bit",
),
},
"unsloth/granite-4.0-350m-base-unsloth-bnb-4bit" : (
"unsloth/granite-4.0-350m-base-unsloth-bnb-4bit",
"unsloth/granite-4.0-350m-base",
"ibm-granite/granite-4.0-350m-base",
"unsloth/granite-4.0-350m-base-bnb-4bit",
),
"unsloth/granite-4.0-350m-unsloth-bnb-4bit" : (
"unsloth/granite-4.0-350m-unsloth-bnb-4bit",
"unsloth/granite-4.0-350m",
"ibm-granite/granite-4.0-350m",
"unsloth/granite-4.0-350m-bnb-4bit",
),
"unsloth/granite-4.0-h-350m-base-unsloth-bnb-4bit" : (
"unsloth/granite-4.0-h-350m-base-unsloth-bnb-4bit",
"unsloth/granite-4.0-h-350m-base",
"ibm-granite/granite-4.0-h-350m-base",
"unsloth/granite-4.0-h-350m-base-bnb-4bit",
),
"unsloth/granite-4.0-h-350m-unsloth-bnb-4bit" : (
"unsloth/granite-4.0-h-350m-unsloth-bnb-4bit",
"unsloth/granite-4.0-h-350m",
"ibm-granite/granite-4.0-h-350m",
"unsloth/granite-4.0-h-350m-bnb-4bit",
),
"unsloth/granite-4.0-1b-base-unsloth-bnb-4bit" : (
"unsloth/granite-4.0-1b-base-unsloth-bnb-4bit",
"unsloth/granite-4.0-1b-base",
"ibm-granite/granite-4.0-1b-base",
"unsloth/granite-4.0-1b-base-bnb-4bit",
),
"unsloth/granite-4.0-1b-unsloth-bnb-4bit" : (
"unsloth/granite-4.0-1b-unsloth-bnb-4bit",
"unsloth/granite-4.0-1b",
"ibm-granite/granite-4.0-1b",
"unsloth/granite-4.0-1b-bnb-4bit",
),
"unsloth/granite-4.0-h-1b-base-unsloth-bnb-4bit" : (
"unsloth/granite-4.0-h-1b-base-unsloth-bnb-4bit",
"unsloth/granite-4.0-h-1b-base",
"ibm-granite/granite-4.0-h-1b-base",
"unsloth/granite-4.0-h-1b-base-bnb-4bit",
),
"unsloth/granite-4.0-h-1b-unsloth-bnb-4bit" : (
"unsloth/granite-4.0-h-1b-unsloth-bnb-4bit",
"unsloth/granite-4.0-h-1b",
"ibm-granite/granite-4.0-h-1b",
"unsloth/granite-4.0-h-1b-bnb-4bit",
),
@ -1105,8 +1256,32 @@ __INT_TO_FLOAT_MAPPER = \
INT_TO_FLOAT_MAPPER = {}
FLOAT_TO_INT_MAPPER = {}
MAP_TO_UNSLOTH_16bit = {}
FLOAT_TO_FP8_BLOCK_MAPPER = {}
FLOAT_TO_FP8_ROW_MAPPER = {}
for key, values in __INT_TO_FLOAT_MAPPER.items():
block, row = None, None
if type(values) is dict:
assert "16" in values
float16_values = values["16"]
# Float8 and other quantized types
if "8" in values:
float8_values = values["8"]
assert len(float8_values) == 3
official, block, row = float8_values
FLOAT_TO_FP8_BLOCK_MAPPER[key.lower()] = block
FLOAT_TO_FP8_ROW_MAPPER[key.lower()] = row
FLOAT_TO_FP8_BLOCK_MAPPER[official.lower() + "-dynamic"] = block
FLOAT_TO_FP8_ROW_MAPPER[official.lower()] = row
FLOAT_TO_FP8_ROW_MAPPER[official.lower() + "-dynamic"] = row
FLOAT_TO_FP8_BLOCK_MAPPER[float16_values[0]] = block
FLOAT_TO_FP8_BLOCK_MAPPER[float16_values[0].lower()] = block
FLOAT_TO_FP8_ROW_MAPPER[float16_values[0]] = block
FLOAT_TO_FP8_ROW_MAPPER[float16_values[0].lower()] = block
for k in float8_values:
FLOAT_TO_FP8_BLOCK_MAPPER[k.lower()] = block
FLOAT_TO_FP8_ROW_MAPPER[k.lower()] = row
values = float16_values
INT_TO_FLOAT_MAPPER[key] = values[0]
for value in values:
@ -1117,6 +1292,12 @@ for key, values in __INT_TO_FLOAT_MAPPER.items():
if values[0].startswith("unsloth"):
MAP_TO_UNSLOTH_16bit[values[1]] = values[0]
MAP_TO_UNSLOTH_16bit[values[1].lower()] = values[0]
if block is not None:
MAP_TO_UNSLOTH_16bit[block] = values[0]
MAP_TO_UNSLOTH_16bit[block.lower()] = values[0]
if row is not None:
MAP_TO_UNSLOTH_16bit[row] = values[0]
MAP_TO_UNSLOTH_16bit[row.lower()] = values[0]
elif len(values) == 3:
# Dynamic Unsloth quantization
if values[0].startswith("unsloth"):
@ -1124,6 +1305,12 @@ for key, values in __INT_TO_FLOAT_MAPPER.items():
MAP_TO_UNSLOTH_16bit[values[1].lower()] = values[0]
MAP_TO_UNSLOTH_16bit[values[2]] = values[0]
MAP_TO_UNSLOTH_16bit[values[2].lower()] = values[0]
if block is not None:
MAP_TO_UNSLOTH_16bit[block] = values[0]
MAP_TO_UNSLOTH_16bit[block.lower()] = values[0]
if row is not None:
MAP_TO_UNSLOTH_16bit[row] = values[0]
MAP_TO_UNSLOTH_16bit[row.lower()] = values[0]
pass
# Get lowercased

View file

@ -229,6 +229,7 @@ from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling
from transformers.training_args import ParallelMode
# Wrap trainer with padding to right and enable training mode
# Also patches W&B since multiple runs must use wandb.finish()
import functools
from types import MethodType
def prepare_for_training_mode(f):
@ -241,6 +242,12 @@ def prepare_for_training_mode(f):
# Return inference mode
if hasattr(self, 'model') and hasattr(self.model, "for_inference"):
self.model.for_inference()
# Patch W&B to enable logging on future runs, otherwise it'll overwrite the first run
try:
import wandb
wandb.finish()
except:
pass
return output
return wrapper
pass

View file

@ -35,6 +35,7 @@ from ..device_type import (
ALLOW_PREQUANTIZED_MODELS,
)
import textwrap
from ._utils import _get_inference_mode_context_manager
RL_EXTRA_ARGS = defaultdict(list)
RL_FUNCTIONS = defaultdict(list)
@ -536,7 +537,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
)
with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype):
with torch.inference_mode():
with _get_inference_mode_context_manager(model):
if pixel_values is None:
attention_mask = input_ids != self.processing_class.pad_token_id
attention_mask = attention_mask.to(attention_mask.dtype)
@ -603,6 +604,9 @@ RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(UnslothEfficientGRPO))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss))
RL_PRE_ITEMS["grpo_trainer"].append(grpo_compute_loss_slow)
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_update_SamplingParams))
RL_PRE_ITEMS["grpo_trainer"].append(
inspect.getsource(_get_inference_mode_context_manager)
)
# Edit _get_per_token_logps to handle mixed precision

View file

@ -618,16 +618,24 @@ class FastBaseModel:
else:
quantizer = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_method]
quantizer_kwargs = {}
# We cannot dequantize since gpt-oss-20b MXFP4 will now be gpt-oss-20b-BF16
if (
load_in_16bit
and "dequantize" in inspect.signature(quantizer).parameters
):
quantizer_kwargs["dequantize"] = True
quantization_config = quantizer.from_dict(
quantization_config, **quantizer_kwargs
)
kwargs["quantization_config"] = quantization_config
if quant_method == "compressed-tensors":
# Ignore these
pass
else:
# We cannot dequantize since gpt-oss-20b MXFP4 will now be gpt-oss-20b-BF16
if (
load_in_16bit
and "dequantize" in inspect.signature(quantizer).parameters
):
quantizer_kwargs["dequantize"] = True
try:
# Sometimes this fails so we wrap it in a try except
quantization_config = quantizer.from_dict(
quantization_config, **quantizer_kwargs
)
except:
pass
kwargs["quantization_config"] = quantization_config
# Check if using forced float32 - we load it in bfloat16, then cast to float16!
torch_dtype = dtype

View file

@ -2769,7 +2769,13 @@ def unsloth_save_pretrained_torchao(
for _ in range(3):
gc.collect()
from transformers import AutoModel, AutoTokenizer, TorchAoConfig
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TorchAoConfig,
AutoModelForImageTextToText,
AutoProcessor,
)
from torchao import quantize_
if torchao_config is None:
@ -2781,14 +2787,25 @@ def unsloth_save_pretrained_torchao(
torchao_config = Int8DynamicActivationInt8WeightConfig()
quantization_config = TorchAoConfig(quant_type = torchao_config)
tokenizer = AutoTokenizer.from_pretrained(arguments["save_directory"])
is_vlm = False
if hasattr(self, "config") and hasattr(self.config, "architectures"):
is_vlm = any(
x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
for x in self.config.architectures
)
is_vlm = is_vlm or hasattr(self.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"])
# TorchAO must only use bfloat16 for loading (float16 fails)
if HAS_TORCH_DTYPE:
kwargs = {"torch_dtype": torch.bfloat16}
else:
kwargs = {"dtype": torch.bfloat16}
model = AutoModel.from_pretrained(
model = auto_model.from_pretrained(
arguments["save_directory"],
device_map = "auto",
quantization_config = quantization_config,
@ -2812,6 +2829,13 @@ def unsloth_save_pretrained_torchao(
torchao_save_directory, safe_serialization = safe_serialization
)
tokenizer.save_pretrained(torchao_save_directory)
if os.path.exists(save_directory):
try:
import shutil
shutil.rmtree(save_directory)
except:
pass
for _ in range(3):
gc.collect()