chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] refactor: rename quantized model path server arg (#19142)

Browse Source
This commit is contained in:
Mick
2026-02-22 23:18:35 +08:00
committed by GitHub
parent c6a99e43b9
commit 45095bac70
7 changed files with 185 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
python/sglang/multimodal_gen/configs/quantization.py
View File

@@ -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(

35
python/sglang/multimodal_gen/docs/quantization.md
View File

@@ -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

8
python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py
View File

@@ -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

160
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
View File

@@ -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(

13
python/sglang/multimodal_gen/runtime/server_args.py
View File

@@ -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):

42
python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
View File

@@ -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

10
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -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(