minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jimmy Shong
2026-08-19 15:37:45 -07:00
committed by GitHub
parent c7e2c08d14
commit 5375babbac
2 changed files with 247 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
+76 -3
View File
@@ -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,6 +990,7 @@ class CompressedTensorsConfig(QuantizationConfig):
# Will be empty for models with only sparsity
if self.target_scheme_map:
if matched_target is None:
matched_target = find_matched_target(
layer_name=layer_name,
module=layer,
test/registered/unit/layers/quantization/test_compressed_tensors_lm_head.py
+165
View File
@@ -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()