From b26695a26e4863265cd2c0267140f1a7026b9064 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 22 Aug 2026 09:33:43 +0800 Subject: [PATCH] [diffusion] feat: reject unsupported quantized component checkpoints (#35873) --- docs/docs/sglang-diffusion/api/cli.mdx | 30 +++- .../sglang-diffusion/compatibility_matrix.mdx | 13 +- docs/docs/sglang-diffusion/index.mdx | 2 +- docs/docs/sglang-diffusion/quantization.mdx | 57 +++++-- .../component_loaders/adapter_loader.py | 9 +- .../loader/component_loaders/bridge_loader.py | 9 +- .../component_loaders/component_loader.py | 34 ++++ .../diffusion_decoder_loader.py | 9 +- .../sound_tokenizer_loader.py | 9 +- .../component_loaders/upsampler_loader.py | 31 ++-- .../component_loaders/vocoder_loader.py | 9 +- .../test_component_quantization_admission.py | 146 ++++++++++++++++++ 12 files changed, 297 insertions(+), 61 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index e47a5d830..a869f96b2 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -122,21 +122,31 @@ For frame interpolation and upscaling, see [Post-Processing](./post_processing). ### Quantization -For quantized transformer checkpoints, prefer: +For a pre-quantized main transformer checkpoint, prefer: - `--model-path` for the base pipeline - `--transformer-path` for a quantized `transformers` transformer component folder - `--transformer-weights-path` for a quantized safetensors file, directory, repo, or a supported GGUF transformer file -- `--quantization` for online quantization (apply quantization to unquantized models at load time, activations are quantized dynamically) -- `--quantization-ignored-layers` layer name patterns to keep unquantized (e.g. `attention.to_`) +- `--quantization` to override the quantization method used by the transformer loader +- `--quantization-ignored-layers` transformer layer name patterns to keep unquantized during online quantization (e.g. `attention.to_`) Component checkpoint paths are selected separately, so changing DiT precision -never silently changes prompt embeddings. For a native text encoder: +never silently changes prompt embeddings. Use +`--component-paths. {MODEL}` or its shorter +`---path {MODEL}` form for any component key from `model_index.json` +or a native pipeline's registered module name. +For a native text encoder: - `--component-paths.text_encoder {MODEL}` replaces the text-encoder checkpoint; `--text-encoder-path {MODEL}` is its shorter alias - Quantization metadata is auto-detected from that checkpoint. Each native encoder must explicitly support the serialized format; this is not blanket quantization support for every component, and unsupported combinations fail before weight loading. +The same contract applies to every weighted component: path routing is generic, +while quantized materialization is capability-based. Native auxiliary loaders +that only understand unquantized state dicts reject quantization metadata before +model construction. See [Quantized Component Repositories](../quantization#quantized-component-repositories) +for the current component matrix. + For supported realtime causal video models, `--kv-cache-quant {off|int4|int2}` compresses completed KV-cache chunks independently of transformer weight quantization. It is lossy and disabled by default. @@ -332,7 +342,17 @@ sglang serve \ --vae-path fal/FLUX.2-Tiny-AutoEncoder ``` -The component key must match the key in the model's `model_index.json`, and the path must be either a Hugging Face repo ID or a complete component directory. +The component key must match a key in the model's `model_index.json` or a native +pipeline's registered module name. The path must be either a Hugging Face repo +ID or a complete component directory. + +Path selection and quantized checkpoint support are separate capabilities. A +pre-quantized component repo must carry its quantization metadata, and the +selected loader must support that serialized format. Native plain-state loaders +fail closed; library-managed components inherit the corresponding Transformers +or Diffusers support. The transformer-specific `--quantization` flag does not +select the format of component checkpoints; their own metadata does. See +[Quantized Component Repositories](../quantization#quantized-component-repositories). ## Component Attention Backend Overrides diff --git a/docs/docs/sglang-diffusion/compatibility_matrix.mdx b/docs/docs/sglang-diffusion/compatibility_matrix.mdx index 9329d2fbc..2426e96c4 100644 --- a/docs/docs/sglang-diffusion/compatibility_matrix.mdx +++ b/docs/docs/sglang-diffusion/compatibility_matrix.mdx @@ -716,6 +716,13 @@ component directory. The same overrides can also be provided in config files through `component_paths.`. +Component path overrides are generic, but quantized checkpoint support is +loader-specific. Native SGLang loaders that only materialize plain state dicts +reject detected quantization metadata before weight loading instead of silently +casting the checkpoint into an unquantized module. Library-managed components +inherit the corresponding Transformers or Diffusers support. See +[Quantized Component Repositories](./quantization#quantized-component-repositories). + ### Common Syntax CLI: @@ -818,7 +825,7 @@ The table below lists concrete Hugging Face component repos that are already use ### Transformer / DiT - `--transformer-path` is the standard override for the main denoising transformer. -- For quantized transformers, prefer `--transformer-path` or `--transformer-weights-path`; see `quantization.md`. +- For a pre-quantized main transformer, prefer `--transformer-path` or `--transformer-weights-path`; see [Quantization](./quantization). - `--video-dit-path` and `--audio-dit-path` are only for pipelines that split denoisers by modality. ### Text Encoders and Preprocessors @@ -838,6 +845,10 @@ The table below lists concrete Hugging Face component repos that are already use that component. 2. The override key should match the component name in the pipeline's `model_index.json` or the native pipeline's registered module name. +3. An override path does not enable quantization by itself. A pre-quantized + component must include metadata for a format supported by its selected + loader. `--quantization` overrides the transformer loader; other component + checkpoints select their serialized format through their own metadata. ## Verified LoRA Examples diff --git a/docs/docs/sglang-diffusion/index.mdx b/docs/docs/sglang-diffusion/index.mdx index 84cfd4b6d..0d2c54795 100644 --- a/docs/docs/sglang-diffusion/index.mdx +++ b/docs/docs/sglang-diffusion/index.mdx @@ -35,7 +35,7 @@ sglang serve --model-path Qwen/Qwen-Image --port 30010 - [OpenAI-Compatible API](/docs/sglang-diffusion/api/openai_api): send image and video requests to the HTTP server - [Performance Overview](/docs/sglang-diffusion/performance-optimization): choose speed, memory, parallelism, caching, and quality-tradeoff levers - [Caching Acceleration](/docs/sglang-diffusion/caching-acceleration): use Cache-DiT, TeaCache, or Spectrum to reduce denoising cost -- [Quantization](/docs/sglang-diffusion/quantization): configure transformer weight and causal KV-cache quantization +- [Quantization](/docs/sglang-diffusion/quantization): configure component checkpoint and causal KV-cache quantization - [Realtime and Causal Video Models](/docs/sglang-diffusion/realtime_models): understand session state, causal caches, and realtime-only controls - [Contributing](/docs/sglang-diffusion/contributing): contribution workflow, adding new models, and CI perf baselines diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index a3b25bf80..88bed2a8e 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -2,11 +2,13 @@ title: "Quantization" tag: "approx" metatags: - description: "Configure transformer and native text-encoder checkpoint quantization, plus Quant-VideoGen causal KV-cache quantization in SGLang-Diffusion." + description: "Configure quantized component repositories, transformer and native text-encoder checkpoint quantization, plus Quant-VideoGen causal KV-cache quantization in SGLang-Diffusion." --- -SGLang-Diffusion supports quantized transformer checkpoints and selected -quantized native text-encoder checkpoints. Transformer and text-encoder +SGLang-Diffusion treats component path selection and quantized checkpoint +materialization as separate capabilities. Every loaded component can use an +independent checkpoint path, but that checkpoint is quantized only when its +selected loader supports the serialized format. Transformer, encoder, and VAE precision are resolved independently. ## Quick Reference @@ -14,12 +16,13 @@ precision are resolved independently. Use these paths: - `--model-path`: the base or original model +- `--component-paths.` / `---path`: replace a component from `model_index.json` or a native registered module with an independent repo or local directory - `--transformer-path`: a quantized transformers-style transformer component directory that already contains its own `config.json` - `--transformer-weights-path`: replacement transformer weights in safetensors format (file, directory, or Hub repository/file) or a supported GGUF file (local or Hub) -- `--quantization`: apply online quantization to unquantized models at load time (activations are quantized dynamically) -- `--quantization-ignored-layers` layer name patterns to keep unquantized (e.g. `attention.to_`) +- `--quantization`: override the quantization method used by the transformer loader +- `--quantization-ignored-layers`: transformer layer name patterns to keep unquantized during online quantization (e.g. `attention.to_`) - `--component-paths.text_encoder`: replace a native text encoder with a checkpoint whose `quantization_config` is auto-detected - `--text-encoder-path`: shorter alias for `--component-paths.text_encoder` - `--kv-cache-quant`: compress completed causal KV-cache chunks for supported realtime models @@ -53,6 +56,33 @@ path, `owner/repo/path/file.safetensors`, or direct Hugging Face file URL; do not combine it with `--quantization`. MiniMax-H3 GGUF usage is documented in the [MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#pre-quantized-gguf-transformer). +## Quantized Component Repositories + +Every loaded component can point to an independent repository, but path routing +does not imply that every loader can materialize every quantization format. +SGLang resolves quantized component checkpoints through one of three explicit +paths: + +- load through an SGLang quantization implementation; +- delegate a standard component to Transformers or Diffusers; +- fail closed in a native plain-state loader when the format cannot be restored. + +| Component path | Quantized checkpoint behavior | +| --- | --- | +| `transformer`, `transformer_2`, `unconditional_transformer`, `audio_dit`, `video_dit` | Uses the SGLang transformer quantization adapters documented below. | +| `text_encoder*`, `image_encoder*` | Requires the native encoder class to declare support for the detected format. MiniMax-H3 FP8 and model-managed integrations such as Ideogram are supported; unknown combinations fail closed. | +| `vae`, `video_vae`, `audio_vae` | A standard top-level Diffusers `quantization_config` is delegated to `AutoModel.from_pretrained`. Native-only VAEs and nested/compression metadata fail closed. | +| Library-managed Transformers or Diffusers components | Delegates to the upstream `from_pretrained` path and inherits its format support and validation behavior. The local PE model uses this path; compatible formats remain model-specific. | +| Native auxiliary components that load raw state dicts | Quantized checkpoints are rejected before model construction until that component has a quantized materialization implementation. This includes connectors, duration heads, bridges, diffusion decoders, sound tokenizers, spatial upsamplers, and vocoders. | + +`--quantization` is the explicit override for the transformer loader; it is not +the boundary of component quantization support. Other pre-quantized component +repositories select their format through their own metadata and the capability +of the selected loader. A generic string override without a matching +materialization backend would advertise support that the component does not +have, while a quantized weight file without matching config metadata cannot be +identified or restored generically. + ## Quant Families Here, `quant_family` means a checkpoint and loading family with shared CLI @@ -99,7 +129,7 @@ backend. fp8 (offline quantization) Quantized transformer component folder, or safetensors with quantization_config metadata --transformer-path or --transformer-weights-path - ALL + Native DiTs whose linear layers support the selected FP8 method; validate quality per model None Component-folder and single-file flows are both supported @@ -268,11 +298,15 @@ support LongLive2 pinned sinks, global sinks, or dynamically growing caches. ## Online Quantization -Online quantization applies quantization to unquantized models at load time. This is useful for when pre-quantized checkpoints are not available. +This section describes the online methods currently implemented by the +transformer loader. They are useful when a pre-quantized transformer checkpoint +is not available. Encoders, VAEs, and auxiliary components are independent: +their repositories may still carry serialized quantized weights, which are +restored when the selected component loader supports that format. ### FP8 Online Quantization -Apply FP8 quantization to any unquantized model: +Apply FP8 quantization to a supported unquantized DiT checkpoint: ```bash sglang generate \ @@ -346,8 +380,9 @@ thresholds with `SGLANG_KITCHEN_INT8_MAX_ROWS` and ### Skipping Layers -By default, online quantization quantizes every linear layer in -the transformer. However, `--quantization-ignored-layers` can be used to keep specific layers in their original precision: +By default, transformer online quantization quantizes every supported linear +layer in that component. However, `--quantization-ignored-layers` can keep +specific transformer layers in their original precision: ```bash sglang generate \ @@ -365,7 +400,7 @@ sglang generate \ --save-output ``` -Each pattern is matched against the full layer prefix (e.g. `layers.0.attention.to_q`). A layer is skipped and left unquantizd if its prefix contains any of the given patterns. +Each pattern is matched against the full layer prefix (e.g. `layers.0.attention.to_q`). A layer is skipped and left unquantized if its prefix contains any of the given patterns. ## MiniMax-H3 Text Encoder FP8 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py index c260a73a2..2361071d3 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py @@ -7,7 +7,7 @@ from sglang.multimodal_gen.configs.models.adapter.ltx_2_duration_head import ( LTX2DurationHeadConfig, ) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( - ComponentLoader, + PlainStateDictComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( load_safetensors_state_dict, @@ -16,13 +16,10 @@ from sglang.multimodal_gen.runtime.loader.utils import ( ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( - get_diffusers_component_config, -) from sglang.multimodal_gen.runtime.utils.precision import resolve_precision -class AdapterLoader(ComponentLoader): +class AdapterLoader(PlainStateDictComponentLoader): """Loader for small adapter-style modules (e.g., LTX-2 connectors). This loader intentionally avoids FSDP sharding and just: @@ -46,7 +43,7 @@ class AdapterLoader(ComponentLoader): component_name: str = "connectors", *args, ): - config = get_diffusers_component_config(component_path=component_model_path) + config = self.load_component_config(component_model_path, component_name) cls_name = config.pop("_class_name", None) if cls_name is None: diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py index e9d15b86f..d648b00a8 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py @@ -4,7 +4,7 @@ import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( - ComponentLoader, + PlainStateDictComponentLoader, ) from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files @@ -14,16 +14,13 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( - get_diffusers_component_config, -) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.precision import resolve_precision logger = init_logger(__name__) -class BridgeLoader(ComponentLoader): +class BridgeLoader(PlainStateDictComponentLoader): """Loader for MOVA dual tower bridge with FSDP support.""" pipeline_bridge_config_attr: str = "bridge_config" @@ -34,7 +31,7 @@ class BridgeLoader(ComponentLoader): def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): - config = get_diffusers_component_config(component_path=component_model_path) + config = self.load_component_config(component_model_path, component_name) hf_config = deepcopy(config) class_name = config.pop("_class_name", None) if class_name is None: 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 6e6ce828b..f1755aa0c 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 @@ -34,11 +34,15 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency_ from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + get_diffusers_component_config, get_hf_config, prepare_diffusers_component_path_for_loading, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.precision import resolve_component_precision +from sglang.srt.model_loader.checkpoint_quantization import ( + resolve_checkpoint_quant_spec, +) logger = init_logger(__name__) @@ -441,6 +445,36 @@ class ComponentLoader(ABC): return GenericComponentLoader(transformers_or_diffusers, component_architecture) +class PlainStateDictComponentLoader(ComponentLoader): + """Base for native loaders whose current materializer expects plain weights.""" + + @staticmethod + def ensure_plain_state_dict_checkpoint(config: object, component_name: str) -> None: + try: + quant_spec = resolve_checkpoint_quant_spec(config) + except (TypeError, ValueError) as error: + raise ComponentCheckpointUnsupportedError( + f"Cannot parse checkpoint quantization metadata for " + f"{component_name!r}: {error}" + ) from error + if quant_spec is None: + return + + method = quant_spec.declared_method or "unspecified" + raise ComponentCheckpointUnsupportedError( + f"{component_name!r} checkpoint declares quantization metadata in " + f"{quant_spec.source} (quant_method={method!r}), which its current " + "plain state-dict materializer cannot restore." + ) + + def load_component_config( + self, component_model_path: str, component_name: str + ) -> dict[str, Any]: + config = get_diffusers_component_config(component_path=component_model_path) + self.ensure_plain_state_dict_checkpoint(config, component_name) + return config + + class ImageProcessorLoader(ComponentLoader): """Loader for image processor.""" diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py index 90a0bbcea..f78c55e28 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/diffusion_decoder_loader.py @@ -4,7 +4,7 @@ from sglang.multimodal_gen.configs.models.decoders.ltx_2_5_diffusion_decoder imp LTX25DiffusionDecoderConfig, ) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( - ComponentLoader, + PlainStateDictComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( load_safetensors_state_dict, @@ -13,13 +13,10 @@ from sglang.multimodal_gen.runtime.loader.utils import ( ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( - get_diffusers_component_config, -) from sglang.multimodal_gen.runtime.utils.precision import resolve_precision -class DiffusionDecoderLoader(ComponentLoader): +class DiffusionDecoderLoader(PlainStateDictComponentLoader): """Loader for the standalone, replicated LTX-2.5 diffusion decoder.""" component_names = ["diffusion_decoder"] @@ -32,7 +29,7 @@ class DiffusionDecoderLoader(ComponentLoader): component_name: str = "diffusion_decoder", *args, ): - config = get_diffusers_component_config(component_path=component_model_path) + config = self.load_component_config(component_model_path, component_name) class_name = config.pop("_class_name", None) if class_name is None: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py index 981d15c11..6ebf7926d 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py @@ -2,7 +2,7 @@ from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( - ComponentLoader, + PlainStateDictComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( _list_safetensors_files, @@ -11,23 +11,20 @@ from sglang.multimodal_gen.runtime.loader.utils import ( ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( - get_diffusers_component_config, -) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) -class SoundTokenizerLoader(ComponentLoader): +class SoundTokenizerLoader(PlainStateDictComponentLoader): component_names = ["sound_tokenizer"] expected_library = "diffusers" def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): - config = get_diffusers_component_config(component_path=component_model_path) + config = self.load_component_config(component_model_path, component_name) class_name = config.pop("_class_name", None) or self.component_architecture assert ( class_name is not None diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py index 6f119a9ce..4c9cf7a5f 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py @@ -8,7 +8,7 @@ import torch from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( - ComponentLoader, + PlainStateDictComponentLoader, ) from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import ( LatentUpsampler, @@ -158,28 +158,26 @@ def _infer_config_from_state_dict(state_dict: dict[str, torch.Tensor]) -> dict: return config -def _load_config( +def _load_explicit_config( safetensors_path: str, original_path: str, - state_dict: dict[str, torch.Tensor], -) -> dict: - """Load upsampler config with fallback chain: +) -> dict | None: + """Load an explicit upsampler config with this fallback chain: 1. safetensors metadata ("config" key) - original LTX-2 repo format 2. sibling config.json - diffusers format 3. config.json from HF (if original_path was a URL) - 4. infer from state dict shapes (always works) """ with safetensors.safe_open(safetensors_path, framework="pt") as f: meta = f.metadata() if meta and "config" in meta: logger.info("Using config from safetensors metadata") - return _normalize_config(json.loads(meta["config"])) + return json.loads(meta["config"]) config_json_path = os.path.join(os.path.dirname(safetensors_path), "config.json") if os.path.isfile(config_json_path): with open(config_json_path) as fp: logger.info("Using config from sibling config.json") - return _normalize_config(json.load(fp)) + return json.load(fp) hf = _parse_hf_url(original_path) if hf: @@ -189,15 +187,14 @@ def _load_config( local = _download_hf_file(repo_id, config_filename, revision) with open(local) as fp: logger.info("Using config from HF config.json") - return _normalize_config(json.load(fp)) + return json.load(fp) except Exception: pass - logger.info("No explicit config found, inferring from state dict") - return _infer_config_from_state_dict(state_dict) + return None -class UpsamplerLoader(ComponentLoader): +class UpsamplerLoader(PlainStateDictComponentLoader): component_names = ["spatial_upsampler"] expected_library = "diffusers" @@ -208,8 +205,16 @@ class UpsamplerLoader(ComponentLoader): component_name: str, ): safetensors_path = _find_safetensors_file(component_model_path) + raw_config = _load_explicit_config(safetensors_path, component_model_path) + if raw_config is not None: + self.ensure_plain_state_dict_checkpoint(raw_config, component_name) + state_dict = safetensors_load_file(safetensors_path) - config = _load_config(safetensors_path, component_model_path, state_dict) + if raw_config is None: + logger.info("No explicit config found, inferring from state dict") + config = _infer_config_from_state_dict(state_dict) + else: + config = _normalize_config(raw_config) logger.info("Loading LatentUpsampler with config: %s", config) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py index e2a8bc378..e91499bbf 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vocoder_loader.py @@ -3,7 +3,7 @@ import re from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( - ComponentLoader, + PlainStateDictComponentLoader, ) from sglang.multimodal_gen.runtime.loader.utils import ( _list_safetensors_files, @@ -12,9 +12,6 @@ from sglang.multimodal_gen.runtime.loader.utils import ( ) from sglang.multimodal_gen.runtime.models.registry import ModelRegistry from sglang.multimodal_gen.runtime.server_args import ServerArgs -from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( - get_diffusers_component_config, -) 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 @@ -22,14 +19,14 @@ from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) -class VocoderLoader(ComponentLoader): +class VocoderLoader(PlainStateDictComponentLoader): component_names = ["vocoder"] expected_library = "diffusers" def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): - config = get_diffusers_component_config(component_path=component_model_path) + config = self.load_component_config(component_model_path, component_name) class_name = config.pop("_class_name", None) or self.component_architecture assert ( class_name is not None diff --git a/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py b/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py new file mode 100644 index 000000000..180f54905 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_component_quantization_admission.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 + +import re +import unittest +from unittest.mock import patch + +from sglang.multimodal_gen.runtime.loader.component_loaders.adapter_loader import ( + AdapterLoader, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.bridge_loader import ( + BridgeLoader, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentCheckpointUnsupportedError, + PlainStateDictComponentLoader, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.diffusion_decoder_loader import ( + DiffusionDecoderLoader, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.sound_tokenizer_loader import ( + SoundTokenizerLoader, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.upsampler_loader import ( + UpsamplerLoader, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.vocoder_loader import ( + VocoderLoader, +) + + +class _TestLoader(PlainStateDictComponentLoader): + pass + + +class TestComponentQuantizationAdmission(unittest.TestCase): + def test_plain_checkpoint_config_is_accepted(self): + config = {"_class_name": "TestModel"} + + with patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_diffusers_component_config", + return_value=config, + ): + loaded = _TestLoader().load_component_config("/model/component", "test") + + self.assertIs(loaded, config) + + def test_all_quantization_metadata_layouts_fail_closed(self): + configs = { + "quantization_config": { + "quantization_config": {"quant_method": "bitsandbytes"} + }, + "text_config.quantization_config": { + "text_config": {"quantization_config": {"quant_method": "fp8"}} + }, + "compression_config": { + "compression_config": {"quant_method": "compressed-tensors"} + }, + } + + for source, config in configs.items(): + with ( + self.subTest(source=source), + self.assertRaisesRegex( + ComponentCheckpointUnsupportedError, + rf"{re.escape(source)}.*quant_method=.*cannot restore", + ), + ): + _TestLoader.ensure_plain_state_dict_checkpoint(config, "test_component") + + with self.assertRaisesRegex( + ComponentCheckpointUnsupportedError, + "Cannot parse checkpoint quantization metadata", + ): + _TestLoader.ensure_plain_state_dict_checkpoint( + {"quantization_config": "invalid"}, "test_component" + ) + + def test_native_raw_state_loaders_share_the_admission_boundary(self): + loader_classes = ( + AdapterLoader, + BridgeLoader, + DiffusionDecoderLoader, + SoundTokenizerLoader, + UpsamplerLoader, + VocoderLoader, + ) + + for loader_class in loader_classes: + with self.subTest(loader=loader_class.__name__): + self.assertTrue(issubclass(loader_class, PlainStateDictComponentLoader)) + + def test_adapter_rejects_quantization_before_model_construction(self): + config = { + "_class_name": "LTX2ConnectorModel", + "quantization_config": {"quant_method": "bitsandbytes"}, + } + + with ( + patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "component_loader.get_diffusers_component_config", + return_value=config, + ), + patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "adapter_loader.ModelRegistry.resolve_model_cls" + ) as resolve_model, + self.assertRaises(ComponentCheckpointUnsupportedError), + ): + AdapterLoader().load_customized("/model/connectors", None, "connectors") + + resolve_model.assert_not_called() + + def test_upsampler_rejects_quantization_before_loading_weights(self): + config = { + "_class_name": "LatentUpsampler", + "quantization_config": {"quant_method": "bitsandbytes"}, + } + + with ( + patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "upsampler_loader._find_safetensors_file", + return_value="/model/spatial_upsampler/model.safetensors", + ), + patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "upsampler_loader._load_explicit_config", + return_value=config, + ), + patch( + "sglang.multimodal_gen.runtime.loader.component_loaders." + "upsampler_loader.safetensors_load_file" + ) as load_weights, + self.assertRaises(ComponentCheckpointUnsupportedError), + ): + UpsamplerLoader().load_customized( + "/model/spatial_upsampler", None, "spatial_upsampler" + ) + + load_weights.assert_not_called() + + +if __name__ == "__main__": + unittest.main()