[Bugfix] compressed-tensors: mixed-precision checkpoints silently load unquantized (#32736)
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
This commit is contained in:
Hert4
co-authored by
Mohammad Miadh Angkad
Mohammad Miadh Angkad
parent
fddfc1fb5e
commit
f69af7b7ad
@@ -305,7 +305,18 @@ class CompressedTensorsConfig(QuantizationConfig):
|
||||
)
|
||||
|
||||
target_scheme_map[target]["input_activations"] = None
|
||||
if is_activation_quantization_format(quant_format):
|
||||
|
||||
# A config_group may carry its own format. When several groups
|
||||
# use different formats, compressed-tensors sets the top-level
|
||||
# format to "mixed-precision" and the real format lives on each
|
||||
# group, so the per-group value must win.
|
||||
group_format = quant_config.get("format")
|
||||
target_scheme_map[target]["format"] = group_format
|
||||
act_quant_format = is_activation_quantization_format(
|
||||
group_format if group_format is not None else quant_format
|
||||
)
|
||||
|
||||
if act_quant_format:
|
||||
input_activations = quant_config.get("input_activations")
|
||||
# When activation quant format is set but no
|
||||
# input_activations provided: valid for w8a16fp8 (FLOAT
|
||||
@@ -572,13 +583,20 @@ class CompressedTensorsConfig(QuantizationConfig):
|
||||
return is_w4 and weight_quant.symmetric and is_token and is_dynamic
|
||||
|
||||
def _get_scheme_from_parts(
|
||||
self, weight_quant: BaseModel, input_quant: BaseModel
|
||||
self,
|
||||
weight_quant: BaseModel,
|
||||
input_quant: BaseModel,
|
||||
format: Optional[str] = None,
|
||||
) -> CompressedTensorsLinearScheme:
|
||||
# The format of the config_group this layer matched, when it declares
|
||||
# one. Falls back to the top-level format, which is "mixed-precision"
|
||||
# (i.e. meaningless on its own) for multi-group checkpoints.
|
||||
quant_format = format if format is not None else self.quant_format
|
||||
|
||||
# Detect If Mixed Precision
|
||||
if self._is_wNa16_group_channel(weight_quant, input_quant):
|
||||
if (
|
||||
self.quant_format == CompressionFormat.pack_quantized.value
|
||||
quant_format == CompressionFormat.pack_quantized.value
|
||||
and weight_quant.num_bits in WNA16_SUPPORTED_BITS
|
||||
):
|
||||
return CompressedTensorsWNA16(
|
||||
@@ -593,7 +611,7 @@ class CompressedTensorsConfig(QuantizationConfig):
|
||||
"Other method (CompressedTensorsW4A16Sparse24) is not supported now"
|
||||
)
|
||||
|
||||
if is_activation_quantization_format(self.quant_format):
|
||||
if is_activation_quantization_format(quant_format):
|
||||
if self._is_fp4a4_nvfp4(weight_quant, input_quant):
|
||||
is_fp4a4_nvfp4_supported = self._check_scheme_supported(
|
||||
CompressedTensorsW4A4Fp4.get_min_capability(), error=False
|
||||
@@ -796,9 +814,11 @@ class CompressedTensorsConfig(QuantizationConfig):
|
||||
scheme_dict = self.get_scheme_dict(layer, layer_name)
|
||||
weight_quant = None
|
||||
input_quant = None
|
||||
scheme_format = None
|
||||
if scheme_dict:
|
||||
weight_quant = scheme_dict.get("weights")
|
||||
input_quant = scheme_dict.get("input_activations")
|
||||
scheme_format = scheme_dict.get("format")
|
||||
|
||||
# Find the sparsity scheme of the layer
|
||||
# assume that fused layers inerhit first component's sparsity scheme
|
||||
@@ -834,6 +854,7 @@ class CompressedTensorsConfig(QuantizationConfig):
|
||||
scheme = self._get_scheme_from_parts( # type: ignore
|
||||
weight_quant=weight_quant,
|
||||
input_quant=input_quant,
|
||||
format=scheme_format,
|
||||
)
|
||||
|
||||
# Raise error if device does not support the scheme
|
||||
|
||||
@@ -80,9 +80,22 @@ def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool:
|
||||
"""
|
||||
Checks whether a layer_name is exactly equal or a regex match for
|
||||
if target starts with 're:' to any target in list.
|
||||
|
||||
A plain (non-regex) target additionally matches a dotted-path *suffix* of
|
||||
layer_name, so a target may be written as a module suffix, e.g.
|
||||
"self_attn.kv_b_proj" matches "model.layers.0.self_attn.kv_b_proj".
|
||||
|
||||
It must never match a *prefix*: llm-compressor writes parent modules into
|
||||
the `ignore` list (e.g. "model.layers.0.mlp.experts.0"), and a plain
|
||||
substring match would let such an entry silently swallow its quantized
|
||||
children ("model.layers.0.mlp.experts.0.gate_proj"), dropping the layer
|
||||
back to an unquantized method.
|
||||
"""
|
||||
for target in targets:
|
||||
if _is_equal_or_regex_match(layer_name, target, check_contains=True):
|
||||
if target.startswith("re:"):
|
||||
if _is_equal_or_regex_match(layer_name, target):
|
||||
return True
|
||||
elif target == layer_name or layer_name.endswith("." + target):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@@ -0,0 +1,192 @@
|
||||
"""CPU regression tests for multi-group ("mixed-precision") compressed-tensors configs.
|
||||
|
||||
Such a checkpoint used to load completely unquantized. ``ignore`` was matched by
|
||||
substring, so a parent module entry swallowed its quantized children, and the
|
||||
activation-quantization gate read the top-level format -- which compressed-tensors
|
||||
sets to ``mixed-precision`` when groups disagree -- dropping ``input_activations``
|
||||
for every group. Both are config-parsing paths, so these tests run on CPU.
|
||||
"""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
|
||||
CompressedTensorsConfig,
|
||||
)
|
||||
from sglang.srt.layers.quantization.compressed_tensors.schemes import (
|
||||
CompressedTensorsWNA16,
|
||||
)
|
||||
from sglang.srt.layers.quantization.compressed_tensors.utils import (
|
||||
check_equal_or_regex_match,
|
||||
should_ignore_layer,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
EXPERTS_LAYER = "model.language_model.layers.0.mlp.experts"
|
||||
GATE_PROJ = f"{EXPERTS_LAYER}.0.gate_proj"
|
||||
MLP_LAYER = "model.language_model.layers.0.mlp.gate_proj"
|
||||
|
||||
FP8_TARGET = "re:.*self_attn\\.(q|k|v|o)_proj$"
|
||||
NVFP4_TARGET = "re:.*mlp\\.experts\\.\\d+\\.(gate|up|down)_proj$"
|
||||
WNA16_TARGET = "re:.*mlp\\.(gate|up|down)_proj$"
|
||||
|
||||
# FP8 W8A8 attention projections.
|
||||
FP8_GROUP = {
|
||||
"format": "float-quantized",
|
||||
"targets": [FP8_TARGET],
|
||||
"weights": {
|
||||
"num_bits": 8,
|
||||
"type": "float",
|
||||
"symmetric": True,
|
||||
"strategy": "channel",
|
||||
"dynamic": False,
|
||||
},
|
||||
"input_activations": {
|
||||
"num_bits": 8,
|
||||
"type": "float",
|
||||
"symmetric": True,
|
||||
"strategy": "token",
|
||||
"dynamic": True,
|
||||
},
|
||||
}
|
||||
|
||||
# NVFP4 W4A4 expert projections.
|
||||
NVFP4_GROUP = {
|
||||
"format": "nvfp4-pack-quantized",
|
||||
"targets": [NVFP4_TARGET],
|
||||
"weights": {
|
||||
"num_bits": 4,
|
||||
"type": "float",
|
||||
"symmetric": True,
|
||||
"strategy": "tensor_group",
|
||||
"group_size": 16,
|
||||
"dynamic": False,
|
||||
},
|
||||
"input_activations": {
|
||||
"num_bits": 4,
|
||||
"type": "float",
|
||||
"symmetric": True,
|
||||
"strategy": "tensor_group",
|
||||
"group_size": 16,
|
||||
"dynamic": "local",
|
||||
},
|
||||
}
|
||||
|
||||
# Weight-only INT4: a group whose format is not an activation format.
|
||||
WNA16_GROUP = {
|
||||
"format": "pack-quantized",
|
||||
"targets": [WNA16_TARGET],
|
||||
"weights": {
|
||||
"num_bits": 4,
|
||||
"type": "int",
|
||||
"symmetric": True,
|
||||
"strategy": "group",
|
||||
"group_size": 128,
|
||||
"dynamic": False,
|
||||
},
|
||||
"input_activations": None,
|
||||
}
|
||||
|
||||
|
||||
def _mixed_precision_config(*groups, ignore=()):
|
||||
"""Groups disagree, so compressed-tensors writes format="mixed-precision"."""
|
||||
return {
|
||||
"quant_method": "compressed-tensors",
|
||||
"format": "mixed-precision",
|
||||
"config_groups": {f"group_{i}": g for i, g in enumerate(groups)},
|
||||
"ignore": list(ignore),
|
||||
}
|
||||
|
||||
|
||||
class TestIgnoreListPrefixMatching(CustomTestCase):
|
||||
"""A parent module in `ignore` must not de-quantize its children."""
|
||||
|
||||
def test_parent_module_does_not_swallow_children(self):
|
||||
ignore = [f"{EXPERTS_LAYER}.0", "model.language_model.layers.0.linear_attn"]
|
||||
|
||||
# The parent itself is still ignored: semantics are unchanged.
|
||||
self.assertTrue(should_ignore_layer(f"{EXPERTS_LAYER}.0", ignore=ignore))
|
||||
self.assertTrue(
|
||||
should_ignore_layer(
|
||||
"model.language_model.layers.0.linear_attn", ignore=ignore
|
||||
)
|
||||
)
|
||||
|
||||
# Its quantized children are not.
|
||||
self.assertFalse(should_ignore_layer(GATE_PROJ, ignore=ignore))
|
||||
self.assertFalse(
|
||||
should_ignore_layer(
|
||||
"model.language_model.layers.0.linear_attn.out_proj", ignore=ignore
|
||||
)
|
||||
)
|
||||
|
||||
def test_module_suffix_target_still_matches(self):
|
||||
# Targets written as a module suffix keep working (dotted-path boundary).
|
||||
self.assertTrue(
|
||||
check_equal_or_regex_match(
|
||||
"model.layers.0.self_attn.kv_b_proj", ["self_attn.kv_b_proj"]
|
||||
)
|
||||
)
|
||||
self.assertTrue(check_equal_or_regex_match("model.lm_head", ["lm_head"]))
|
||||
# ... but only on a boundary, never mid-token.
|
||||
self.assertFalse(
|
||||
check_equal_or_regex_match("model.layers.0.gate_proj", ["ate"])
|
||||
)
|
||||
self.assertFalse(
|
||||
check_equal_or_regex_match(
|
||||
"model.layers.0.mlp.shared_expert_gate", ["gate"]
|
||||
)
|
||||
)
|
||||
|
||||
def test_exact_and_regex_targets_unchanged(self):
|
||||
self.assertTrue(check_equal_or_regex_match(GATE_PROJ, [GATE_PROJ]))
|
||||
self.assertTrue(check_equal_or_regex_match(GATE_PROJ, ["re:.*gate_proj$"]))
|
||||
self.assertFalse(check_equal_or_regex_match(GATE_PROJ, ["re:.*down_proj$"]))
|
||||
|
||||
|
||||
class TestMixedPrecisionFormat(CustomTestCase):
|
||||
"""The per-group `format` must win over a top-level "mixed-precision"."""
|
||||
|
||||
def test_input_activations_survive_mixed_precision(self):
|
||||
config = _mixed_precision_config(FP8_GROUP, NVFP4_GROUP)
|
||||
quant_config = CompressedTensorsConfig.from_config(config)
|
||||
|
||||
for target, scheme in quant_config.target_scheme_map.items():
|
||||
self.assertIsNotNone(
|
||||
scheme["input_activations"],
|
||||
f"input_activations dropped for target {target}",
|
||||
)
|
||||
|
||||
expert_scheme = quant_config.target_scheme_map[NVFP4_TARGET]
|
||||
self.assertEqual(expert_scheme["format"], "nvfp4-pack-quantized")
|
||||
self.assertEqual(expert_scheme["weights"].num_bits, 4)
|
||||
self.assertEqual(expert_scheme["input_activations"].num_bits, 4)
|
||||
|
||||
def test_linear_scheme_uses_per_group_format(self):
|
||||
# WNA16 is selected only when the format is "pack-quantized". Reading the
|
||||
# top-level format instead of the matched group's would see
|
||||
# "mixed-precision" and resolve no scheme at all.
|
||||
config = _mixed_precision_config(WNA16_GROUP, NVFP4_GROUP)
|
||||
quant_config = CompressedTensorsConfig.from_config(config)
|
||||
|
||||
with mock.patch.object(
|
||||
CompressedTensorsConfig, "_check_scheme_supported", return_value=True
|
||||
):
|
||||
scheme = quant_config.get_linear_scheme(
|
||||
torch.nn.Module(), layer_name=MLP_LAYER
|
||||
)
|
||||
|
||||
self.assertIsInstance(scheme, CompressedTensorsWNA16)
|
||||
self.assertEqual(scheme.pack_factor, 32 // 4)
|
||||
self.assertEqual(scheme.strategy, "group")
|
||||
self.assertEqual(scheme.group_size, 128)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user