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

[diffusion] feat: delegate recognized quantized components to transformers (#36902)

Browse Source
This commit is contained in:
Mick
2026-08-29 14:26:50 +08:00
committed by GitHub
parent 3c1d77be21
commit d1ce017665
5 changed files with 260 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
+48 -17
View File
@@ -19,6 +19,7 @@ from transformers import (
AutoTokenizer,
PretrainedConfig,
)
from transformers.quantizers import AutoHfQuantizer
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.layers.attention.selector import (
@@ -32,6 +33,7 @@ from sglang.multimodal_gen.runtime.loader.utils import (
get_memory_usage_of_component,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
RESIDENT,
ComponentResidencyError,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_strategies import (
@@ -65,15 +67,15 @@ class NativeComponentLoaderRequired(RuntimeError):
"""The customized loader must defer to the native library loader."""
def uses_native_transformers_bnb4(config: object, component_name: str) -> bool:
"""Validate a serialized BnB4 checkpoint owned by Transformers."""
def uses_native_transformers_quantization(config: object, component_name: str) -> bool:
"""Validate quantization metadata that Transformers can restore itself."""
try:
quant_spec = resolve_checkpoint_quant_spec(config)
except (TypeError, ValueError) as error:
raise ComponentCheckpointUnsupportedError(
f"Cannot parse checkpoint quantization for {component_name!r}: {error}"
) from error
if quant_spec is None or quant_spec.declared_method != "bitsandbytes":
if quant_spec is None:
return False
if quant_spec.source != "quantization_config":
raise ComponentCheckpointUnsupportedError(
@@ -82,16 +84,18 @@ def uses_native_transformers_bnb4(config: object, component_name: str) -> bool:
f"got metadata from {quant_spec.source!r}"
)
load_in_4bit = quant_spec.config.get(
"load_in_4bit", quant_spec.config.get("_load_in_4bit")
)
load_in_8bit = quant_spec.config.get(
"load_in_8bit", quant_spec.config.get("_load_in_8bit", False)
)
if load_in_4bit is not True or load_in_8bit is True:
try:
supported = AutoHfQuantizer.supports_quant_method(dict(quant_spec.config))
except (TypeError, ValueError) as error:
raise ComponentCheckpointUnsupportedError(
f"Transformers-managed {component_name!r} quantization supports only "
"serialized BitsAndBytes 4-bit checkpoints"
f"Cannot configure Transformers-managed quantization for "
f"{component_name!r}: {error}"
) from error
if not supported:
method = quant_spec.declared_method or "unspecified"
raise ComponentCheckpointUnsupportedError(
f"Transformers does not support quant_method={method!r} declared by "
f"{component_name!r}"
)
return True
@@ -147,6 +151,7 @@ class ComponentLoader(ABC):
def __init__(self, device=None) -> None:
self.device = device
self.component_architecture: str | None = None
self._native_load_manages_placement = False
@staticmethod
def target_device(component_starts_on_cpu: bool) -> torch.device:
@@ -230,6 +235,7 @@ class ComponentLoader(ABC):
If all of the above methods failed, an error will be thrown
"""
self._native_load_manages_placement = False
component_quantization = server_args.component_quantizations.get(component_name)
if (
component_quantization is not None
@@ -317,7 +323,10 @@ class ComponentLoader(ABC):
else:
if isinstance(component, nn.Module):
component = component.eval()
if not is_fsdp_managed_module(component):
if (
not is_fsdp_managed_module(component)
and not self._native_load_manages_placement
):
component = component.to(
self.target_device(
server_args.should_start_component_on_cpu(component_name)
@@ -356,16 +365,38 @@ class ComponentLoader(ABC):
load_kwargs["torch_dtype"] = precision
if transformers_or_diffusers == "transformers":
self._native_load_manages_placement = False
config = get_hf_config(
component_model_path,
trust_remote_code=server_args.trust_remote_code,
revision=server_args.revision,
)
if uses_native_transformers_bnb4(config, component_name or "component"):
server_args.require_component_resident(
component_name or "component",
feature_name="Transformers bitsandbytes component",
if uses_native_transformers_quantization(
config, component_name or "component"
):
resolved_component_name = component_name or "component"
explicit_residency = server_args.explicit_residency_mode(
resolved_component_name
)
if explicit_residency is not None and explicit_residency != RESIDENT:
raise ComponentCheckpointUnsupportedError(
"Transformers-managed quantized component "
f"{resolved_component_name!r} requires resident placement; "
f"got explicit mode {explicit_residency!r}"
)
server_args.require_component_resident(
resolved_component_name,
feature_name="Transformers quantized component",
)
if server_args.should_use_fsdp_for_component(resolved_component_name):
raise ComponentCheckpointUnsupportedError(
"Transformers-managed quantized components do not support "
"SGLang FSDP loading"
)
load_kwargs["device_map"] = {
"": self.target_device(component_starts_on_cpu=False)
}
self._native_load_manages_placement = True
model_class = self.resolve_native_transformers_model_class(config)
return model_class.from_pretrained(
component_model_path,
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+25 -8
View File
@@ -58,7 +58,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp
ComponentCheckpointUnsupportedError,
ComponentLoader,
NativeComponentLoaderRequired,
uses_native_transformers_bnb4,
uses_native_transformers_quantization,
)
from sglang.multimodal_gen.runtime.loader.gguf_weights import (
gguf_weights_iterator,
@@ -135,15 +135,23 @@ _TRANSFORMERS_ENCODER_ONLY_CLASSES = {
}
def _delegate_standard_bnb4_to_transformers(
def _delegate_quantized_checkpoint_to_transformers(
component_config: dict,
component_name: str,
*,
methods: frozenset[str] | None = None,
) -> None:
"""Use Transformers when it owns a standard serialized BnB4 checkpoint."""
if uses_native_transformers_bnb4(component_config, component_name):
"""Use Transformers when it owns the checkpoint's serialized format."""
quant_spec = resolve_checkpoint_quant_spec(component_config)
if quant_spec is None or (
methods is not None and quant_spec.declared_method not in methods
):
return
if uses_native_transformers_quantization(component_config, component_name):
method = quant_spec.declared_method or "unspecified"
raise NativeComponentLoaderRequired(
f"{component_name!r} delegates serialized bitsandbytes checkpoint "
"loading to Transformers"
f"{component_name!r} delegates serialized quant_method={method!r} "
"checkpoint loading to Transformers"
)
@@ -268,9 +276,10 @@ def _configure_encoder_quantization(
# themselves; running the generic lifecycle as well would process twice.
return
_delegate_standard_bnb4_to_transformers(
_delegate_quantized_checkpoint_to_transformers(
component_config,
component_name,
methods=frozenset({"bitsandbytes"}),
)
try:
quant_config = _get_encoder_quant_config(
@@ -280,6 +289,10 @@ def _configure_encoder_quantization(
model_cls,
)
except (KeyError, NotImplementedError, TypeError, ValueError) as error:
_delegate_quantized_checkpoint_to_transformers(
component_config,
component_name,
)
raise ComponentCheckpointUnsupportedError(
f"Cannot configure checkpoint quantization for {component_name!r}: {error}"
) from error
@@ -305,6 +318,10 @@ def _configure_encoder_quantization(
)
quant_config = model_config.quant_config
if quant_config is None:
_delegate_quantized_checkpoint_to_transformers(
component_config,
component_name,
)
return
if not issubclass(model_cls, EncoderTensorParallelMixin):
raise ComponentCheckpointUnsupportedError(
@@ -327,7 +344,7 @@ def _resolve_and_configure_encoder_quantization(
try:
model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
except Exception as resolution_error:
_delegate_standard_bnb4_to_transformers(
_delegate_quantized_checkpoint_to_transformers(
component_config,
component_name,
)
python/sglang/multimodal_gen/test/unit/test_image_encoder_loader.py
+138 -2
View File
@@ -3,10 +3,12 @@ from types import SimpleNamespace
from unittest import mock
import torch
from torch import nn
from sglang.multimodal_gen.configs.models.encoders.clip import CLIPVisionConfig
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentCheckpointUnsupportedError,
NativeComponentLoaderRequired,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.image_encoder_loader import (
ImageEncoderLoader,
@@ -14,6 +16,9 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.image_encoder_loader
from sglang.multimodal_gen.runtime.loader.component_loaders.text_encoder_loader import (
_configure_encoder_quantization,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
COMPONENT_OFFLOAD,
)
from sglang.multimodal_gen.runtime.models.encoders.clip import CLIPVisionModel
from sglang.srt.layers.quantization.fp8 import Fp8Config as SRTFp8Config
@@ -76,8 +81,26 @@ class TestImageEncoderQuantizationAdmission(unittest.TestCase):
{"qkv_proj": ["q_proj", "k_proj", "v_proj"]},
)
def test_unknown_quantized_architecture_does_not_fall_back(self):
def test_unknown_transformers_quantized_architecture_falls_back(self):
config = self._component_config("UnknownVisionModel", quantized=True)
with self._config_patch(config):
self._load()
self.load_native.assert_called_once()
def test_native_only_quantized_architecture_does_not_fall_back(self):
self.server_args.pipeline_config.native_only_components = ("image_encoder",)
config = self._component_config("UnknownVisionModel", quantized=True)
with self._config_patch(config), self.assertRaises(
NativeComponentLoaderRequired
):
self._load()
self.load_native.assert_not_called()
def test_unknown_unsupported_quantized_architecture_does_not_fall_back(self):
config = {
"architectures": ["UnknownVisionModel"],
"quantization_config": {"quant_method": "not-a-format"},
}
with self._config_patch(config), self.assertRaises(
ComponentCheckpointUnsupportedError
):
@@ -107,7 +130,9 @@ class TestImageEncoderNativeLoading(unittest.TestCase):
)
server_args = SimpleNamespace(
pipeline_config=SimpleNamespace(image_encoder_precision="bf16"),
explicit_residency_mode=mock.Mock(return_value=None),
require_component_resident=mock.Mock(),
should_use_fsdp_for_component=mock.Mock(return_value=False),
revision=None,
trust_remote_code=False,
)
@@ -121,6 +146,10 @@ class TestImageEncoderNativeLoading(unittest.TestCase):
loader,
"resolve_native_transformers_model_class",
return_value=model_class,
), mock.patch.object(
loader,
"target_device",
return_value=torch.device("cuda:0"),
):
component = loader.load_native(
"/model/image_encoder",
@@ -132,7 +161,7 @@ class TestImageEncoderNativeLoading(unittest.TestCase):
self.assertIs(component, loaded_encoder)
server_args.require_component_resident.assert_called_once_with(
"image_encoder",
feature_name="Transformers bitsandbytes component",
feature_name="Transformers quantized component",
)
model_class.from_pretrained.assert_called_once_with(
"/model/image_encoder",
@@ -140,4 +169,111 @@ class TestImageEncoderNativeLoading(unittest.TestCase):
trust_remote_code=False,
revision=None,
torch_dtype=torch.bfloat16,
device_map={"": torch.device("cuda:0")},
)
def test_explicit_offload_is_rejected_before_transformers_load(self):
component_config = SimpleNamespace(
is_encoder_decoder=False,
architectures=["ThirdPartyVisionModel"],
quantization_config={"quant_method": "fp8"},
)
model_class = SimpleNamespace(from_pretrained=mock.Mock())
server_args = SimpleNamespace(
pipeline_config=SimpleNamespace(image_encoder_precision="bf16"),
explicit_residency_mode=mock.Mock(return_value=COMPONENT_OFFLOAD),
require_component_resident=mock.Mock(),
should_use_fsdp_for_component=mock.Mock(return_value=False),
revision=None,
trust_remote_code=False,
)
loader = ImageEncoderLoader()
with mock.patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.get_hf_config",
return_value=component_config,
), mock.patch.object(
loader,
"resolve_native_transformers_model_class",
return_value=model_class,
), self.assertRaisesRegex(
ComponentCheckpointUnsupportedError, "requires resident placement"
):
loader.load_native(
"/model/image_encoder",
server_args,
"transformers",
"image_encoder",
)
model_class.from_pretrained.assert_not_called()
server_args.require_component_resident.assert_not_called()
def test_full_loader_does_not_move_quantized_component_again(self):
class RejectMoveModule(nn.Module):
def to(self, *args, **kwargs):
raise AssertionError("quantized component must not be moved again")
component_config = SimpleNamespace(
is_encoder_decoder=False,
architectures=["ThirdPartyVisionModel"],
quantization_config={"quant_method": "fp8"},
)
loaded_encoder = RejectMoveModule()
model_class = SimpleNamespace(
from_pretrained=mock.Mock(return_value=loaded_encoder)
)
server_args = SimpleNamespace(
component_quantizations={},
pipeline_config=SimpleNamespace(
image_encoder_precision="bf16",
native_only_components=(),
),
resolve_component_attention_backend=lambda _name: (None, None),
explicit_residency_mode=lambda _name: None,
require_component_resident=mock.Mock(),
should_use_fsdp_for_component=lambda _name: False,
should_start_component_on_cpu=lambda _name: False,
revision=None,
trust_remote_code=False,
)
loader = ImageEncoderLoader()
with mock.patch.object(
loader,
"load_customized",
side_effect=NativeComponentLoaderRequired("use Transformers"),
), mock.patch.object(
loader,
"resolve_native_transformers_model_class",
return_value=model_class,
), mock.patch.object(
loader,
"target_device",
return_value=torch.device("cuda:0"),
), mock.patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.get_hf_config",
return_value=component_config,
), mock.patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.current_platform.get_available_gpu_memory",
return_value=10.0,
), mock.patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.get_memory_usage_of_component",
return_value=0.0,
), mock.patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.format_component_residency",
return_value="resident",
):
component, _ = loader.load(
"/model/image_encoder",
server_args,
"image_encoder",
"transformers",
)
self.assertIs(component, loaded_encoder)
python/sglang/multimodal_gen/test/unit/test_text_encoder_loader.py
+34 -6
View File
@@ -151,7 +151,9 @@ class TestTextEncoderClassResolution(unittest.TestCase):
)
server_args = SimpleNamespace(
pipeline_config=SimpleNamespace(text_encoder_precisions=["bf16"]),
explicit_residency_mode=mock.Mock(return_value=None),
require_component_resident=mock.Mock(),
should_use_fsdp_for_component=mock.Mock(return_value=False),
revision=None,
trust_remote_code=False,
)
@@ -162,16 +164,21 @@ class TestTextEncoderClassResolution(unittest.TestCase):
}
}
loader = TextEncoderLoader()
with mock.patch.object(
TextEncoderLoader,
"resolve_native_transformers_model_class",
return_value=transformers_model_class,
), mock.patch.object(
loader,
"target_device",
return_value=torch.device("cuda:0"),
), mock.patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.get_hf_config",
return_value=component_config,
):
encoder = TextEncoderLoader().load_native(
encoder = loader.load_native(
"/model/text_encoder",
server_args,
"transformers",
@@ -181,7 +188,7 @@ class TestTextEncoderClassResolution(unittest.TestCase):
self.assertIs(encoder, loaded_encoder)
server_args.require_component_resident.assert_called_once_with(
"text_encoder",
feature_name="Transformers bitsandbytes component",
feature_name="Transformers quantized component",
)
transformers_model_class.from_pretrained.assert_called_once_with(
"/model/text_encoder",
@@ -189,6 +196,7 @@ class TestTextEncoderClassResolution(unittest.TestCase):
trust_remote_code=False,
revision=None,
torch_dtype=torch.bfloat16,
device_map={"": torch.device("cuda:0")},
)
@@ -712,7 +720,7 @@ class TestTextEncoderQuantization(unittest.TestCase):
):
with self.subTest(architecture=architecture), self.assertRaisesRegex(
NativeComponentLoaderRequired,
"delegates serialized bitsandbytes checkpoint loading to Transformers",
"delegates serialized quant_method='bitsandbytes' checkpoint",
):
_resolve_and_configure_encoder_quantization(
SimpleNamespace(architectures=[architecture], quant_config=None),
@@ -742,10 +750,10 @@ class TestTextEncoderQuantization(unittest.TestCase):
"text_encoder",
)
def test_rejects_bitsandbytes_8bit(self):
def test_bitsandbytes_8bit_delegates_to_transformers(self):
with self.assertRaisesRegex(
ComponentCheckpointUnsupportedError,
"supports only serialized BitsAndBytes 4-bit checkpoints",
NativeComponentLoaderRequired,
"delegates serialized quant_method='bitsandbytes' checkpoint",
):
_resolve_and_configure_encoder_quantization(
SimpleNamespace(
@@ -763,6 +771,26 @@ class TestTextEncoderQuantization(unittest.TestCase):
"text_encoder",
)
def test_unknown_fp8_architecture_delegates_to_transformers(self):
with self.assertRaisesRegex(
NativeComponentLoaderRequired,
"delegates serialized quant_method='fp8' checkpoint",
):
_resolve_and_configure_encoder_quantization(
SimpleNamespace(
architectures=["ThirdPartyTextEncoder"], quant_config=None
),
{
"quantization_config": {
"quant_method": "fp8",
"activation_scheme": "dynamic",
}
},
"/model/text_encoder",
"/model/text_encoder",
"text_encoder",
)
def test_model_managed_quantization_bypasses_generic_lifecycle(self):
model_config = SimpleNamespace(quant_config=None)
with mock.patch.object(