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

[diffusion] fix: route quantized vae component repos safely (#35184)

Browse Source
This commit is contained in:
Mick
2026-08-19 21:02:46 +08:00
committed by GitHub
parent 29f5d1c7c3
commit e3445ed2bd
3 changed files with 181 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py
+10 -1
View File
@@ -46,6 +46,10 @@ class ComponentCheckpointUnsupportedError(ValueError):
"""A component checkpoint is unsupported and must not use native fallback."""
class NativeComponentLoaderRequired(RuntimeError):
"""The customized loader must defer to the native library loader."""
def _load_auto_tokenizer_with_roberta_processing_compat(*args, **kwargs):
from tokenizers import processors
@@ -198,13 +202,18 @@ class ComponentLoader(ABC):
except (ComponentCheckpointUnsupportedError, ComponentResidencyError):
raise
except Exception as e:
native_loader_required = isinstance(e, NativeComponentLoaderRequired)
if self.should_raise_customized_load_error(server_args, component_name):
if native_loader_required:
raise
traceback.print_exc()
raise RuntimeError(
f"Failed to load customized {component_name}; native fallback "
"is disabled for this component configuration."
) from e
if "Unsupported model architecture" in str(e):
if native_loader_required:
logger.info("%s", e)
elif "Unsupported model architecture" in str(e):
logger.info(
f"Component: {component_name} doesn't have a customized version yet, using native version"
)
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+40 -5
View File
@@ -11,7 +11,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
)
from sglang.multimodal_gen.configs.pipeline_configs.wan import WanT2V480PConfig
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentCheckpointUnsupportedError,
ComponentLoader,
NativeComponentLoaderRequired,
)
from sglang.multimodal_gen.runtime.loader.utils import (
_list_safetensors_files,
@@ -28,11 +30,41 @@ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
from sglang.srt.model_loader.checkpoint_quantization import (
resolve_checkpoint_quant_spec,
)
logger = init_logger(__name__)
VAE_CHANNELS_LAST_3D_ENV = "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D"
def _require_native_loader_for_quantized_vae(
config: dict, component_name: str, *, native_only: bool = False
) -> None:
quant_spec = resolve_checkpoint_quant_spec(config)
if quant_spec is None:
return
method = quant_spec.declared_method or "unspecified"
if native_only:
raise ComponentCheckpointUnsupportedError(
f"{component_name} uses a native-only SGLang implementation that "
f"cannot restore quant_method={method!r}; Diffusers fallback is disabled."
)
if quant_spec.source != "quantization_config":
raise ComponentCheckpointUnsupportedError(
f"{component_name} checkpoint declares quantization metadata in "
f"{quant_spec.source} (quant_method={method!r}), which the Diffusers "
"component loader does not restore automatically."
)
raise NativeComponentLoaderRequired(
f"{component_name} checkpoint declares quant_method={method!r}; routing "
"through Diffusers from_pretrained because the SGLang VAE loader cannot "
"restore serialized quantized state."
)
def _backfill_ltx2_audio_vae_latent_stats(
loaded: dict[str, torch.Tensor], component_name: str
) -> None:
@@ -119,13 +151,19 @@ class VAELoader(ComponentLoader):
):
"""Load the VAE based on the model path, and inference args."""
config = get_diffusers_component_config(component_path=component_model_path)
server_args.model_paths[component_name] = component_model_path
native_only = component_name in getattr(
server_args.pipeline_config, "native_only_components", ()
)
_require_native_loader_for_quantized_vae(
config, component_name, native_only=native_only
)
class_name = config.pop("_class_name", None)
assert (
class_name is not None
), "Model config does not contain a _class_name attribute. Only diffusers format is supported."
server_args.model_paths[component_name] = component_model_path
if component_name in ("vae", "video_vae"):
pipeline_vae_config_attr = "vae_config"
pipeline_vae_precision = "vae_precision"
@@ -155,9 +193,6 @@ class VAELoader(ComponentLoader):
)
target_device = self.target_device(component_starts_on_cpu)
native_only = component_name in getattr(
server_args.pipeline_config, "native_only_components", ()
)
auto_map = config.get("auto_map", {})
auto_model_map = auto_map.get("AutoModel")
if auto_model_map and not native_only:
python/sglang/multimodal_gen/test/unit/test_vae_loader.py
+131
View File
@@ -1,7 +1,9 @@
import unittest
from tempfile import TemporaryDirectory
from unittest.mock import patch
import torch
import torch.nn as nn
from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
@@ -13,8 +15,12 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import (
WanT2V480PConfig,
)
from sglang.multimodal_gen.runtime.loader.component_loaders import vae_loader
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentCheckpointUnsupportedError,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import (
_backfill_ltx2_audio_vae_latent_stats,
_require_native_loader_for_quantized_vae,
_should_use_channels_last_3d,
)
from sglang.multimodal_gen.runtime.models.vaes import wanvae
@@ -24,9 +30,134 @@ class _FakeServerArgs:
def __init__(self, pipeline_config, num_gpus=1):
self.pipeline_config = pipeline_config
self.num_gpus = num_gpus
self.model_paths = {}
self.revision = "test-revision"
self.trust_remote_code = True
def resolve_component_attention_backend(self, _component_name):
return None, None
def should_start_component_on_cpu(self, _component_name):
return False
class TestVAELoader(unittest.TestCase):
def test_quantized_vae_admission_leaves_plain_configs_unchanged(self):
_require_native_loader_for_quantized_vae(
{"_class_name": "AutoencoderKL"}, "vae"
)
with self.assertRaisesRegex(
ComponentCheckpointUnsupportedError, "compression_config"
):
_require_native_loader_for_quantized_vae(
{
"_class_name": "AutoencoderKL",
"compression_config": {"quant_method": "compressed-tensors"},
},
"vae",
)
with self.assertRaisesRegex(
ComponentCheckpointUnsupportedError,
r"text_config\.quantization_config",
):
_require_native_loader_for_quantized_vae(
{
"_class_name": "AutoencoderKL",
"text_config": {
"quantization_config": {
"quant_method": "bitsandbytes",
"load_in_4bit": True,
}
},
},
"vae",
)
def test_quantized_vae_routes_to_diffusers_native_loader(self):
loader = vae_loader.VAELoader()
server_args = _FakeServerArgs(QwenImagePipelineConfig())
native_vae = nn.Linear(1, 1)
with (
TemporaryDirectory() as component_path,
patch.object(
vae_loader,
"get_diffusers_component_config",
return_value={
"_class_name": "AutoencoderKL",
"quantization_config": {
"quant_method": "bitsandbytes",
"load_in_4bit": True,
},
},
),
patch(
"diffusers.AutoModel.from_pretrained",
return_value=native_vae,
) as native_load,
patch.object(loader, "target_device", return_value=torch.device("cpu")),
patch.object(native_vae, "to", wraps=native_vae.to) as module_to,
patch.object(
vae_loader.current_platform,
"get_available_gpu_memory",
side_effect=[10.0, 9.0],
),
patch(
"sglang.multimodal_gen.runtime.loader.component_loaders."
"component_loader.get_memory_usage_of_component",
return_value=1.0,
),
):
loaded, consumed = loader.load(
component_path, server_args, "vae", "diffusers"
)
self.assertIs(loaded, native_vae)
self.assertFalse(loaded.training)
self.assertEqual(consumed, 1.0)
self.assertEqual(server_args.model_paths["vae"], component_path)
native_load.assert_called_once_with(
component_path,
revision="test-revision",
trust_remote_code=True,
torch_dtype=torch.bfloat16,
)
module_to.assert_called_once_with(torch.device("cpu"))
def test_native_only_quantized_vae_fails_closed(self):
pipeline_config = QwenImagePipelineConfig()
pipeline_config.native_only_components = ("vae",)
server_args = _FakeServerArgs(pipeline_config)
loader = vae_loader.VAELoader()
with (
patch.object(
vae_loader,
"get_diffusers_component_config",
return_value={
"_class_name": "AutoencoderKL",
"quantization_config": {
"quant_method": "bitsandbytes",
"load_in_4bit": True,
},
},
),
patch("diffusers.AutoModel.from_pretrained") as native_load,
patch.object(
vae_loader.current_platform,
"get_available_gpu_memory",
return_value=10.0,
),
):
with self.assertRaisesRegex(
ComponentCheckpointUnsupportedError, "native-only SGLang"
):
loader.load("/quantized/vae", server_args, "vae", "diffusers")
native_load.assert_not_called()
def test_backfill_ltx2_audio_vae_latent_stats_maps_official_keys(self):
loaded = {
"per_channel_statistics.mean-of-means": torch.tensor([1.0, 2.0]),