[Quant] Load compressed-tensors quantized lm_head instead of value-casting it

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):
@@ -822,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:
@@ -844,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
@@ -898,12 +913,69 @@ 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
+        # 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 @@ 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]
 

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()