From 38aa148aba47b66dbef54e6e125cce1bc129a13f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 25 Nov 2025 07:23:26 -0800 Subject: [PATCH] Float8 GRPO, RL (#3640) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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} ``` Screenshot 2025-11-11 at 4 10 50 PM 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 * 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 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> --- README.md | 18 +- pyproject.toml | 46 +++- unsloth/__init__.py | 4 +- unsloth/_auto_install.py | 5 +- unsloth/kernels/fp8.py | 45 +++- unsloth/kernels/utils.py | 42 ++- unsloth/models/_utils.py | 24 +- unsloth/models/llama.py | 7 +- unsloth/models/loader.py | 66 ++++- unsloth/models/loader_utils.py | 235 ++++++++++++++++- unsloth/models/mapper.py | 411 ++++++++++++++++++++++-------- unsloth/models/rl.py | 7 + unsloth/models/rl_replacements.py | 6 +- unsloth/models/vision.py | 28 +- unsloth/save.py | 30 ++- 15 files changed, 818 insertions(+), 156 deletions(-) diff --git a/README.md b/README.md index b2d86e2f1d..fd28429f85 100644 --- a/README.md +++ b/README.md @@ -197,9 +197,9 @@ pip install unsloth ### 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. diff --git a/pyproject.toml b/pyproject.toml index 59fc94ee4b..34d46e37e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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'", diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 8b48ce3ba0..c2311a93bf 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -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 diff --git a/unsloth/_auto_install.py b/unsloth/_auto_install.py index f6cf6cfc94..ed9b918412 100644 --- a/unsloth/_auto_install.py +++ b/unsloth/_auto_install.py @@ -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"') \ No newline at end of file +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') \ No newline at end of file diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index 46ebac1bda..435a6e83d2 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -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 diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 4ce3b2e1e9..5dcc7c232c 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -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 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index d63986db49..6ca128bd83 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2025.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() diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6e61b08c93..9e3896244c 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 9d8a2ce9fa..951bcb1e0c 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -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 diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 631f614cb4..7bcddde2d9 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -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 diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index f2dad24225..dcfe890e92 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -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 diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index bfbd0bf103..1a06277d85 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -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 diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 8746399afa..38973fe30e 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -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 diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index c0084f7d55..30d381c791 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -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 diff --git a/unsloth/save.py b/unsloth/save.py index 810bf9062a..0a8f02d90f 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -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()