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

[diffusion] feat: support component-scoped quantization overrides (#36035)

Browse Source
This commit is contained in:
Mick
2026-08-25 09:20:47 +08:00
committed by GitHub
parent ddea7b9156
commit f8f9226cd2
8 changed files with 124 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+7 -1
View File
@@ -248,6 +248,7 @@ def _configure_encoder_quantization(
component_weights_path: str,
component_name: str,
explicit_quantization: str | None = None,
ignored_layers: list[str] | None = None,
) -> None:
if getattr(model_cls, "manages_checkpoint_quantization", False):
if explicit_quantization is not None:
@@ -292,7 +293,9 @@ def _configure_encoder_quantization(
get_quantization_config,
)
model_config.quant_config = get_quantization_config(explicit_quantization)()
model_config.quant_config = get_quantization_config(explicit_quantization)(
ignored_layers=ignored_layers
)
quant_config = model_config.quant_config
if quant_config is None:
return
@@ -311,6 +314,7 @@ def _resolve_and_configure_encoder_quantization(
component_weights_path: str,
component_name: str,
explicit_quantization: str | None = None,
ignored_layers: list[str] | None = None,
) -> type[nn.Module]:
architectures = getattr(model_config, "architectures", [])
try:
@@ -351,6 +355,7 @@ def _resolve_and_configure_encoder_quantization(
component_weights_path,
component_name,
explicit_quantization,
ignored_layers,
)
return model_cls
@@ -725,6 +730,7 @@ class TextEncoderLoader(ComponentLoader):
component_weights_path,
component_name,
server_args.component_quantizations.get(component_name),
server_args.component_quantization_ignored_layers.get(component_name),
)
if issubclass(model_cls, EncoderTensorParallelMixin):
model_cls.configure_component_paths(
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+10 -1
View File
@@ -101,7 +101,14 @@ def _server_args_for_transformer_component(
"""Mask global quantized override flags for secondary transformer components."""
component_weights_path = server_args.component_weights_paths.get(component_name)
component_quantization = server_args.component_quantizations.get(component_name)
if component_weights_path is not None or component_quantization is not None:
component_ignored_layers = server_args.component_quantization_ignored_layers.get(
component_name
)
if (
component_weights_path is not None
or component_quantization is not None
or component_ignored_layers is not None
):
component_server_args = copy.copy(server_args)
if component_weights_path is not None:
component_server_args.transformer_weights_path = component_weights_path
@@ -118,6 +125,8 @@ def _server_args_for_transformer_component(
component_quantization,
component_name,
)
if component_ignored_layers is not None:
component_server_args.quantization_ignored_layers = component_ignored_layers
return component_server_args
if component_name not in ("transformer_2", "unconditional_transformer"):
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+79
View File
@@ -309,6 +309,10 @@ class ServerArgs(DisaggServerArgsMixin):
# Explicit quantization override for one component. Self-describing
# checkpoints remain auto-detected and do not need this override.
component_quantizations: dict[str, str] = field(default_factory=dict)
# Component-local layer name patterns to skip during online quantization.
component_quantization_ignored_layers: dict[str, list[str]] = field(
default_factory=dict
)
# Optional LTX-2.5 decoder is large enough to load only when requested.
load_diffusion_decoder: bool = False
@@ -1741,6 +1745,31 @@ class ServerArgs(DisaggServerArgsMixin):
)
normalized_quantizations[component] = quantization
self.component_quantizations = normalized_quantizations
normalized_ignored_layers: dict[str, list[str]] = {}
for component, layers in self.component_quantization_ignored_layers.items():
component = str(component).strip().replace("-", "_")
if isinstance(layers, str):
layers = [layers]
if (
not component
or not isinstance(layers, (list, tuple))
or not all(isinstance(layer, str) and layer.strip() for layer in layers)
):
raise ValueError(
"Component quantization ignored layers require a component "
"and non-empty layer patterns"
)
normalized_ignored_layers[component] = [layer.strip() for layer in layers]
missing_quantization = set(normalized_ignored_layers) - set(
self.component_quantizations
)
if missing_quantization:
raise ValueError(
"Component quantization ignored layers require a matching "
"quantization override for: "
f"{', '.join(sorted(missing_quantization))}"
)
self.component_quantization_ignored_layers = normalized_ignored_layers
# Convert string disagg_role to enum (from CLI/config)
if isinstance(self.disagg_role, str):
@@ -2867,6 +2896,46 @@ class ServerArgs(DisaggServerArgsMixin):
alias_suffix="-quantization",
)
@staticmethod
def _extract_component_quantization_ignored_layers(
unknown_args: list[str],
) -> tuple[dict[str, list[str]], list[str]]:
ignored_layers: dict[str, list[str]] = {}
remaining: list[str] = []
i = 0
prefixes = (
"--component-quantization-ignored-layers.",
"--component_quantization_ignored_layers.",
)
while i < len(unknown_args):
arg = unknown_args[i]
key_part = arg.split("=", 1)[0] if "=" in arg else arg
prefix = next(
(candidate for candidate in prefixes if key_part.startswith(candidate)),
None,
)
if prefix is None:
remaining.append(arg)
i += 1
continue
component = key_part[len(prefix) :].replace("-", "_")
if "=" in arg:
values = [arg.split("=", 1)[1]]
else:
values = []
while i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith(
"-"
):
i += 1
values.append(unknown_args[i])
if component and values:
ignored_layers[component] = values
else:
remaining.append(arg)
i += 1
return ignored_layers, remaining
@staticmethod
def _extract_component_attention_backends(
unknown_args: list[str],
@@ -2917,6 +2986,9 @@ class ServerArgs(DisaggServerArgsMixin):
dynamic_quantizations, remaining = cls._extract_component_quantizations(
unknown_args
)
dynamic_ignored_layers, remaining = (
cls._extract_component_quantization_ignored_layers(remaining)
)
# Extract the more specific weights suffix before the generic path alias.
dynamic_weights_paths, remaining = cls._extract_component_weights_paths(
remaining
@@ -2957,6 +3029,13 @@ class ServerArgs(DisaggServerArgsMixin):
existing.update(dynamic_quantizations)
provided_args["component_quantizations"] = existing
explicit_arg_names.add("component_quantizations")
if dynamic_ignored_layers:
existing = dict(
provided_args.get("component_quantization_ignored_layers") or {}
)
existing.update(dynamic_ignored_layers)
provided_args["component_quantization_ignored_layers"] = existing
explicit_arg_names.add("component_quantization_ignored_layers")
if dynamic_attention_backends:
existing = cls._parse_component_attention_backend_map(
provided_args.get("component_attention_backends")
python/sglang/multimodal_gen/test/unit/test_ideogram4.py
+5
View File
@@ -511,6 +511,7 @@ class TestIdeogram4(unittest.TestCase):
nunchaku_config={"enabled": True},
component_weights_paths={},
component_quantizations={},
component_quantization_ignored_layers={},
)
component_args = _server_args_for_transformer_component(
server_args, "unconditional_transformer"
@@ -532,6 +533,9 @@ class TestIdeogram4(unittest.TestCase):
)
},
component_quantizations={"unconditional_transformer": "fp8"},
component_quantization_ignored_layers={
"unconditional_transformer": ["lm_head"]
},
)
component_args = _server_args_for_transformer_component(
@@ -546,6 +550,7 @@ class TestIdeogram4(unittest.TestCase):
)
self.assertIsNone(component_args.nunchaku_config)
self.assertEqual(component_args.quantization, "fp8")
self.assertEqual(component_args.quantization_ignored_layers, ["lm_head"])
def test_ideogram_nvfp4_unconditional_transformer_path_uses_sibling_file(self):
self.assertEqual(
python/sglang/multimodal_gen/test/unit/test_server_args.py
+7
View File
@@ -496,6 +496,9 @@ class TestServerArgsPathExpansion(unittest.TestCase):
"--image-encoder-weights-path=/custom/image_encoder.safetensors",
"--component-quantizations.text_encoder",
"kitchen_int8",
"--component-quantization-ignored-layers.text_encoder",
"model.layers.0",
"lm_head",
"--transformer-quantization=fp8",
"--component-attention-backends.transformer",
"fa3",
@@ -543,6 +546,10 @@ class TestServerArgsPathExpansion(unittest.TestCase):
{"text_encoder": "kitchen_int8", "transformer": "fp8"},
server_args.component_quantizations,
)
self.assertEqual(
{"text_encoder": ["model.layers.0", "lm_head"]},
server_args.component_quantization_ignored_layers,
)
def test_serve_cli_defaults_warmup_on(self):
from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py
+2
View File
@@ -363,10 +363,12 @@ class TestTextEncoderQuantization(unittest.TestCase):
"/model/text_encoder",
"text_encoder",
explicit_quantization="kitchen_int8",
ignored_layers=["lm_head"],
)
self.assertIsInstance(model_config.quant_config, KitchenInt8Config)
self.assertFalse(model_config.quant_config.is_checkpoint_int8_serialized)
self.assertEqual(model_config.quant_config.ignored_layers, ["lm_head"])
def test_weight_file_metadata_configures_native_encoder(self):
model_config = SimpleNamespace(quant_config=None)