From 9856b58de43dc5730b69c4aae24e48789d871f02 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 24 Aug 2026 19:59:35 +0800 Subject: [PATCH] [diffusion] feat: support loading serialized fp8 clip image encoders (#36056) --- .../configs/models/encoders/base.py | 11 ++- .../component_loaders/text_encoder_loader.py | 79 ++++++++++--------- .../runtime/models/encoders/base.py | 2 + .../runtime/models/encoders/clip.py | 5 +- .../test/unit/test_image_encoder_loader.py | 22 ++++++ .../test/unit/test_text_encoder_loader.py | 21 +++++ 6 files changed, 98 insertions(+), 42 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/models/encoders/base.py b/python/sglang/multimodal_gen/configs/models/encoders/base.py index a43b713f8..5ad7c04f4 100644 --- a/python/sglang/multimodal_gen/configs/models/encoders/base.py +++ b/python/sglang/multimodal_gen/configs/models/encoders/base.py @@ -1,8 +1,10 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + from dataclasses import dataclass, field -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal import torch @@ -10,6 +12,11 @@ from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +if TYPE_CHECKING: + from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig as SRTQuantizationConfig, + ) + @dataclass class EncoderArchConfig(ArchConfig): @@ -73,7 +80,7 @@ class EncoderConfig(ModelConfig): arch_config: ArchConfig = field(default_factory=EncoderArchConfig) prefix: str = "" - quant_config: QuantizationConfig | None = None + quant_config: QuantizationConfig | SRTQuantizationConfig | None = None lora_config: Any | None = None # Parallel folding: during the encoding stage the whole DiT replica is idle, 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 952bfcc21..3aaefe75a 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 @@ -3,7 +3,6 @@ import glob import os import re from collections.abc import Callable, Generator, Iterable -from itertools import chain from typing import cast import torch @@ -96,9 +95,14 @@ from sglang.multimodal_gen.runtime.weights.source import ( from sglang.multimodal_gen.utils import PRECISION_TO_TYPE from sglang.srt.environ import envs from sglang.srt.layers.linear import LinearBase as SrtLinearBase +from sglang.srt.layers.quantization.fp8 import Fp8Config as SrtFp8Config from sglang.srt.layers.quantization.unquant import ( UnquantizedLinearMethod as SrtUnquantizedLinearMethod, ) +from sglang.srt.model_loader.checkpoint_quantization import ( + resolve_checkpoint_quant_spec, +) +from sglang.srt.model_loader.post_load import stage_module_for_post_load logger = init_logger(__name__) @@ -129,12 +133,42 @@ def _delegate_standard_bnb4_to_transformers( ) +def _get_srt_encoder_quant_config( + component_config: dict, + model_cls: type[EncoderTensorParallelMixin], +) -> SrtFp8Config | None: + quant_spec = resolve_checkpoint_quant_spec(component_config) + if quant_spec is None: + return None + if quant_spec.declared_method != "fp8": + raise ComponentCheckpointUnsupportedError( + "The SRT encoder checkpoint adapter supports only serialized 'fp8', " + f"got {quant_spec.declared_method!r}" + ) + + config = dict(quant_spec.config) + config["packed_modules_mapping"] = model_cls.packed_modules_mapping + return SrtFp8Config.from_config(config) + + def _get_encoder_quant_config( component_config: dict, component_model_path: str, component_weights_path: str, model_cls: type[nn.Module] | None = None, ): + if ( + model_cls is not None + and issubclass(model_cls, EncoderTensorParallelMixin) + and model_cls.checkpoint_quantization_backend == "srt" + ): + srt_quant_config = _get_srt_encoder_quant_config( + component_config, + model_cls, + ) + if srt_quant_config is not None: + return srt_quant_config + quant_config = get_quant_config(component_config, component_model_path) name_mapper = None if model_cls is not None: @@ -290,28 +324,6 @@ def _resolve_and_configure_encoder_quantization( return model_cls -def _module_tensor_device(module: nn.Module) -> torch.device | None: - """Return the device of a module's own tensors. - - Quantized linear layers are expected to keep their parameters and buffers - together. Failing explicitly is safer than staging only part of a layer. - """ - - devices = { - tensor.device - for tensor in chain( - module.parameters(recurse=False), - module.buffers(recurse=False), - ) - } - if len(devices) > 1: - raise ValueError( - f"Cannot stage {type(module).__name__} with tensors on multiple " - f"devices: {sorted(map(str, devices))}" - ) - return next(iter(devices), None) - - def _process_quantized_encoder_weights( model: nn.Module, process_device: torch.device | None, @@ -327,23 +339,12 @@ def _process_quantized_encoder_weights( (UnquantizedLinearMethod, SrtUnquantizedLinearMethod), ): continue - - origin_device = _module_tensor_device(module) - should_stage = ( - process_device is not None - and origin_device is not None - and origin_device != process_device - ) - if should_stage: - module.to(process_device) - try: + if process_device is None: quant_method.process_weights_after_loading(module) - processed_layers += 1 - finally: - # Post-load methods may replace parameters or register buffers. Move - # the complete layer back so component residency remains authoritative. - if should_stage: - module.to(origin_device) + else: + with stage_module_for_post_load(module, process_device): + quant_method.process_weights_after_loading(module) + processed_layers += 1 if processed_layers == 0: raise ValueError( f"The {component_name!r} checkpoint declares quantization, but the " diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/base.py b/python/sglang/multimodal_gen/runtime/models/encoders/base.py index 1f5147d02..bb8454ffa 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/base.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/base.py @@ -158,6 +158,8 @@ class EncoderTensorParallelMixin: """Keep an encoder on the TP group that was used to build its shards.""" _encoder_tp_group: GroupCoordinator | None = None + checkpoint_quantization_backend = "diffusion" + packed_modules_mapping: dict[str, list[str]] = {} # 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 diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py index 037cd3513..75df48a7f 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py @@ -17,12 +17,12 @@ from sglang.multimodal_gen.configs.models.encoders import ( CLIPTextConfig, CLIPVisionConfig, ) -from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader from sglang.multimodal_gen.runtime.models.encoders.base import ImageEncoder, TextEncoder from sglang.multimodal_gen.runtime.models.encoders.vision import ( resolve_visual_encoder_outputs, ) +from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.models.clip import ( CLIPEncoder, CLIPTextEmbeddings, @@ -337,6 +337,9 @@ class CLIPVisionModel(ImageEncoder): config_class = CLIPVisionConfig main_input_name = "pixel_values" packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + # CLIPEncoder and its linears come from SRT, so their serialized checkpoint + # format must use the matching SRT quantization implementation. + checkpoint_quantization_backend = "srt" def __init__(self, config: CLIPVisionConfig) -> None: super().__init__(config) 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 index 520bc75df..17c8b4952 100644 --- a/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py @@ -11,6 +11,11 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp from sglang.multimodal_gen.runtime.loader.component_loaders.image_encoder_loader import ( ImageEncoderLoader, ) +from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import ( + _configure_encoder_quantization, +) +from sglang.multimodal_gen.runtime.models.encoders.clip import CLIPVisionModel +from sglang.srt.layers.quantization.fp8 import Fp8Config as SRTFp8Config class TestImageEncoderQuantizationAdmission(unittest.TestCase): @@ -54,6 +59,23 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase): "/model/image_encoder", self.server_args, "image_encoder", "transformers" ) + def test_clip_serialized_fp8_uses_srt_quantization(self): + encoder_config = CLIPVisionConfig() + _configure_encoder_quantization( + encoder_config, + CLIPVisionModel, + self._component_config("CLIPVisionModelWithProjection", quantized=True), + "/model/image_encoder", + "/model/image_encoder", + "image_encoder", + ) + + self.assertIsInstance(encoder_config.quant_config, SRTFp8Config) + self.assertEqual( + encoder_config.quant_config.packed_modules_mapping, + {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}, + ) + def test_unknown_quantized_architecture_does_not_fall_back(self): config = self._component_config("UnknownVisionModel", quantized=True) with self._config_patch(config), self.assertRaises( 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 18ce38b3c..d9a470094 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 @@ -37,6 +37,7 @@ from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import ( MiniMaxH3Qwen3VLEncoder, ) from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import Qwen3VLTextModel +from sglang.srt.layers.linear import LinearBase as SrtLinearBase class TestTextEncoderClassResolution(unittest.TestCase): @@ -640,7 +641,27 @@ class _QuantizedEncoder(nn.Module): self.unquantized = nn.Linear(2, 2, bias=False) +class _SRTQuantizedLinear(SrtLinearBase): + def __init__(self, quant_method): + nn.Module.__init__(self) + self.weight = nn.Parameter(torch.empty(2, 2), requires_grad=False) + self.quant_method = quant_method + + class TestQuantizedTextEncoderPostprocess(unittest.TestCase): + def test_processes_srt_quantized_linear(self): + quant_method = _RecordingQuantMethod() + model = _SRTQuantizedLinear(quant_method) + + processed = _process_quantized_encoder_weights( + model, + torch.device("cpu"), + "image_encoder", + ) + + self.assertEqual(processed, 1) + self.assertEqual(quant_method.devices, [torch.device("cpu")]) + def test_rejects_native_encoder_without_quantized_layers(self): with self.assertRaisesRegex( ComponentCheckpointUnsupportedError,