From a14b032d79ddd8f10bd874e741d3909945d1163f Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Fri, 17 Jul 2026 16:30:45 -0700 Subject: [PATCH] Propagate fp8 block_size before the early return in get_lora_parameters_bias (#7189) * Propagate fp8 block_size before the early return in get_lora_parameters_bias get_lora_parameters_bias set the fp8 block_size on W/W_quant only after the disable_adapters/merged early return, so on the merged or disabled path (merged inference, DPO reference model) a block-fp8 weight lost its real block_size and downstream fp8 kernels fell back to [128, 128]. The non-bias sibling get_lora_parameters already sets block_size before its early return; move the block so both behave the same. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard the fp8 block_size against a missing quant state A decompressed compressed-tensors layer keeps quant_method == "fp8" while its weight is back to bf16, so it has no quant state and get_lora_parameters_bias must still return W_quant None for fast_linear_forward to fall back to a plain matmul. Only attach block_size when a quant state was actually found. * Guard the sibling get_lora_parameters fp8 block_size against a missing quant state Mirror the get_lora_parameters_bias guard so a decompressed compressed-tensors layer (quant_method fp8, bf16 weight, no quant state) does not raise AttributeError on the fused-LoRA path. Add a CPU-only regression test. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- ...get_lora_parameters_bias_fp8_block_size.py | 82 +++++++++++++++++++ ...test_get_lora_parameters_fp8_block_size.py | 81 ++++++++++++++++++ unsloth/kernels/utils.py | 18 ++-- 3 files changed, 175 insertions(+), 6 deletions(-) create mode 100644 tests/python/test_get_lora_parameters_bias_fp8_block_size.py create mode 100644 tests/python/test_get_lora_parameters_fp8_block_size.py diff --git a/tests/python/test_get_lora_parameters_bias_fp8_block_size.py b/tests/python/test_get_lora_parameters_bias_fp8_block_size.py new file mode 100644 index 0000000000..835ad19542 --- /dev/null +++ b/tests/python/test_get_lora_parameters_bias_fp8_block_size.py @@ -0,0 +1,82 @@ +import ast +from pathlib import Path + + +def _load_function(name): + # Extract a function from kernels/utils.py without importing unsloth (which + # needs a GPU / torch / bitsandbytes). The functions exercised here only use + # getattr and the _FP8_WEIGHT_DTYPES name on the paths under test. + source = Path(__file__).parents[2] / "unsloth" / "kernels" / "utils.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + funcs = [ + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name + ] + assert len(funcs) == 1, (name, funcs) + namespace = {"getattr": getattr, "_FP8_WEIGHT_DTYPES": ()} + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace[name] + + +class _Obj: + pass + + +def _make_disabled_block_fp8_proj(block_size): + # A merged/disabled projection whose base layer is a block-fp8 weight that + # ships a non-default block size on its checkpoint. + weight = _Obj() + weight.quant_state = _Obj() + base_layer = _Obj() + base_layer.weight = weight + base_layer.quant_method = "fp8" + base_layer.block_size = block_size + base_layer.bias = None + proj = _Obj() + proj.base_layer = base_layer + proj.merged = True + proj.disable_adapters = True + return proj, weight.quant_state + + +def test_bias_variant_propagates_fp8_block_size_on_disabled_path(): + # Downstream fp8 kernels read getattr(weight_scale, "block_size", [128, 128]), + # so the checkpoint's real block size must survive the merged/disabled path, + # exactly as it does for the non-bias sibling get_lora_parameters. + get_lora_parameters_bias = _load_function("get_lora_parameters_bias") + + proj, weight_scale = _make_disabled_block_fp8_proj([64, 128]) + get_lora_parameters_bias(proj) + + assert getattr(weight_scale, "block_size", [128, 128]) == [64, 128] + + +def _make_decompressed_merged_proj(): + # A merged compressed-tensors layer that was decompressed back to bf16. It keeps + # quant_method == "fp8" from the checkpoint metadata, but the live weight is bf16 + # so there is no quant state to attach a block size to. + weight = _Obj() + weight.dtype = "bfloat16" + base_layer = _Obj() + base_layer.weight = weight + base_layer.quant_method = "fp8" + base_layer.block_size = [128, 128] + base_layer.bias = None + proj = _Obj() + proj.base_layer = base_layer + proj.merged = True + proj.disable_adapters = True + return proj + + +def test_bias_variant_keeps_none_quant_state_for_decompressed_layer(): + # Such a layer has no quant state, and fast_linear_forward relies on getting + # W_quant None back so it can fall back to a plain matmul, so setting the block + # size must not assume a quant state is present. + get_lora_parameters_bias = _load_function("get_lora_parameters_bias") + + W, W_quant = get_lora_parameters_bias(_make_decompressed_merged_proj())[:2] + + assert W_quant is None + assert getattr(W, "block_size", None) == [128, 128] diff --git a/tests/python/test_get_lora_parameters_fp8_block_size.py b/tests/python/test_get_lora_parameters_fp8_block_size.py new file mode 100644 index 0000000000..f5f1359125 --- /dev/null +++ b/tests/python/test_get_lora_parameters_fp8_block_size.py @@ -0,0 +1,81 @@ +import ast +from pathlib import Path + + +def _load_function(name): + # Extract a function from kernels/utils.py without importing unsloth (which + # needs a GPU / torch / bitsandbytes). get_lora_parameters only uses getattr, + # hasattr and the _FP8_WEIGHT_DTYPES name on the paths under test. + source = Path(__file__).parents[2] / "unsloth" / "kernels" / "utils.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + funcs = [ + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name + ] + assert len(funcs) == 1, (name, funcs) + namespace = {"getattr": getattr, "hasattr": hasattr, "_FP8_WEIGHT_DTYPES": ()} + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace[name] + + +class _Obj: + pass + + +def _make_disabled_block_fp8_proj(block_size): + # A merged/disabled projection whose base layer is a block-fp8 weight that + # ships a non-default block size on its checkpoint. + weight = _Obj() + weight.quant_state = _Obj() + base_layer = _Obj() + base_layer.weight = weight + base_layer.quant_method = "fp8" + base_layer.block_size = block_size + proj = _Obj() + proj.base_layer = base_layer + proj.merged = True + proj.disable_adapters = True + return proj, weight.quant_state + + +def test_propagates_fp8_block_size_on_disabled_path(): + # get_lora_parameters already sets block_size before its early return; downstream + # fp8 kernels read getattr(weight_scale, "block_size", [128, 128]), so the + # checkpoint's real block size must survive the merged/disabled path. + get_lora_parameters = _load_function("get_lora_parameters") + + proj, weight_scale = _make_disabled_block_fp8_proj([64, 128]) + get_lora_parameters(proj) + + assert getattr(weight_scale, "block_size", [128, 128]) == [64, 128] + + +def _make_decompressed_merged_proj(): + # A merged compressed-tensors layer that was decompressed back to bf16. It keeps + # quant_method == "fp8" from the checkpoint metadata, but the live weight is bf16 + # so there is no quant state to attach a block size to. + weight = _Obj() + weight.dtype = "bfloat16" + base_layer = _Obj() + base_layer.weight = weight + base_layer.quant_method = "fp8" + base_layer.block_size = [128, 128] + proj = _Obj() + proj.base_layer = base_layer + proj.merged = True + proj.disable_adapters = True + return proj + + +def test_keeps_none_quant_state_for_decompressed_layer(): + # Mirrors the get_lora_parameters_bias guard: with no quant state, assigning + # W_quant.block_size must not assume one is present, or it raises AttributeError + # on None. fast_lora relies on getting W_quant None back to fall back to a plain + # matmul, so this path must stay crash-free. + get_lora_parameters = _load_function("get_lora_parameters") + + W, W_quant = get_lora_parameters(_make_decompressed_merged_proj())[:2] + + assert W_quant is None + assert getattr(W, "block_size", None) == [128, 128] diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 1b0b5ce12e..ccfedfdef0 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -325,7 +325,10 @@ def get_lora_parameters(proj): if getattr(base_layer, "quant_method", None) == "fp8": # we need to somehow store and pass this information :) W.block_size = getattr(base_layer, "block_size", [128, 128]) - W_quant.block_size = W.block_size + # A decompressed compressed-tensors layer keeps quant_method == "fp8" while its + # weight is back to bf16, so it has no quant state to carry the block size. + if W_quant is not None: + W_quant.block_size = W.block_size # if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged: if getattr(proj, "disable_adapters", True) or proj.merged: @@ -375,14 +378,17 @@ def get_lora_parameters_bias(proj): if W_quant is None: W_quant = getattr(base_layer, "weight_scale", None) - # if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged: - if getattr(proj, "disable_adapters", True) or proj.merged: - return W, W_quant, None, None, None, base_layer.bias - if getattr(base_layer, "quant_method", None) == "fp8": # we need to somehow store and pass this information :) W.block_size = getattr(base_layer, "block_size", [128, 128]) - W_quant.block_size = W.block_size + # A decompressed compressed-tensors layer keeps quant_method == "fp8" while its + # weight is back to bf16, so it has no quant state to carry the block size. + if W_quant is not None: + W_quant.block_size = W.block_size + + # if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged: + if getattr(proj, "disable_adapters", True) or proj.merged: + return W, W_quant, None, None, None, base_layer.bias adapter = getattr(proj, "active_adapters", None) if adapter is None: