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

[diffusion] refactor: gate native encoder quantized checkpoints (#35183)

Browse Source 
Co-authored-by: Yiqi Yang <yangyiqi8787@gmail.com>
This commit is contained in:
Mick
2026-08-19 16:18:36 +08:00
committed by GitHub
co-authored by Yiqi Yang
parent ccbe380028
commit 73e5fa4724
7 changed files with 241 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
+5 -1
View File
@@ -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):
python/sglang/multimodal_gen/runtime/loader/component_loaders/image_encoder_loader.py
+7
View File
@@ -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,
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+86 -30
View File
@@ -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:
python/sglang/multimodal_gen/runtime/models/encoders/base.py
+14 -1
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py
+10 -3
View File
@@ -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:
python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py
+73
View File
@@ -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()
python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py
+46 -10
View File
@@ -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"))