From 45095bac70ef1382425cb86f4b7af66dc6e7641c Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 22 Feb 2026 23:18:35 +0800 Subject: [PATCH] [diffusion] refactor: rename quantized model path server arg (#19142) --- .../multimodal_gen/configs/quantization.py | 37 ++-- .../multimodal_gen/docs/quantization.md | 35 ++-- .../quantization/configs/nunchaku_config.py | 8 +- .../component_loaders/transformer_loader.py | 160 +++++++++++------- .../multimodal_gen/runtime/server_args.py | 13 +- .../runtime/utils/hf_diffusers_utils.py | 42 ++--- .../test/server/testcase_configs.py | 10 ++ 7 files changed, 185 insertions(+), 120 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/quantization.py b/python/sglang/multimodal_gen/configs/quantization.py index 5d93fe944..f0bd7f9c8 100644 --- a/python/sglang/multimodal_gen/configs/quantization.py +++ b/python/sglang/multimodal_gen/configs/quantization.py @@ -28,23 +28,25 @@ class NunchakuSVDQuantArgs: """ enable_svdquant: bool = False - quantized_model_path: str | None = None + transformer_weights_path: str | None = None quantization_precision: str | None = None # "int4" or "nvfp4" quantization_rank: int | None = None quantization_act_unsigned: bool = False def _adjust_config(self) -> None: """infer precision and rank from filename if not provided""" - if self.quantized_model_path and not self.enable_svdquant: - self.enable_svdquant = True + if self.transformer_weights_path and not self.enable_svdquant: + filename = os.path.basename(self.transformer_weights_path) + if re.search(r"svdq-(int4|fp4)_r(\d+)", filename): + self.enable_svdquant = True - if not self.enable_svdquant or not self.quantized_model_path: + if not self.enable_svdquant or not self.transformer_weights_path: return inferred_precision = None inferred_rank = None - filename = os.path.basename(self.quantized_model_path) + filename = os.path.basename(self.transformer_weights_path) # Expected pattern: svdq-{precision}_r{rank}-... # e.g., svdq-int4_r32-qwen-image.safetensors match = re.search(r"svdq-(int4|fp4)_r(\d+)", filename) @@ -59,7 +61,7 @@ class NunchakuSVDQuantArgs: if inferred_precision: logger.info( f"inferred --quantization-precision: {self.quantization_precision} " - f"from --quantized-model-path: {self.quantized_model_path}" + f"from --transformer-weights-path: {self.transformer_weights_path}" ) if self.quantization_rank is None: @@ -67,7 +69,7 @@ class NunchakuSVDQuantArgs: if inferred_rank: logger.info( f"inferred --quantization-rank: {self.quantization_rank} " - f"from --quantized-model-path: {self.quantized_model_path}" + f"from --transformer-weights-path: {self.transformer_weights_path}" ) def validate(self) -> None: @@ -101,9 +103,9 @@ class NunchakuSVDQuantArgs: "Disable it with --enable-svdquant false." ) - if not self.quantized_model_path: + if not self.transformer_weights_path: raise ValueError( - "--enable-svdquant requires --quantized-model-path to be set" + "--enable-svdquant requires --transformer-weights-path to be set" ) if not is_nunchaku_available(): @@ -131,12 +133,12 @@ class NunchakuSVDQuantArgs: help="Enable Nunchaku SVDQuant (W4A4-style) inference.", ) parser.add_argument( - "--quantized-model-path", + "--transformer-weights-path", type=str, - default=NunchakuSVDQuantArgs.quantized_model_path, + default=NunchakuSVDQuantArgs.transformer_weights_path, help=( - "Path to pre-quantized Nunchaku weights. Can be a single .safetensors " - "file or a directory containing .safetensors." + "Path to pre-quantized transformer weights. Can be a single .safetensors " + "file, a directory, or a HuggingFace repo ID. Used by Nunchaku (SVDQuant) and quantized single-file checkpoints." ), ) parser.add_argument( @@ -161,11 +163,14 @@ class NunchakuSVDQuantArgs: @classmethod def from_dict(cls, kwargs: dict[str, Any]) -> "NunchakuSVDQuantArgs": # Map CLI/config keys to dataclass fields (keep backwards compatibility). + path = ( + kwargs.get("transformer_weights_path") + or kwargs.get("transformer_quantized_path") + or kwargs.get("quantized_model_path") + ) return cls( enable_svdquant=bool(kwargs.get("enable_svdquant", cls.enable_svdquant)), - quantized_model_path=kwargs.get( - "quantized_model_path", cls.quantized_model_path - ), + transformer_weights_path=path, quantization_precision=kwargs.get("quantization_precision"), quantization_rank=kwargs.get("quantization_rank"), quantization_act_unsigned=bool( diff --git a/python/sglang/multimodal_gen/docs/quantization.md b/python/sglang/multimodal_gen/docs/quantization.md index ffb6822cc..542cd0292 100644 --- a/python/sglang/multimodal_gen/docs/quantization.md +++ b/python/sglang/multimodal_gen/docs/quantization.md @@ -70,7 +70,7 @@ sglang generate \ --model-path Qwen/Qwen-Image \ --prompt "change the raccoon to a cute cat" \ --save-output \ - --quantized-model-path /path/to/svdq-int4_r32-qwen-image.safetensors + --transformer-weights-path /path/to/svdq-int4_r32-qwen-image.safetensors ``` **Manual Override (If needed):** @@ -89,7 +89,7 @@ sglang generate \ --model-path Qwen/Qwen-Image \ --prompt "a beautiful sunset" \ --enable-svdquant \ - --quantized-model-path /path/to/custom_model.safetensors \ + --transformer-weights-path /path/to/custom_model.safetensors \ --quantization-precision int4 \ --quantization-rank 128 ``` @@ -108,7 +108,7 @@ Choose the appropriate configuration based on your hardware and requirements: ### Notes -1. Model Path Correspondence: `--model-path` should point to the original non-quantized model (for loading config and tokenizer, etc.), while `--quantized-model-path` points to the quantized weight file. +1. Model Path Correspondence: `--model-path` should point to the original non-quantized model (for loading config and tokenizer, etc.), while `--transformer-weights-path` points to the quantized weight file / folder / Huggingface Repo ID. 2. Auto-Detection Requirements: For auto-detection to work, the filename must contain the pattern `svdq-{precision}_r{rank}` (e.g., `svdq-int4_r32`). @@ -122,26 +122,24 @@ Choose the appropriate configuration based on your hardware and requirements: If you want to quantize your own models, you can use the [DeepCompressor](https://github.com/mit-han-lab/deepcompressor) tool. For detailed instructions, please refer to the Nunchaku official documentation. -## FP8 Quantization +## Quantization ### Usage -#### Option 1: Use Pre-quantized Models (Recommended) +#### Option 1: Pre-quantized folder (has `config.json`) -If available, you can directly use pre-quantized FP8 models from Hugging Face or other sources. Simply load them with SGLang: +For quantized checkpoints that include a `config.json` with a `quantization_config` field (e.g., models converted via `convert_hf_to_fp8.py`), where the transformer's `config.json` already encodes the `quantization_config`, use the component override: ```bash sglang generate \ - --model-path /path/to/FLUX.1-dev-FP8/ \ + --model-path /path/to/FLUX.1-dev \ + --transformer-path /path/to/FLUX.1-dev/transformer-FP8 \ --prompt "A Logo With Bold Large Text: SGL Diffusion" \ --save-output ``` -#### Option 2: Convert Your Own Models -If you need to convert a model to FP8 format, use the provided conversion script: - -**Step 1: Convert the Model** +If you need to convert a model to FP8 format yourself, use the provided conversion script: ```bash # convert transformer to FP8 with block quantization @@ -152,13 +150,20 @@ python -m sglang.multimodal_gen.tools.convert_hf_to_fp8 \ --block-size 128 128 ``` -**Step 2: Run Inference** +#### Option 2: Pre-quantized single-file checkpoint (no `config.json`) + + + +Some providers (e.g., [black-forest-labs/FLUX.2-klein-9b-fp8](https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-fp8)) distribute a single `.safetensors` file without a companion `config.json`. Use `--transformer-weights-path` to point to this file (or HuggingFace repo ID) while keeping `--model-path` for the base model: ```bash sglang generate \ - --model-path /path/to/FLUX.1-dev/ - # override transformer component with path to converted model - --transformer-path /path/to/FLUX.1-dev/transformer-FP8 + --model-path black-forest-labs/FLUX.2-klein-9B \ + --transformer-weights-path black-forest-labs/FLUX.2-klein-9b-fp8 \ --prompt "A Logo With Bold Large Text: SGL Diffusion" \ --save-output ``` + +SGLang-Diffusion will automatically read the `quantization_config` metadata embedded in the safetensors file header (if present). For the quant config to be auto-detected, the file's metadata must contain a JSON-encoded `quantization_config` key with at least a `quant_method` field (e.g. `"fp8"`). + +Note: this feature is a WIP diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py index 839692a84..a8c22b407 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py @@ -40,7 +40,7 @@ class NunchakuConfig(QuantizationConfig): rank: SVD low-rank dimension for absorbing outliers group_size: Quantization group size (automatically set based on precision) act_unsigned: Use unsigned activation quantization - quantized_model_path: Path to pre-quantized model weights (.safetensors) + transformer_weights_path: Path to pre-quantized transformer weights (.safetensors) model_cls: DiT model class that provides quantization rules via get_nunchaku_quant_rules() """ @@ -48,7 +48,7 @@ class NunchakuConfig(QuantizationConfig): rank: int = 32 group_size: Optional[int] = None act_unsigned: bool = False - quantized_model_path: Optional[str] = None + transformer_weights_path: Optional[str] = None model_cls: Optional[type] = None @classmethod @@ -75,7 +75,7 @@ class NunchakuConfig(QuantizationConfig): rank=int(config.get("rank", 32)), group_size=config.get("group_size"), act_unsigned=bool(config.get("act_unsigned", False)), - quantized_model_path=config.get("quantized_model_path"), + transformer_weights_path=config.get("transformer_weights_path"), ) def get_quant_method( @@ -158,7 +158,7 @@ class NunchakuConfig(QuantizationConfig): "rank": self.rank, "group_size": self.group_size, "act_unsigned": self.act_unsigned, - "quantized_model_path": self.quantized_model_path, + "transformer_weights_path": self.transformer_weights_path, } @classmethod diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 20078b0f4..344fc6d38 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -2,13 +2,13 @@ import inspect import json import logging import os -from copy import deepcopy -from typing import Any +from typing import Any, Dict, List, Optional import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( + NunchakuConfig, _patch_nunchaku_scales, ) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( @@ -25,6 +25,8 @@ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( get_diffusers_component_config, get_metadata_from_safetensors_file, get_quant_config, + get_quant_config_from_safetensors_metadata, + maybe_download_model, ) from sglang.multimodal_gen.runtime.utils.logging_utils import get_log_level, init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE @@ -44,35 +46,105 @@ class TransformerLoader(ComponentLoader): """ get list of safetensors to load. - For some quantization framework, if --quantized-model-path is provided, load from this path instead of main model + If --transformer-weights-path is provided, load weights from that path + instead of the base model's component directory. """ - nunchaku_config = server_args.nunchaku_config + quantized_path = server_args.transformer_weights_path - if nunchaku_config is not None and nunchaku_config.quantized_model_path: - # load from quantized_model_path if applicable - weights_path = nunchaku_config.quantized_model_path - - logger.info("Using quantized model weights from: %s", weights_path) - if os.path.isfile(weights_path) and weights_path.endswith(".safetensors"): - safetensors_list = [weights_path] + if quantized_path: + quantized_path = maybe_download_model(quantized_path) + logger.info("using quantized transformer weights from: %s", quantized_path) + if os.path.isfile(quantized_path) and quantized_path.endswith( + ".safetensors" + ): + safetensors_list = [quantized_path] else: - safetensors_list = _list_safetensors_files(weights_path) + safetensors_list = _list_safetensors_files(quantized_path) else: - weights_path = component_model_path - safetensors_list = _list_safetensors_files(weights_path) + safetensors_list = _list_safetensors_files(component_model_path) if not safetensors_list: - raise ValueError(f"No safetensors files found in {weights_path}") + raise ValueError( + f"no safetensors files found in " + f"{quantized_path or component_model_path}" + ) return safetensors_list + def _resolve_quant_config( + self, + hf_config: Dict[str, List[str]], + server_args: ServerArgs, + safetensors_list: list[str], + ) -> Optional[dict]: + # priority: model config.json → safetensors metadata → nunchaku config + quant_config = get_quant_config(hf_config) + if quant_config is None and server_args.transformer_weights_path: + # try to read quantization_config from the safetensors metadata header + for safetensors_file in safetensors_list: + quant_config = get_quant_config_from_safetensors_metadata( + safetensors_file + ) + if quant_config: + break + return quant_config + + def _resolve_target_param_dtype( + self, + quant_config: Optional[dict], + nunchaku_config: Optional[NunchakuConfig], + model_cls, + server_args: ServerArgs, + ) -> Optional[torch.dtype]: + if quant_config is not None or nunchaku_config is not None: + # TODO: improve the condition + # respect dtype from checkpoint + param_dtype = None + else: + param_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] + + if nunchaku_config is not None: + nunchaku_config.model_cls = model_cls + # verify that the nunchaku checkpoint matches the selected model class + original_dit_cls_name = json.loads( + get_metadata_from_safetensors_file( + nunchaku_config.transformer_weights_path + ).get("config") + )["_class_name"] + specified_dit_cls_name = str(model_cls.__name__) + if original_dit_cls_name != specified_dit_cls_name: + raise Exception( + f"Class name of DiT specified in nunchaku transformer_weights_path: {original_dit_cls_name} does not match that of specified DiT name: {specified_dit_cls_name}" + ) + + return param_dtype + def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): """Load the transformer based on the model path, and inference args.""" + # 1. hf config config = get_diffusers_component_config(component_path=component_model_path) - hf_config = deepcopy(config) + # 2. quant config + safetensors_list = self.get_list_of_safetensors_to_load( + server_args, component_model_path + ) + quant_config = self._resolve_quant_config(config, server_args, safetensors_list) + + # 3. dit config + # Config from Diffusers supersedes sgl_diffusion's model config + component_name = _normalize_component_type(component_name) + server_args.model_paths[component_name] = component_model_path + if component_name in ("transformer", "video_dit"): + pipeline_dit_config_attr = "dit_config" + elif component_name in ("audio_dit",): + pipeline_dit_config_attr = "audio_dit_config" + else: + raise ValueError(f"Invalid module name: {component_name}") + dit_config = getattr(server_args.pipeline_config, pipeline_dit_config_attr) + dit_config.update_model_arch(config) + cls_name = config.pop("_class_name") if cls_name is None: raise ValueError( @@ -80,46 +152,11 @@ class TransformerLoader(ComponentLoader): "Only diffusers format is supported." ) - component_name = _normalize_component_type(component_name) - server_args.model_paths[component_name] = component_model_path - - if component_name in ("transformer", "video_dit"): - pipeline_dit_config_attr = "dit_config" - elif component_name in ("audio_dit",): - pipeline_dit_config_attr = "audio_dit_config" - else: - raise ValueError(f"Invalid module name: {component_name}") - # Config from Diffusers supersedes sgl_diffusion's model config - dit_config = getattr(server_args.pipeline_config, pipeline_dit_config_attr) - dit_config.update_model_arch(config) - model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) nunchaku_config = server_args.nunchaku_config - - if nunchaku_config is not None: - nunchaku_config.model_cls = model_cls - - # respect dtype from checkpoint - # TODO: improve the condition - param_dtype = None - - # check if the specified nunchaku quantized model path matches with the specified model path - original_dit_cls_name = json.loads( - get_metadata_from_safetensors_file( - nunchaku_config.quantized_model_path - ).get("config") - )["_class_name"] - specified_dit_cls_name = str(model_cls.__name__) - if original_dit_cls_name != specified_dit_cls_name: - raise Exception( - f"Class name of DiT specified in nunchaku quantized model_path: {original_dit_cls_name} does not match that of specified DiT name: {specified_dit_cls_name}" - ) - else: - param_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] - - safetensors_list = self.get_list_of_safetensors_to_load( - server_args, component_model_path + param_dtype = self._resolve_target_param_dtype( + quant_config, nunchaku_config, model_cls, server_args ) logger.info( @@ -130,13 +167,20 @@ class TransformerLoader(ComponentLoader): param_dtype, ) - init_params: dict[str, Any] = {"config": dit_config, "hf_config": hf_config} - + init_params: dict[str, Any] = {"config": dit_config, "hf_config": config} + # prepare init_param if "quant_config" in inspect.signature(model_cls.__init__).parameters: - quant_config = get_quant_config(config) - init_params["quant_config"] = ( - quant_config if quant_config else nunchaku_config + init_params.update( + { + "quant_config": (quant_config if quant_config else nunchaku_config), + } ) + if init_params["quant_config"] is None: + logger.warning( + f"transformer_weights_path provided, but quantization config not resolved, which is unexpected and likely to cause errors" + ) + else: + logger.debug("quantization config: %s", init_params["quant_config"]) # Load the model using FSDP loader model = maybe_load_fsdp_model( diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 51ddfc479..2539b87c3 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -289,6 +289,9 @@ class ServerArgs: # Component path overrides (key = model_index.json component name, value = path) component_paths: dict[str, str] = field(default_factory=dict) + + # path to pre-quantized transformer weights (single .safetensors or directory). + transformer_weights_path: str | None = None # can restrict layers to adapt, e.g. ["q_proj"] # Will adapt only q, k, v, o by default. lora_target_modules: list[str] | None = None @@ -400,15 +403,19 @@ class ServerArgs: if ncfg is None or isinstance(ncfg, NunchakuConfig): return ncfg.validate() - if not ncfg.enable_svdquant or not ncfg.quantized_model_path: - # if nunchaku is not applied + + # propagate the path to server_args + if ncfg.transformer_weights_path: + self.transformer_weights_path = ncfg.transformer_weights_path + + if not ncfg.enable_svdquant or not ncfg.transformer_weights_path: self.nunchaku_config = None else: self.nunchaku_config = NunchakuConfig( precision=self.nunchaku_config.quantization_precision, rank=self.nunchaku_config.quantization_rank, act_unsigned=self.nunchaku_config.quantization_act_unsigned, - quantized_model_path=self.nunchaku_config.quantized_model_path, + transformer_weights_path=self.nunchaku_config.transformer_weights_path, ) def _adjust_offload(self): diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index f47e5af63..4f0ba6c5a 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -913,48 +913,42 @@ def get_quant_config_from_safetensors_metadata( file_path: str, ) -> Optional[QuantizationConfig]: """Extract quantization config from a safetensors file's metadata header. - - Safetensors files can embed a flat string→string metadata dict in their header. - We expect a ``quantization_config`` key containing a JSON-encoded dict with at - least a ``quant_method`` field (e.g. ``"fp8"``), matching the format written by - ``convert_hf_to_fp8.py`` when embedded into a config.json. - Returns None if no recognizable quantization metadata is found. """ metadata = get_metadata_from_safetensors_file(file_path) if not metadata: return None - quant_config_str = metadata.get("quantization_config") + quant_config_str = metadata.get("_quantization_metadata") if not quant_config_str: return None - try: quant_config_dict = json.loads(quant_config_str) - except Exception as e: - logger.warning( - "failed to parse quantization_config from safetensors metadata: %s", e - ) + except Exception as _e: return None + # handle diffusers fp8 safetensors metadata format + if ( + "quant_method" not in quant_config_dict + and "format_version" in quant_config_dict + and "layers" in quant_config_dict + ): + layers = quant_config_dict.get("layers", {}) + if any( + isinstance(v, dict) and "float8" in v.get("format", "") + for v in layers.values() + ): + quant_config_dict["quant_method"] = "fp8" + quant_config_dict["activation_scheme"] = "dynamic" + quant_method = quant_config_dict.get("quant_method") if not quant_method: - logger.warning( - "quantization_config in safetensors metadata is missing 'quant_method'" - ) return None try: quant_cls = get_quantization_config(quant_method) config = quant_cls.from_config(quant_config_dict) - logger.info( - "loaded quantization config (%s) from safetensors metadata: %s", - quant_method, - file_path, - ) + logger.debug(f"Get quantization config from safetensors file: {file_path}") return config - except Exception as e: - logger.warning( - "failed to build QuantizationConfig from safetensors metadata: %s", e - ) + except Exception as _e: return None diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 4cec678d3..7bf5baad3 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -368,6 +368,16 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [ ), T2I_sampling_params, ), + # TODO: modeling of flux different from official flux, so weights can't be loaded + # consider opting for a different quantized hf-repo + # DiffusionTestCase( + # "flux_image_t2i_override_transformer_weights_path_fp8", + # DiffusionServerArgs( + # model_path="black-forest-labs/FLUX.1-dev", modality="image", + # extras=["--transformer-weights-path black-forest-labs/FLUX.1-dev-FP8"] + # ), + # T2I_sampling_params, + # ), DiffusionTestCase( "flux_2_image_t2i", DiffusionServerArgs(