diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index b90b50b18..878fd6212 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -131,6 +131,8 @@ pipeline's registered module name: | --- | --- | --- | --- | | Replace a component | `--component-paths. {MODEL}` | `---path {MODEL}` | Load the replacement component's configuration and weights | | Replace only its weights | `--component-weights-paths. {WEIGHTS}` | `---weights-path {WEIGHTS}` | Retain the base component configuration and replace its weights | +| Quantize an unquantized component online | `--component-quantizations. {METHOD}` | `---quantization {METHOD}` | Apply a method supported by that component's native loader | +| Keep selected component layers unquantized | `--component-quantization-ignored-layers. {PATTERN...}` | None | Pass component-local ignored-layer patterns to its online quantizer | For example, pair a replacement text-encoder configuration with a separate single-file checkpoint as follows: @@ -152,6 +154,7 @@ unquantized during online quantization. For a native text encoder: - `--component-paths.text_encoder {MODEL}` replaces the text-encoder checkpoint; `--text-encoder-path {MODEL}` is its shorter alias +- `--component-quantizations.text_encoder {METHOD}` applies supported online quantization to an unquantized native encoder; pair it with `--component-quantization-ignored-layers.text_encoder {PATTERN...}` when selected layers must remain unquantized - Quantization metadata is auto-detected from that checkpoint. The native loader accepts compatible serialized formats without a model-name allowlist and rejects implementations that do not construct the required quantized layers. The same contract applies to every weighted component: path routing is generic, diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index 44b56b49b..88248d748 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -23,6 +23,8 @@ Use these paths: (local or Hub) - `--quantization`: override the quantization method used by the transformer loader - `--quantization-ignored-layers`: transformer layer name patterns to keep unquantized during online quantization (e.g. `attention.to_`) +- `--component-quantizations.`: explicitly apply supported online quantization to an unquantized component +- `--component-quantization-ignored-layers.`: component-local layer patterns to keep unquantized during that online quantization - `--component-paths.text_encoder`: replace a native text encoder with a checkpoint whose `quantization_config` is auto-detected - `--text-encoder-path`: shorter alias for `--component-paths.text_encoder` - `--kv-cache-quant`: compress completed causal KV-cache chunks for supported realtime models @@ -78,13 +80,15 @@ paths: | Library-managed Transformers or Diffusers components | Delegates to the upstream `from_pretrained` path and inherits its format support and validation behavior. The local PE model uses this path; compatible formats remain model-specific. | | Native auxiliary components that load raw state dicts | Quantized checkpoints are rejected before model construction until that component has a quantized materialization implementation. This includes connectors, duration heads, bridges, diffusion decoders, sound tokenizers, spatial upsamplers, and vocoders. | -`--quantization` is the explicit override for the transformer loader; it is not -the boundary of component quantization support. Other pre-quantized component -repositories select their format through their own metadata and the capability -of the selected loader. A generic string override without a matching -materialization backend would advertise support that the component does not -have, while a quantized weight file without matching config metadata cannot be -identified or restored generically. +`--quantization` is the explicit override for the primary transformer loader; +`--component-quantizations.` expresses the same intent for one +supported component. Pair the latter with +`--component-quantization-ignored-layers.` to keep matching layers +unquantized. Pre-quantized component repositories instead select their format +through their own metadata and the capability of the selected loader. A generic +string override without a matching materialization backend would advertise +support that the component does not have, while a quantized weight file without +matching config metadata cannot be identified or restored generically. ## Quant Families diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 7c3140acd..f86f80204 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -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( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 34d0eda5d..4b4050c7d 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -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"): diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 0fff4346b..df84ff700 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -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") diff --git a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py b/python/sglang/multimodal_gen/test/unit/test_ideogram4.py index 031af5f86..64155e085 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ideogram4.py +++ b/python/sglang/multimodal_gen/test/unit/test_ideogram4.py @@ -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( diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index de211cfbd..03b5c6a86 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -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 ( diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py index 8b467c29d..0d9453fc9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py @@ -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)