diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index 645128ca0..8eb926981 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -42,6 +42,10 @@ from sglang.multimodal_gen.runtime.utils.precision import resolve_component_prec logger = init_logger(__name__) +class ComponentCheckpointUnsupportedError(ValueError): + """A component checkpoint is unsupported and must not use native fallback.""" + + def _load_auto_tokenizer_with_roberta_processing_compat(*args, **kwargs): from tokenizers import processors @@ -191,7 +195,7 @@ class ComponentLoader(ABC): component_attn_name, ) source = "sgl-diffusion" - except ComponentResidencyError: + except (ComponentCheckpointUnsupportedError, ComponentResidencyError): raise except Exception as e: if self.should_raise_customized_load_error(server_args, component_name): diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py index 462870623..f4a9fbb54 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py @@ -1,5 +1,6 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import ( TextEncoderLoader, + _resolve_and_configure_encoder_quantization, ) from sglang.multimodal_gen.runtime.models.encoders.base import finalize_encoder_folding from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -34,6 +35,12 @@ class ImageEncoderLoader(TextEncoderLoader): encoder_config = server_args.pipeline_config.image_encoder_config encoder_config.update_model_arch(model_config) + _resolve_and_configure_encoder_quantization( + encoder_config, + model_config, + component_model_path, + component_name, + ) # real dims are populated now; resolve fold vs replicate finalize_encoder_folding( encoder_config, 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 25d4cdac5..69ee4800f 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 @@ -26,6 +26,7 @@ from sglang.multimodal_gen.runtime.layers.linear import ( UnquantizedLinearMethod, ) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentCheckpointUnsupportedError, ComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( @@ -61,11 +62,12 @@ from sglang.srt.environ import envs logger = init_logger(__name__) -def _configure_text_encoder_quantization( +def _configure_encoder_quantization( model_config: EncoderConfig, model_cls: type[nn.Module], component_config: dict, component_model_path: str, + component_name: str, ) -> None: if getattr(model_cls, "manages_checkpoint_quantization", False): # Preserve model-owned formats such as Ideogram's bitsandbytes state. @@ -73,27 +75,79 @@ def _configure_text_encoder_quantization( # themselves; running the generic lifecycle as well would process twice. return - quant_config = get_quant_config( - component_config, - component_model_path, - ) + try: + quant_config = get_quant_config( + component_config, + component_model_path, + ) + except (KeyError, ValueError) as error: + raise ComponentCheckpointUnsupportedError( + f"Cannot configure checkpoint quantization for {component_name!r}: {error}" + ) from error model_config.quant_config = quant_config if quant_config is None: return - if not issubclass(model_cls, TextEncoder): - raise ValueError( - "A quantized text-encoder checkpoint requires an in-tree native " - "TextEncoder; " + if not issubclass(model_cls, EncoderTensorParallelMixin): + raise ComponentCheckpointUnsupportedError( + f"A quantized {component_name!r} checkpoint requires an in-tree " + "native encoder; " f"got {model_cls.__name__}" ) - quant_method = quant_config.get_name() - supported_methods = model_cls.supported_checkpoint_quantization_methods - if quant_method not in supported_methods: - raise ValueError( - f"{model_cls.__name__} does not support text-encoder checkpoints " - f"quantized with {quant_method!r}; supported methods: " - f"{sorted(supported_methods)}" + + capability = model_cls.checkpoint_quantization_capability + if capability is None: + raise ComponentCheckpointUnsupportedError( + f"{model_cls.__name__} does not support quantized checkpoints for " + f"{component_name!r}: no checkpoint quantization capability is declared" ) + if capability.backend != "diffusion": + raise ComponentCheckpointUnsupportedError( + f"{model_cls.__name__} declares the {capability.backend!r} checkpoint " + f"quantization backend for {component_name!r}, but the native encoder " + "loader currently supports only the 'diffusion' backend" + ) + + quant_method = quant_config.get_name() + if quant_method not in capability.methods: + raise ComponentCheckpointUnsupportedError( + f"{model_cls.__name__} does not support {component_name!r} checkpoints " + f"quantized with {quant_method!r}; supported methods for the " + f"{capability.backend!r} backend: {sorted(capability.methods)}" + ) + + +def _resolve_and_configure_encoder_quantization( + model_config: EncoderConfig, + component_config: dict, + component_model_path: str, + component_name: str, +) -> type[nn.Module]: + architectures = getattr(model_config, "architectures", []) + try: + model_cls, _ = ModelRegistry.resolve_model_cls(architectures) + except Exception as resolution_error: + try: + quant_config = get_quant_config(component_config, component_model_path) + except Exception as quantization_error: + raise ComponentCheckpointUnsupportedError( + f"Cannot parse checkpoint quantization for {component_name!r}: " + f"{quantization_error}" + ) from quantization_error + if quant_config is None: + raise + raise ComponentCheckpointUnsupportedError( + f"A quantized {component_name!r} checkpoint requires an in-tree " + f"native encoder; unsupported architectures: {architectures}" + ) from resolution_error + + _configure_encoder_quantization( + model_config, + model_cls, + component_config, + component_model_path, + component_name, + ) + return model_cls def _module_tensor_device(module: nn.Module) -> torch.device | None: @@ -118,9 +172,10 @@ def _module_tensor_device(module: nn.Module) -> torch.device | None: return next(iter(devices), None) -def _process_quantized_text_encoder_weights( +def _process_quantized_encoder_weights( model: nn.Module, process_device: torch.device, + component_name: str, ) -> int: processed_layers = 0 for module in model.modules(): @@ -144,8 +199,8 @@ def _process_quantized_text_encoder_weights( module.to(origin_device) if processed_layers == 0: raise ValueError( - "The text-encoder checkpoint declares quantization, but the model " - "did not construct any quantized linear layers" + f"The {component_name!r} checkpoint declares quantization, but the " + "model did not construct any quantized linear layers" ) return processed_layers @@ -412,14 +467,11 @@ class TextEncoderLoader(ComponentLoader): ) if post_diffusers_config_update is not None: post_diffusers_config_update() - model_cls, _ = ModelRegistry.resolve_model_cls( - getattr(encoder_config, "architectures", []) - ) - _configure_text_encoder_quantization( + model_cls = _resolve_and_configure_encoder_quantization( encoder_config, - model_cls, model_config, component_model_path, + component_name, ) encoder_dp_group = get_encoder_data_parallel_group() prefer_dp = ( @@ -485,13 +537,14 @@ class TextEncoderLoader(ComponentLoader): if quant_config is not None: if param_dtype not in quant_config.get_supported_act_dtypes(): raise ValueError( - f"Text-encoder quantization method {quant_config.get_name()!r} " + f"{component_name!r} quantization method " + f"{quant_config.get_name()!r} " f"does not support activation dtype {param_dtype}" ) if current_platform.is_mps(): raise ValueError( - f"Text-encoder quantization method {quant_config.get_name()!r} " - "is not supported on MPS" + f"{component_name!r} quantization method " + f"{quant_config.get_name()!r} is not supported on MPS" ) if current_platform.is_cuda(): capability = current_platform.get_device_capability() @@ -500,7 +553,8 @@ class TextEncoderLoader(ComponentLoader): and capability.to_int() < quant_config.get_min_capability() ): raise ValueError( - f"Text-encoder quantization method {quant_config.get_name()!r} " + f"{component_name!r} quantization method " + f"{quant_config.get_name()!r} " "requires CUDA compute capability " f">= {quant_config.get_min_capability() / 10:.1f}; got " f"{capability.to_int() / 10:.1f}" @@ -575,14 +629,16 @@ class TextEncoderLoader(ComponentLoader): ) if quant_config is not None: - processed_layers = _process_quantized_text_encoder_weights( + processed_layers = _process_quantized_encoder_weights( model, local_torch_device, + component_name, ) logger.info( - "Processed %d %s text-encoder linear layers", + "Processed %d %s linear layers for %s", processed_layers, quant_config.get_name(), + component_name, ) if component_starts_on_cpu: diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/base.py b/python/sglang/multimodal_gen/runtime/models/encoders/base.py index fc8b486ed..1053e24bc 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/base.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/base.py @@ -2,7 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod -from dataclasses import field +from dataclasses import dataclass, field +from typing import Literal import torch from torch import nn @@ -154,10 +155,22 @@ def finalize_encoder_folding( config.parallel_folding_mode = None +@dataclass(frozen=True) +class CheckpointQuantizationCapability: + """Quantized-checkpoint contract implemented by a native encoder.""" + + backend: Literal["diffusion", "srt"] + methods: frozenset[str] + + class EncoderTensorParallelMixin: """Keep an encoder on the TP group that was used to build its shards.""" _encoder_tp_group: GroupCoordinator | None = None + checkpoint_quantization_capability: CheckpointQuantizationCapability | None = None + # Some encoders own checkpoint quantization end to end because their weight + # states or sharding contract cannot use the generic loader lifecycle. + manages_checkpoint_quantization = False def bind_encoder_tp_group(self, tp_group: GroupCoordinator) -> None: self._encoder_tp_group = tp_group diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py index c5c9a694d..a23835a3e 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py @@ -17,7 +17,10 @@ from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import ( ) from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader -from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder +from sglang.multimodal_gen.runtime.models.encoders.base import ( + CheckpointQuantizationCapability, + TextEncoder, +) from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLModel MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 @@ -41,11 +44,15 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder): eight otherwise-idle ranks during encoding. """ - supports_dp_encode = True # The inherited text-layer list covers Qwen's language stack; reference # modes also execute the embedded visual tower. layer_names = [*TextEncoder.layer_names, "model.visual.blocks"] - supported_checkpoint_quantization_methods = frozenset({"fp8"}) + + supports_dp_encode = True + checkpoint_quantization_capability = CheckpointQuantizationCapability( + backend="diffusion", + methods=frozenset({"fp8"}), + ) @staticmethod def should_materialize_checkpoint_weight(name: str) -> bool: diff --git a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py new file mode 100644 index 000000000..30b2fd775 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py @@ -0,0 +1,73 @@ +import unittest +from types import SimpleNamespace +from unittest import mock + +from sglang.multimodal_gen.configs.models.encoders.clip import CLIPVisionConfig +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentCheckpointUnsupportedError, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.image_encoder_loader import ( + ImageEncoderLoader, +) + + +class TestImageEncoderQuantizationAdmission(unittest.TestCase): + def setUp(self): + self.loader = ImageEncoderLoader() + load_native_patcher = mock.patch.object( + self.loader, "load_native", return_value=object() + ) + self.load_native = load_native_patcher.start() + self.addCleanup(load_native_patcher.stop) + self.server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + image_encoder_config=CLIPVisionConfig(), + image_encoder_precision="bf16", + native_only_components=(), + ), + encoder_parallel="replicate", + resolve_component_attention_backend=lambda _name: (None, None), + ) + + def _component_config(self, architecture, *, quantized): + config = {"architectures": [architecture]} + if quantized: + config["quantization_config"] = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + } + return config + + def _config_patch(self, config): + return mock.patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "image_encoder_loader.get_diffusers_component_config", + return_value=config, + ) + + def _load(self): + return self.loader.load( + "/model/image_encoder", self.server_args, "image_encoder", "transformers" + ) + + def test_quantized_clip_checkpoint_is_not_silently_enabled(self): + config = self._component_config("CLIPVisionModelWithProjection", quantized=True) + with self._config_patch(config), self.assertRaisesRegex( + ComponentCheckpointUnsupportedError, + "CLIPVisionModel.*image_encoder.*no checkpoint quantization capability", + ): + self.loader.load_customized("/model/image_encoder", self.server_args) + + def test_unknown_quantized_architecture_does_not_fall_back(self): + config = self._component_config("UnknownVisionModel", quantized=True) + with self._config_patch(config), self.assertRaises( + ComponentCheckpointUnsupportedError + ): + self._load() + self.load_native.assert_not_called() + + def test_unknown_unquantized_architecture_keeps_native_fallback(self): + config = self._component_config("UnknownVisionModel", quantized=False) + with self._config_patch(config): + self._load() + self.load_native.assert_called_once() 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 72f0407d7..542d95b7e 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 @@ -8,12 +8,18 @@ from torch import nn from sglang.multimodal_gen.runtime.layers.linear import LinearBase from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentCheckpointUnsupportedError, +) from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import ( TextEncoderLoader, - _configure_text_encoder_quantization, - _process_quantized_text_encoder_weights, + _configure_encoder_quantization, + _process_quantized_encoder_weights, +) +from sglang.multimodal_gen.runtime.models.encoders.base import ( + CheckpointQuantizationCapability, + TextEncoder, ) -from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import ( MiniMaxH3Qwen3VLEncoder, ) @@ -147,22 +153,48 @@ class TestTextEncoderQuantization(unittest.TestCase): def test_serialized_fp8_checkpoint_configures_h3_encoder(self): model_config = SimpleNamespace(quant_config=None) - _configure_text_encoder_quantization( + _configure_encoder_quantization( model_config, MiniMaxH3Qwen3VLEncoder, {}, "/model/text_encoder", + "text_encoder", ) self.assertIs(model_config.quant_config, self.serialized) def test_encoder_class_must_opt_in(self): model_config = SimpleNamespace(quant_config=None) - with self.assertRaisesRegex(ValueError, "does not support"): - _configure_text_encoder_quantization( + with self.assertRaisesRegex( + ComponentCheckpointUnsupportedError, "does not support" + ): + _configure_encoder_quantization( model_config, TextEncoder, {}, "/model/text_encoder", + "text_encoder", + ) + + def test_srt_backend_is_not_admitted_without_an_adapter(self): + model_config = SimpleNamespace(quant_config=None) + capability = CheckpointQuantizationCapability( + backend="srt", + methods=frozenset({"fp8"}), + ) + with mock.patch.object( + MiniMaxH3Qwen3VLEncoder, + "checkpoint_quantization_capability", + capability, + ), self.assertRaisesRegex( + ComponentCheckpointUnsupportedError, + "'srt'.*only the 'diffusion' backend", + ): + _configure_encoder_quantization( + model_config, + MiniMaxH3Qwen3VLEncoder, + {}, + "/model/text_encoder", + "text_encoder", ) def test_model_managed_quantization_bypasses_generic_lifecycle(self): @@ -172,11 +204,12 @@ class TestTextEncoderQuantization(unittest.TestCase): "manages_checkpoint_quantization", True, ): - _configure_text_encoder_quantization( + _configure_encoder_quantization( model_config, TextEncoder, {}, "/model/text_encoder", + "text_encoder", ) self.assertIsNone(model_config.quant_config) @@ -213,9 +246,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase): quant_method = _RecordingQuantMethod() model = _QuantizedEncoder(quant_method) - processed = _process_quantized_text_encoder_weights( + processed = _process_quantized_encoder_weights( model, torch.device("cpu"), + "text_encoder", ) self.assertEqual(processed, 1) @@ -227,9 +261,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase): quant_method = _RecordingQuantMethod() model = _QuantizedEncoder(quant_method) - processed = _process_quantized_text_encoder_weights( + processed = _process_quantized_encoder_weights( model, torch.device("cuda", torch.cuda.current_device()), + "text_encoder", ) self.assertEqual(processed, 1) @@ -242,9 +277,10 @@ class TestQuantizedTextEncoderPostprocess(unittest.TestCase): model = _QuantizedEncoder(_RecordingQuantMethod(error=RuntimeError("boom"))) with self.assertRaisesRegex(RuntimeError, "boom"): - _process_quantized_text_encoder_weights( + _process_quantized_encoder_weights( model, torch.device("cuda", torch.cuda.current_device()), + "text_encoder", ) self.assertEqual(model.quantized.weight.device, torch.device("cpu"))