From 5375babbac9977cdb8f061cec77b6efd0987a1fd Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:37:45 -0700 Subject: [PATCH] [Quant] Load compressed-tensors quantized lm_head instead of value-casting it (#35228) --- .../compressed_tensors/compressed_tensors.py | 91 +++++++++- .../test_compressed_tensors_lm_head.py | 165 ++++++++++++++++++ 2 files changed, 247 insertions(+), 9 deletions(-) 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 3490fc6a1..aeecaee62 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 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import ( 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 @@ class CompressedTensorsConfig(QuantizationConfig): 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): @@ -822,7 +834,10 @@ class CompressedTensorsConfig(QuantizationConfig): ) 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: @@ -844,7 +859,7 @@ class CompressedTensorsConfig(QuantizationConfig): # 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 @@ -898,12 +913,69 @@ class CompressedTensorsConfig(QuantizationConfig): 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 + # 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. + # 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 + 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 + 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, @@ -918,12 +990,13 @@ class CompressedTensorsConfig(QuantizationConfig): # 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 new file mode 100644 index 000000000..c5564a4bb --- /dev/null +++ b/test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.py @@ -0,0 +1,165 @@ +"""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", matched_target="re:.*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)) + + 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 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()