diff --git a/tests/test_fast_gemv_dispatch.py b/tests/test_fast_gemv_dispatch.py new file mode 100644 index 0000000000..7758db2cd9 --- /dev/null +++ b/tests/test_fast_gemv_dispatch.py @@ -0,0 +1,63 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""`get_lora_parameters` must not treat a `weight_scale` as a quant state for a weight that is +already dequantized to bf16 (e.g. a compressed-tensors layer at forward time). Otherwise the +bnb fast_gemv / fast_dequantize path reads a missing `absmax` and crashes. +""" + +from types import SimpleNamespace + +import pytest +import torch + +# unsloth.kernels.utils imports bitsandbytes unconditionally, so skip the whole module up +# front on runners without it (e.g. CPU-only) before importing unsloth, otherwise collection +# errors instead of producing a skip. Any other import error still surfaces as a failure. +pytest.importorskip("bitsandbytes") + +import unsloth # noqa: F401 (sets UNSLOTH_IS_PRESENT before transformers) +from unsloth.kernels.utils import get_lora_parameters_bias, _FP8_WEIGHT_DTYPES + +_FP8 = _FP8_WEIGHT_DTYPES[0] if _FP8_WEIGHT_DTYPES else None + + +def _proj(weight, weight_scale = None): + proj = SimpleNamespace(weight = weight, bias = None, merged = False) + if weight_scale is not None: + proj.weight_scale = weight_scale + return proj + + +def test_bf16_weight_scale_not_used_as_quant_state(): + """A bf16 weight carrying a weight_scale (compressed-tensors) -> quant state must be None.""" + proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16), torch.rand(2, 2)) + W, W_quant = get_lora_parameters_bias(proj)[:2] + assert W_quant is None + + +def test_fp8_weight_keeps_scale(): + """An actual fp8 weight still resolves its weight_scale as the quant state.""" + if _FP8 is None: + pytest.skip("no float8 dtype in this torch build") + scale = torch.rand(2, 2) + proj = _proj(torch.randn(4, 4).to(_FP8), scale) + W, W_quant = get_lora_parameters_bias(proj)[:2] + assert W_quant is scale + + +def test_plain_bf16_has_no_quant_state(): + proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16)) + W, W_quant = get_lora_parameters_bias(proj)[:2] + assert W_quant is None diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 43ed198a4a..1b0b5ce12e 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -282,6 +282,21 @@ def QUANT_STATE(W): return getattr(W, "quant_state", None) +# fp8 weight dtypes. A `weight_scale` / `weight_scale_inv` should only be treated as a +# quant state when the weight itself is still fp8. compressed-tensors layers expose an +# already-dequantized bf16 weight at forward time while keeping a `weight_scale` around; +# reading that as a quant state routes a bf16 weight into the bitsandbytes fast_gemv / +# fast_dequantize path, which then reads a missing `absmax` and crashes. +_FP8_WEIGHT_DTYPES = tuple( + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), + ) + if dtype is not None +) + + def get_lora_parameters(proj): """Return (weight, weight quant_state, lora A, lora B, lora scale). With QAT enabled, also fake-quantizes the base layer and lora weights. @@ -298,9 +313,11 @@ def get_lora_parameters(proj): if weight_fake_quantizer is not None: W = weight_fake_quantizer(W) - # Get quant state for 4bit or FP8 + # Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the + # weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer) + # must not carry a scale as its quant state or fast_gemv will crash on it. W_quant = getattr(W, "quant_state", None) - if W_quant is None: + if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES: W_quant = getattr(base_layer, "weight_scale_inv", None) if W_quant is None: W_quant = getattr(base_layer, "weight_scale", None) @@ -349,9 +366,11 @@ def get_lora_parameters_bias(proj): ) # (proj.base_layer if hasattr(proj, "base_layer") else proj) W = base_layer.weight - # Get quant state for 4bit or FP8 + # Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the + # weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer) + # must not carry a scale as its quant state or fast_gemv will crash on it. W_quant = getattr(W, "quant_state", None) - if W_quant is None: + if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES: W_quant = getattr(base_layer, "weight_scale_inv", None) if W_quant is None: W_quant = getattr(base_layer, "weight_scale", None)