From 45fe1cb8ea2fe422866587b92d535184fbb43145 Mon Sep 17 00:00:00 2001 From: Jiminator Date: Tue, 18 Aug 2026 00:15:40 +0000 Subject: [PATCH 1/3] [Quant] Load compressed-tensors quantized lm_head instead of value-casting it CompressedTensorsConfig.get_quant_method returned None for ParallelLMHead, so a checkpoint that quantizes the head (e.g. FP8 channel weights targeted via 're:.*lm_head') fell back to UnquantizedEmbeddingMethod: the fp8 weight was value-cast into a bf16 param, weight_scale was dropped, and the model emitted gibberish. Resolve a linear scheme for the head when a config target names it by layer name; module-type targets like 'Linear' are not consulted so conventional checkpoints keep their unquantized head. --- .../compressed_tensors/compressed_tensors.py | 38 ++++++ .../test_compressed_tensors_lm_head.py | 121 ++++++++++++++++++ 2 files changed, 159 insertions(+) create mode 100644 test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.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 3490fc6a184e..063b0d8f53ff 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -58,6 +58,7 @@ NPUCompressedTensorsW8A8Int8DynamicMoE, ) from sglang.srt.layers.quantization.compressed_tensors.utils import ( + check_equal_or_regex_match, find_matched_target, is_activation_quantization_format, should_ignore_layer, @@ -175,6 +176,17 @@ def get_quant_method( return UnquantizedLinearMethod() layer.scheme = scheme return CompressedTensorsLinearMethod(self) + + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + + if isinstance(layer, ParallelLMHead): + scheme = self.get_lm_head_scheme(layer=layer, layer_name=prefix) + if scheme is None: + # Unquantized head: fall back to the embedding default. + return None + layer.scheme = scheme + return CompressedTensorsLinearMethod(self) + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE if isinstance(layer, FusedMoE): @@ -898,6 +910,32 @@ def get_linear_scheme( logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name) return scheme + def get_lm_head_scheme( + self, layer: torch.nn.Module, layer_name: Optional[str] = None + ) -> Optional[CompressedTensorsLinearScheme]: + """Resolve the scheme for a ParallelLMHead, or None if the checkpoint + stores the head unquantized. + + The head is treated as quantized only when a config target names it by + layer name (exact or ``re:`` regex, e.g. ``re:.*lm_head``). Module-type + targets like ``Linear`` are not consulted: llm-compressor emits those + for decoder linears, and checkpoints following the common convention + leave the head out of both ``targets`` and ``ignore`` — matching by + name keeps such heads on the unquantized path instead of tripping + ``find_matched_target``'s unmatched-layer error. + """ + if layer_name is None or not self.target_scheme_map: + return None + if should_ignore_layer( + layer_name, ignore=self.ignore, fused_mapping=self.packed_modules_mapping + ): + return None + if not check_equal_or_regex_match( + layer_name=layer_name, targets=self.target_scheme_map.keys() + ): + return None + return self.get_linear_scheme(layer=layer, layer_name=layer_name) + def get_scheme_dict( self, layer: torch.nn.Module, layer_name: str | None = None ) -> dict[str, QuantizationArgs | str | None] | None: diff --git a/test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.py b/test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.py new file mode 100644 index 000000000000..259f551dd949 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.py @@ -0,0 +1,121 @@ +"""Unit tests for compressed-tensors lm_head scheme resolution — CPU-only.""" + +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.mock import patch + +import torch + +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, + CompressedTensorsLinearMethod, +) +from sglang.test.test_utils import CustomTestCase + +_FP8_WEIGHTS = { + "num_bits": 8, + "type": "float", + "strategy": "channel", + "symmetric": True, + "dynamic": False, +} +_FP8_DYNAMIC_ACTS = { + "num_bits": 8, + "type": "float", + "strategy": "token", + "symmetric": True, + "dynamic": True, +} + + +def _config(targets, ignore=()): + return CompressedTensorsConfig.from_config( + { + "format": "float-quantized", + "quant_method": "compressed-tensors", + "ignore": list(ignore), + "config_groups": { + "group_0": { + "targets": list(targets), + "weights": _FP8_WEIGHTS, + "input_activations": _FP8_DYNAMIC_ACTS, + } + }, + } + ) + + +class _Head(torch.nn.Module): + pass + + +_GET_LINEAR_SCHEME = ( + "sglang.srt.layers.quantization.compressed_tensors.compressed_tensors." + "CompressedTensorsConfig.get_linear_scheme" +) + + +class TestGetLmHeadScheme(CustomTestCase): + """The head resolves a scheme only when a config target names it by + layer name; module-type targets and ignored heads stay unquantized.""" + + def test_regex_target_resolves(self): + config = _config(["re:.*lm_head", "re:.*mlp\\.down_proj$"]) + head = _Head() + with patch(_GET_LINEAR_SCHEME, return_value="scheme") as mock_resolve: + scheme = config.get_lm_head_scheme(head, "lm_head") + self.assertEqual(scheme, "scheme") + mock_resolve.assert_called_once_with(layer=head, layer_name="lm_head") + + def test_exact_target_resolves(self): + config = _config(["lm_head"]) + with patch(_GET_LINEAR_SCHEME, return_value="scheme"): + self.assertEqual(config.get_lm_head_scheme(_Head(), "lm_head"), "scheme") + + def test_ignored_head_is_none(self): + config = _config(["re:.*lm_head"], ignore=["lm_head"]) + with patch(_GET_LINEAR_SCHEME) as mock_resolve: + self.assertIsNone(config.get_lm_head_scheme(_Head(), "lm_head")) + mock_resolve.assert_not_called() + + def test_module_type_target_is_none(self): + # llm-compressor emits "Linear" for decoder linears; an unmentioned + # head must stay on the unquantized path instead of tripping + # find_matched_target's unmatched-layer error. + config = _config(["Linear"]) + with patch(_GET_LINEAR_SCHEME) as mock_resolve: + self.assertIsNone(config.get_lm_head_scheme(_Head(), "lm_head")) + mock_resolve.assert_not_called() + + def test_no_layer_name_is_none(self): + config = _config(["re:.*lm_head"]) + self.assertIsNone(config.get_lm_head_scheme(_Head(), None)) + + +class TestGetQuantMethodLmHead(CustomTestCase): + def _head(self): + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + + # __new__ is enough: get_quant_method only isinstance-checks the + # layer and attaches `scheme` to it. + return ParallelLMHead.__new__(ParallelLMHead) + + def test_quantized_head_gets_linear_method(self): + config = _config(["re:.*lm_head"]) + head = self._head() + with patch.object(config, "get_lm_head_scheme", return_value="scheme"): + method = config.get_quant_method(head, "lm_head") + self.assertIsInstance(method, CompressedTensorsLinearMethod) + self.assertEqual(head.scheme, "scheme") + + def test_unquantized_head_falls_back(self): + config = _config(["Linear"]) + head = self._head() + self.assertIsNone(config.get_quant_method(head, "lm_head")) + + +if __name__ == "__main__": + unittest.main() From 4c238895c9bc2ca97670379e6b00e53e91e7ce1f Mon Sep 17 00:00:00 2001 From: Jiminator Date: Tue, 18 Aug 2026 06:05:14 +0000 Subject: [PATCH 2/3] Address review: share the head's target match downstream, reject block scales A dotted-suffix target match ("lm_head" for a "language_model.lm_head" prefix) passed the head guard but find_matched_target's exact/regex name pass could not re-derive it, raising ValueError at load; carry the matched target into get_scheme_dict instead of re-deriving. Block-FP8 heads are rejected loudly: the vocab-parallel weight loader shards output_dim=0 params by vocab index and cannot load a vocab/block_n weight_scale. --- .../compressed_tensors/compressed_tensors.py | 58 ++++++++++++++----- .../test_compressed_tensors_lm_head.py | 46 ++++++++++++++- 2 files changed, 90 insertions(+), 14 deletions(-) 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 063b0d8f53ff..ebd7373be994 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -834,7 +834,10 @@ def get_moe_scheme( ) def get_linear_scheme( - self, layer: torch.nn.Module, layer_name: Optional[str] = None + self, + layer: torch.nn.Module, + layer_name: Optional[str] = None, + matched_target: Optional[str] = None, ) -> Optional[CompressedTensorsLinearScheme]: """ compressed-tensors supports non uniform in the following way: @@ -856,7 +859,7 @@ def get_linear_scheme( # need to make accelerate optional in ct to do this # Use the new get_scheme_dict method to extract QuantizationArgs - scheme_dict = self.get_scheme_dict(layer, layer_name) + scheme_dict = self.get_scheme_dict(layer, layer_name, matched_target) weight_quant = None input_quant = None scheme_format = None @@ -930,18 +933,46 @@ def get_lm_head_scheme( layer_name, ignore=self.ignore, fused_mapping=self.packed_modules_mapping ): return None - if not check_equal_or_regex_match( - layer_name=layer_name, targets=self.target_scheme_map.keys() - ): + # check_equal_or_regex_match also accepts dotted-suffix targets + # (e.g. target "lm_head" for a "language_model.lm_head" prefix), which + # find_matched_target's exact/regex name pass would miss — so the match + # made here is carried through instead of being re-derived downstream. + matched_target = next( + ( + target + for target in self.target_scheme_map + if check_equal_or_regex_match(layer_name=layer_name, targets=[target]) + ), + None, + ) + if matched_target is None: return None - return self.get_linear_scheme(layer=layer, layer_name=layer_name) + weights = self.target_scheme_map[matched_target].get("weights") + if weights is not None and weights.block_structure: + # The vocab-parallel weight loader shards output_dim=0 params by + # vocab index; a block weight_scale's first dim is vocab/block_n, + # which that loader cannot shard or even load at TP=1. + raise NotImplementedError( + "Block-quantized lm_head is not supported; use channel or " + "tensor weight scales for the head." + ) + return self.get_linear_scheme( + layer=layer, layer_name=layer_name, matched_target=matched_target + ) def get_scheme_dict( - self, layer: torch.nn.Module, layer_name: str | None = None + self, + layer: torch.nn.Module, + layer_name: str | None = None, + matched_target: str | None = None, ) -> dict[str, QuantizationArgs | str | None] | None: """ Extract the QuantizationArgs for a given layer. + A caller that already resolved the layer's target (e.g. via + suffix-aware matching) passes it as ``matched_target`` to skip + ``find_matched_target``'s stricter exact/regex lookup. + Returns: dict with { "weights": QuantizationArgs, @@ -956,12 +987,13 @@ def get_scheme_dict( # Will be empty for models with only sparsity if self.target_scheme_map: - matched_target = find_matched_target( - layer_name=layer_name, - module=layer, - targets=self.target_scheme_map.keys(), - fused_mapping=self.packed_modules_mapping, - ) + if matched_target is None: + matched_target = find_matched_target( + layer_name=layer_name, + module=layer, + targets=self.target_scheme_map.keys(), + fused_mapping=self.packed_modules_mapping, + ) return self.target_scheme_map[matched_target] diff --git a/test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.py b/test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.py index 259f551dd949..c5564a4bbbde 100644 --- a/test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.py +++ b/test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.py @@ -68,7 +68,9 @@ def test_regex_target_resolves(self): with patch(_GET_LINEAR_SCHEME, return_value="scheme") as mock_resolve: scheme = config.get_lm_head_scheme(head, "lm_head") self.assertEqual(scheme, "scheme") - mock_resolve.assert_called_once_with(layer=head, layer_name="lm_head") + mock_resolve.assert_called_once_with( + layer=head, layer_name="lm_head", matched_target="re:.*lm_head" + ) def test_exact_target_resolves(self): config = _config(["lm_head"]) @@ -94,6 +96,48 @@ def test_no_layer_name_is_none(self): config = _config(["re:.*lm_head"]) self.assertIsNone(config.get_lm_head_scheme(_Head(), None)) + def test_prefixed_head_with_plain_target_resolves(self): + """Bug regression: `check_equal_or_regex_match` accepts the dotted + suffix ("lm_head" target for a "language_model.lm_head" prefix) but + `find_matched_target`'s name pass is exact/regex only, so re-deriving + the match downstream raised ValueError at load instead of resolving + the scheme. The suffix match must be carried through.""" + from sglang.srt.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsW8A8Fp8, + ) + + config = _config(["lm_head"]) + with patch( + "sglang.srt.layers.quantization.compressed_tensors." + "compressed_tensors.CompressedTensorsConfig._check_scheme_supported", + return_value=True, + ): + scheme = config.get_lm_head_scheme(_Head(), "language_model.lm_head") + self.assertIsInstance(scheme, CompressedTensorsW8A8Fp8) + + def test_block_quantized_head_is_rejected(self): + """Bug regression: a block-FP8 head resolves to a weight_scale whose + first dim is vocab/block_n, which the vocab-parallel weight loader + (asserting dim0 == vocab size on output_dim=0 params) cannot load + even at TP=1. Reject loudly instead of asserting mid-load.""" + block_weights = dict(_FP8_WEIGHTS, strategy="block", block_structure=[128, 128]) + config = CompressedTensorsConfig.from_config( + { + "format": "float-quantized", + "quant_method": "compressed-tensors", + "ignore": [], + "config_groups": { + "group_0": { + "targets": ["re:.*lm_head"], + "weights": block_weights, + "input_activations": _FP8_DYNAMIC_ACTS, + } + }, + } + ) + with self.assertRaises(NotImplementedError): + config.get_lm_head_scheme(_Head(), "lm_head") + class TestGetQuantMethodLmHead(CustomTestCase): def _head(self): From 8a7fa1387b89d9597e8fab1680b100e38ba8c828 Mon Sep 17 00:00:00 2001 From: Jiminator Date: Tue, 18 Aug 2026 06:45:56 +0000 Subject: [PATCH 3/3] Note first-match rule for multi-group head targets --- .../quantization/compressed_tensors/compressed_tensors.py | 3 +++ 1 file changed, 3 insertions(+) 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 ebd7373be994..aeecaee6208c 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -937,6 +937,9 @@ def get_lm_head_scheme( # (e.g. target "lm_head" for a "language_model.lm_head" prefix), which # find_matched_target's exact/regex name pass would miss — so the match # made here is carried through instead of being re-derived downstream. + # When several config groups name the head, the first target in config + # order wins — the same first-match rule find_matched_target applies + # to every other layer. matched_target = next( ( target