mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
* 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
63 lines
2.5 KiB
Python
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
|