[Quant] Load compressed-tensors quantized lm_head instead of value-casting it (#35228)

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]