From f69af7b7ad62f055867732196c13f0dd742097d2 Mon Sep 17 00:00:00 2001 From: Hert4 <98259769+Hert4@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:55:33 +0700 Subject: [PATCH] [Bugfix] compressed-tensors: mixed-precision checkpoints silently load unquantized (#32736) Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> --- .../compressed_tensors/compressed_tensors.py | 29 ++- .../quantization/compressed_tensors/utils.py | 15 +- ...test_compressed_tensors_mixed_precision.py | 192 ++++++++++++++++++ 3 files changed, 231 insertions(+), 5 deletions(-) create mode 100644 test/registered/unit/layers/quantization/test_compressed_tensors_mixed_precision.py diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 76e5c139e..984d7cbc5 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/utils.py b/python/sglang/srt/layers/quantization/compressed_tensors/utils.py index d6c2ca3d2..5369d8d9d 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/utils.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/utils.py @@ -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 diff --git a/test/registered/unit/layers/quantization/test_compressed_tensors_mixed_precision.py b/test/registered/unit/layers/quantization/test_compressed_tensors_mixed_precision.py new file mode 100644 index 000000000..90ca492c3 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_compressed_tensors_mixed_precision.py @@ -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()