From e3445ed2bda12c878b234a1b77be68b87c619679 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 19 Aug 2026 21:02:46 +0800 Subject: [PATCH] [diffusion] fix: route quantized vae component repos safely (#35184) --- .../component_loaders/component_loader.py | 11 +- .../loader/component_loaders/vae_loader.py | 45 +++++- .../test/unit/test_vae_loader.py | 131 ++++++++++++++++++ 3 files changed, 181 insertions(+), 6 deletions(-) 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 8eb926981..ce95535c7 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 @@ -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" ) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 411964068..ab676eda5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -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: diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index 4aa7ee1e3..26c25904e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py @@ -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]),