unsloth/tests/test_fast_gemv_dispatch.py
Daniel Han cd9d251f15
Fix fast inference crash on compressed-tensors FP8 models (#7025)
* Fix fast_gemv crash on compressed-tensors FP8 models

Loading a compressed-tensors FP8 checkpoint (for example
unsloth/Llama-3.2-1B-Instruct-FP8-Block) with fast_inference=False and
running a forward crashed with 'Parameter object has no attribute absmax'
inside fast_gemv.

A compressed-tensors CompressedLinear exposes an already dequantized bf16
weight at forward time while keeping a weight_scale Parameter. The quant
state resolution in get_lora_parameters/get_lora_parameters_bias fell back
to that weight_scale, so a bf16 weight was routed into the bitsandbytes
fast_gemv/fast_dequantize path, which expects a bitsandbytes QuantState
with an absmax attribute.

Only fall back to weight_scale_inv/weight_scale when the weight is still
fp8. A decompressed bf16 weight then resolves to no quant state and flows
through the normal bf16 path, which already handles bias and the LoRA
backward. Real fp8 and bitsandbytes 4bit weights are unchanged.

* Skip the fast_gemv dispatch test before importing unsloth when bitsandbytes is absent
2026-07-09 04:10:59 -07:00

63 lines
2.5 KiB
Python

# 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