diff --git a/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py b/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py index 53bf3aaba..be1dfe57a 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py +++ b/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py @@ -28,6 +28,10 @@ class QwenImageArchConfig(DiTArchConfig): default_factory=lambda: { # LoRA mappings r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1", + # SVDquant mappings + r"(.*)\.add_qkv_proj\.(.+)$": r"\1.to_added_qkv.\2", + r"(transformer_blocks\.\d+\.(img_mlp|txt_mlp)\..*\.(smooth_factor_orig|wcscales))$": r"\1", + r".*\.wtscale$": r"", } ) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index 49f0a83e4..40f42c0c7 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -186,6 +186,13 @@ class QwenImagePipelineConfig(ImagePipelineConfig): ] * batch_size txt_seq_lens = [prompt_embeds[0].shape[1]] + if rotary_emb is None: + return { + "img_shapes": img_shapes, + "txt_seq_lens": txt_seq_lens, + "freqs_cis": None, + } + freqs_cis = self.get_freqs_cis( img_shapes, txt_seq_lens, rotary_emb, device, dtype ) @@ -195,6 +202,7 @@ class QwenImagePipelineConfig(ImagePipelineConfig): return { "txt_seq_lens": txt_seq_lens, "freqs_cis": (img_cache, txt_cache), + "img_shapes": img_shapes, } def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): @@ -254,6 +262,14 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig): ], ] * batch_size txt_seq_lens = [prompt_embeds[0].shape[1]] + + if rotary_emb is None: + return { + "img_shapes": img_shapes, + "txt_seq_lens": txt_seq_lens, + "freqs_cis": None, + } + freqs_cis = QwenImagePipelineConfig.get_freqs_cis( img_shapes, txt_seq_lens, rotary_emb, device, dtype ) @@ -271,6 +287,7 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig): return { "txt_seq_lens": txt_seq_lens, "freqs_cis": (img_cache, txt_cache), + "img_shapes": img_shapes, } def preprocess_condition_image( diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index 88a63e901..e2e180a97 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py @@ -272,7 +272,6 @@ class ZImagePipelineConfig(ImagePipelineConfig): device=device, ).flatten(0, 2) - C = self.dit_config.num_channels_latents F = 1 H = height // self.vae_config.arch_config.spatial_compression_ratio W = width // self.vae_config.arch_config.spatial_compression_ratio diff --git a/python/sglang/multimodal_gen/configs/quantization.py b/python/sglang/multimodal_gen/configs/quantization.py new file mode 100644 index 000000000..5d93fe944 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/quantization.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from typing import Any + +import torch + +from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( + is_nunchaku_available, +) +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import StoreBoolean + +logger = init_logger(__name__) + + +@dataclass +class NunchakuSVDQuantArgs: + """CLI-facing configuration for Nunchaku (SVDQuant) inference. + + This is intentionally lightweight and only contains arguments needed to + construct `runtime.layers.quantization.nunchaku_config.NunchakuConfig`. + """ + + enable_svdquant: bool = False + quantized_model_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 not self.enable_svdquant or not self.quantized_model_path: + return + + inferred_precision = None + inferred_rank = None + + filename = os.path.basename(self.quantized_model_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) + + if match: + p_str, r_str = match.groups() + inferred_precision = "nvfp4" if p_str == "fp4" else "int4" + inferred_rank = int(r_str) + + if self.quantization_precision is None: + self.quantization_precision = inferred_precision or "int4" + if inferred_precision: + logger.info( + f"inferred --quantization-precision: {self.quantization_precision} " + f"from --quantized-model-path: {self.quantized_model_path}" + ) + + if self.quantization_rank is None: + self.quantization_rank = inferred_rank or 32 + if inferred_rank: + logger.info( + f"inferred --quantization-rank: {self.quantization_rank} " + f"from --quantized-model-path: {self.quantized_model_path}" + ) + + def validate(self) -> None: + # TODO: warn if the served model doesn't support nunchaku + self._adjust_config() + + if not self.enable_svdquant: + return + + if not current_platform.is_cuda(): + raise ValueError( + "Nunchaku SVDQuant is only supported on NVIDIA CUDA GPUs " + "(Ampere SM8x or SM12x)." + ) + + device_count = torch.cuda.device_count() + + unsupported: list[str] = [] + for i in range(device_count): + major, minor = torch.cuda.get_device_capability(i) + if major == 9: + unsupported.append(f"cuda:{i} (SM{major}{minor}, Hopper)") + elif major not in (8, 12): + unsupported.append(f"cuda:{i} (SM{major}{minor})") + + if unsupported: + raise ValueError( + "Nunchaku SVDQuant is currently only supported on Ampere (SM8x) or SM12x GPUs; " + "Hopper (SM90) is not supported. " + f"Unsupported devices: {', '.join(unsupported)}. " + "Disable it with --enable-svdquant false." + ) + + if not self.quantized_model_path: + raise ValueError( + "--enable-svdquant requires --quantized-model-path to be set" + ) + + if not is_nunchaku_available(): + raise ValueError( + "Nunchaku is enabled, but not installed. Please refer to https://nunchaku.tech/docs/nunchaku/installation/installation.html for detailed installation methods." + ) + + if self.quantization_precision not in ("int4", "nvfp4"): + raise ValueError( + f"Invalid --quantization-precision: {self.quantization_precision}. " + "Must be one of: int4, nvfp4" + ) + + if self.quantization_rank <= 0: + raise ValueError( + f"Invalid --quantization-rank: {self.quantization_rank}. Must be > 0" + ) + + @staticmethod + def add_cli_args(parser) -> None: + parser.add_argument( + "--enable-svdquant", + action=StoreBoolean, + default=NunchakuSVDQuantArgs.enable_svdquant, + help="Enable Nunchaku SVDQuant (W4A4-style) inference.", + ) + parser.add_argument( + "--quantized-model-path", + type=str, + default=NunchakuSVDQuantArgs.quantized_model_path, + help=( + "Path to pre-quantized Nunchaku weights. Can be a single .safetensors " + "file or a directory containing .safetensors." + ), + ) + parser.add_argument( + "--quantization-precision", + type=str, + default=None, + help="Quantization precision: int4 or nvfp4. If not specified, inferred from model path or defaults to int4.", + ) + parser.add_argument( + "--quantization-rank", + type=int, + default=None, + help="SVD low-rank dimension (e.g., 32). If not specified, inferred from model path or defaults to 32.", + ) + parser.add_argument( + "--quantization-act-unsigned", + action=StoreBoolean, + default=NunchakuSVDQuantArgs.quantization_act_unsigned, + help="Use unsigned activation quantization (if supported).", + ) + + @classmethod + def from_dict(cls, kwargs: dict[str, Any]) -> "NunchakuSVDQuantArgs": + # Map CLI/config keys to dataclass fields (keep backwards compatibility). + return cls( + enable_svdquant=bool(kwargs.get("enable_svdquant", cls.enable_svdquant)), + quantized_model_path=kwargs.get( + "quantized_model_path", cls.quantized_model_path + ), + quantization_precision=kwargs.get("quantization_precision"), + quantization_rank=kwargs.get("quantization_rank"), + quantization_act_unsigned=bool( + kwargs.get("quantization_act_unsigned", cls.quantization_act_unsigned) + ), + ) diff --git a/python/sglang/multimodal_gen/docs/quantization.md b/python/sglang/multimodal_gen/docs/quantization.md new file mode 100644 index 000000000..e9c687d35 --- /dev/null +++ b/python/sglang/multimodal_gen/docs/quantization.md @@ -0,0 +1,123 @@ +# Quantization + +This document introduces the model quantization schemes supported in SGLang and how to use them to reduce memory usage and accelerate inference. + +## Nunchaku (SVDQuant) + +### Introduction + +**SVDQuant** is a Post-Training Quantization (PTQ) technique for diffusion models that quantizes model weights and activations to 4-bit precision (W4A4) while maintaining high visual quality. This method uses Singular Value Decomposition (SVD) to decompose the weight matrix into low-rank components and residuals, effectively absorbing outliers in activations, making 4-bit quantization possible. + +**Nunchaku** is a high-performance inference engine that implements SVDQuant, optimized for low-bit neural networks. It is not Quantization-Aware Training (QAT), but directly quantizes pre-trained models. + +Paper: [SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models](https://arxiv.org/abs/2411.05007) (ICLR 2025 Spotlight) + +### Key Features + +SVDQuant significantly reduces memory usage and accelerates inference while maintaining visual quality: + +- **Memory Optimization**: Reduces memory usage by **3.6×** compared to BF16 models. +- **Inference Acceleration**: + - **3.0×** faster than the NF4 (W4A16) baseline on desktop/laptop RTX 4090 GPUs. + - **8.7×** speedup on laptop RTX 4090 by eliminating CPU offloading compared to 16-bit models. + - **3.1×** faster than BF16 and NF4 models on RTX 5090 GPUs with NVFP4. + +### Supported Precisions + +Nunchaku supports two quantization precisions: + +- **INT4**: Standard INT4 quantization, supported on NVIDIA GPUs with Compute Capability 7.0+ (RTX 20 series and above). +- **NVFP4**: FP4 quantization, providing better image quality on newer cards like the RTX 5090. + +### Usage + +#### 1. Install Nunchaku + +```bash +pip install nunchaku +``` + +For more installation information, please refer to the [Nunchaku Official Documentation](https://nunchaku.tech/docs/nunchaku/installation/installation.html). + +#### 2. Download Quantized Models + +Nunchaku provides pre-quantized model weights available on Hugging Face: + +- [nunchaku-ai/nunchaku-qwen-image](https://huggingface.co/nunchaku-ai/nunchaku-qwen-image) +- [nunchaku-ai/nunchaku-flux](https://huggingface.co/nunchaku-ai/nunchaku-flux) + +Taking Qwen-Image as an example, several quantized models with different configurations are provided: + +| Filename | Precision | Rank | Usage | +|----------|-----------|------|-------| +| `svdq-int4_r32-qwen-image.safetensors` | INT4 | 32 | Standard Version | +| `svdq-int4_r128-qwen-image.safetensors` | INT4 | 128 | High-Quality Version | +| `svdq-fp4_r32-qwen-image.safetensors` | NVFP4 | 32 | RTX 5090 Standard Version | +| `svdq-fp4_r128-qwen-image.safetensors` | NVFP4 | 128 | RTX 5090 High-Quality Version | +| `svdq-int4_r32-qwen-image-lightningv1.0-4steps.safetensors` | INT4 | 32 | Lightning 4-Step Version | +| `svdq-int4_r128-qwen-image-lightningv1.1-8steps.safetensors` | INT4 | 128 | Lightning 8-Step Version | + +> **Note**: Higher Rank usually means better image quality, but with slightly increased memory usage and computation. + +#### 3. Run Quantized Models + +SGLang features **smart auto-detection** for Nunchaku models. In most cases, you only need to provide the path to the quantized weights, and the precision and rank will be automatically inferred from the filename. + +**Simplified Command (Recommended):** + +```bash +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 +``` + +**Manual Override (If needed):** + +If your filename doesn't follow the standard naming convention, or you want to force specific settings: + +- `--enable-svdquant`: Manually enable SVDQuant. +- `--quantization-precision`: Set to `int4` or `nvfp4`. +- `--quantization-rank`: Set the SVD rank (e.g., 32, 128). +- `--quantization-act-unsigned` (Optional): Use unsigned activation quantization. + +Example with manual overrides: + +```bash +sglang generate \ + --model-path Qwen/Qwen-Image \ + --prompt "a beautiful sunset" \ + --enable-svdquant \ + --quantized-model-path /path/to/custom_model.safetensors \ + --quantization-precision int4 \ + --quantization-rank 128 +``` + +#### 4. Configuration Recommendations + +Choose the appropriate configuration based on your hardware and requirements: + +| Scenario | Recommended Config | Description | +|----------|-------------------|-------------| +| Standard Use (20/30/40 Series GPU) | INT4 + Rank 32 | Balanced performance and quality | +| Quality Focus (Sufficient VRAM) | INT4 + Rank 128 | Better image quality | +| RTX 5090 Standard Use | NVFP4 + Rank 32 | Utilizes FP4 hardware acceleration | +| RTX 5090 Quality Focus | NVFP4 + Rank 128 | Best image quality | +| Fast Prototyping/Preview | Lightning 4-Step Version | Extremely fast generation, slightly reduced quality | + +### 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. + +2. Auto-Detection Requirements: For auto-detection to work, the filename must contain the pattern `svdq-{precision}_r{rank}` (e.g., `svdq-int4_r32`). + +3. GPU Compatibility: + - INT4: Supports NVIDIA GPUs with Compute Capability 7.0+ (RTX 20 series and above). + - NVFP4: Optimized mainly for newer cards like the RTX 50 series that support FP4. + +4. Lightning Models: When using Lightning versions, adjust `--num-inference-steps` accordingly (usually 4 or 8 steps). + +### Custom Model Quantization + +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. diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index 73a668ff1..f74309a53 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/linear.py @@ -17,7 +17,7 @@ from sglang.multimodal_gen.runtime.distributed import ( tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce, ) -from sglang.multimodal_gen.runtime.layers.quantization.base_config import ( +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, QuantizeMethodBase, ) @@ -165,7 +165,6 @@ class LinearBase(torch.nn.Module): Args: input_size: input dimension of the linear layer. output_size: output dimension of the linear layer. - bias: If true, add bias. skip_bias_add: If true, skip adding bias but instead return it. params_dtype: Data type for the parameters. quant_config: Quantization configure. diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index 0d6c79797..0324606ae 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -2,7 +2,7 @@ from typing import Literal, get_args -from sglang.multimodal_gen.runtime.layers.quantization.base_config import ( +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) @@ -23,17 +23,7 @@ def register_quantization_config(quantization: str): Args: quantization (str): The quantization method name. - Examples: - >>> from sglang.multimodal_gen.runtime.layers.quantization import register_quantization_config - >>> from sglang.multimodal_gen.runtime.layers.quantization import get_quantization_config - >>> from sglang.multimodal_gen.runtime.layers.quantization.base_config import QuantizationConfig - >>> - >>> @register_quantization_config("my_quant") - ... class MyQuantConfig(QuantizationConfig): - ... pass - >>> - >>> get_quantization_config("my_quant") - + """ # noqa: E501 def _wrapper(quant_config_cls): @@ -63,7 +53,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: return method_to_config[quantization] -all = [ +__all__ = [ "QuantizationMethods", "QuantizationConfig", "get_quantization_config", diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/base_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py similarity index 97% rename from python/sglang/multimodal_gen/runtime/layers/quantization/base_config.py rename to python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py index ffb275a8b..f4afe7d01 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/base_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py @@ -65,6 +65,9 @@ def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> class QuantizationConfig(ABC): """Base class for quantization configs.""" + # for quantization frameworks with a separate quantized model provided, e.g. Nunchaku + quantized_model_path: str | None = None + def __init__(self): super().__init__() # mapping is updated by models as they initialize 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 new file mode 100644 index 000000000..c867a3884 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import os +from dataclasses import dataclass +from functools import lru_cache +from typing import Any, Optional + +import torch +from safetensors.torch import load_file as safetensors_load_file +from torch import nn + +from sglang.multimodal_gen.runtime.layers.linear import LinearBase +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +from .base_config import QuantizationConfig, QuantizeMethodBase + +logger = init_logger(__name__) + +SVDQ_W4A4_LAYER_PATTERNS = [ + "attn.to_qkv", + "attn.to_out", + "attn.add_qkv_proj", + "attn.to_add_out", + "img_mlp", + "txt_mlp", +] + +AWQ_W4A16_LAYER_PATTERNS = [ + "img_mod", + "txt_mod", +] + +SKIP_QUANTIZATION_PATTERNS = [ + "norm", + "embed", + "rotary", + "pos_embed", +] + + +@lru_cache(maxsize=1) +def is_nunchaku_available() -> bool: + try: + import nunchaku # noqa + + return True + except Exception: + return False + + +@dataclass +class NunchakuConfig(QuantizationConfig): + """ + Configuration for Nunchaku (SVDQuant) W4A4-style quantization. + + Attributes: + precision: Quantization precision type. Options: + - "int4": Standard INT4 quantization + - "nvfp4": FP4 quantization + 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) + """ + + precision: str = "int4" # "int4" or "nvfp4" + rank: int = 32 + group_size: Optional[int] = None + act_unsigned: bool = False + quantized_model_path: Optional[str] = None + + @classmethod + def get_name(cls) -> str: + return "svdquant" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 70 + + @staticmethod + def get_config_filenames() -> list[str]: + return ["quantization_config.json", "quant_config.json"] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "NunchakuConfig": + + return cls( + precision=config.get("precision", "int4"), + 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"), + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + + if not isinstance(layer, LinearBase): + return None + + for pattern in SKIP_QUANTIZATION_PATTERNS: + if pattern in prefix.lower(): + return None + + for pattern in SVDQ_W4A4_LAYER_PATTERNS: + if pattern in prefix: + from ..nunchaku_linear import NunchakuSVDQLinearMethod + + return NunchakuSVDQLinearMethod( + precision=self.precision, + rank=self.rank, + act_unsigned=self.act_unsigned, + ) + + for pattern in AWQ_W4A16_LAYER_PATTERNS: + if pattern in prefix: + from ..nunchaku_linear import NunchakuAWQLinearMethod + + return NunchakuAWQLinearMethod( + group_size=64, + ) + + from ..nunchaku_linear import NunchakuSVDQLinearMethod + + return NunchakuSVDQLinearMethod( + precision=self.precision, + rank=self.rank, + act_unsigned=self.act_unsigned, + ) + + def __post_init__(self): + if self.group_size is None: + if self.precision == "nvfp4": + self.group_size = 16 + elif self.precision == "int4": + self.group_size = 64 + else: + raise ValueError( + f"Invalid precision: {self.precision}. Must be 'int4' or 'nvfp4'" + ) + + if self.precision not in ["int4", "nvfp4"]: + raise ValueError( + f"Invalid precision: {self.precision}. Must be 'int4' or 'nvfp4'" + ) + + if self.rank <= 0: + raise ValueError(f"Rank must be positive, got {self.rank}") + + @classmethod + def from_dict(cls, config_dict: dict) -> "NunchakuConfig": + """Create configuration from dictionary.""" + return cls(**config_dict) + + def to_dict(self) -> dict: + """Convert configuration to dictionary.""" + return { + "precision": self.precision, + "rank": self.rank, + "group_size": self.group_size, + "act_unsigned": self.act_unsigned, + "quantized_model_path": self.quantized_model_path, + } + + @classmethod + def from_pretrained(cls, model_path: str) -> Optional["NunchakuConfig"]: + for filename in cls.get_config_filenames(): + config_path = os.path.join(model_path, filename) + if os.path.exists(config_path): + with open(config_path, "r") as f: + config_dict = json.load(f) + if config_dict.get("quant_method") == cls.get_name(): + return cls.from_config(config_dict) + return None + + +def _patch_native_svdq_linear( + module: nn.Module, tensor: Any, svdq_linear_cls: type +) -> bool: + if ( + isinstance(module, svdq_linear_cls) + and getattr(module, "wtscale", None) is not None + ): + module.wtscale = tensor + return True + return False + + +def _patch_sglang_svdq_linear( + module: nn.Module, tensor: Any, svdq_method_cls: type +) -> bool: + quant_method = getattr(module, "quant_method", None) + if not isinstance(quant_method, svdq_method_cls): + return False + + existing = getattr(module, "wtscale", None) + if isinstance(existing, nn.Parameter): + with torch.no_grad(): + existing.data.copy_(tensor.to(existing.data.dtype)) + else: + module.wtscale = tensor + + # Keep alpha in sync (kernel reads `layer._nunchaku_alpha`) + try: + module._nunchaku_alpha = float(tensor.detach().cpu().item()) + except Exception: + module._nunchaku_alpha = None + return True + + +def _patch_sglang_svdq_wcscales( + module: nn.Module, tensor: Any, svdq_method_cls: type +) -> bool: + quant_method = getattr(module, "quant_method", None) + if not isinstance(quant_method, svdq_method_cls): + return False + + existing = getattr(module, "wcscales", None) + if isinstance(existing, nn.Parameter): + with torch.no_grad(): + existing.data.copy_(tensor.to(existing.data.dtype)) + else: + module.wcscales = tensor + return True + + +def _patch_nunchaku_scales( + model: nn.Module, + safetensors_list: list[str], +) -> None: + """Patch transformer module with Nunchaku scale tensors from safetensors weights. + + For NVFP4 checkpoints, correctness depends on `wtscale` and attention + `wcscales`. The FSDP loader may skip some of these metadata tensors. + """ + + if not safetensors_list: + return + + if len(safetensors_list) != 1: + logger.warning( + "Nunchaku scale patch expects a single safetensors file, " + "but got %d files. Skipping.", + len(safetensors_list), + ) + return + + from nunchaku.models.linear import SVDQW4A4Linear # type: ignore[import] + + state_dict = safetensors_load_file(safetensors_list[0]) + if state_dict is None: + return + + num_wtscale = 0 + num_wcscales = 0 + + from ..nunchaku_linear import NunchakuSVDQLinearMethod + + for name, module in model.named_modules(): + wt = state_dict.get(f"{name}.wtscale") + if wt is not None: + if _patch_native_svdq_linear(module, wt, SVDQW4A4Linear): + num_wtscale += 1 + elif _patch_sglang_svdq_linear(module, wt, NunchakuSVDQLinearMethod): + num_wtscale += 1 + + wc = state_dict.get(f"{name}.wcscales") + if wc is not None: + # Some modules may have wcscales as a direct attribute/Parameter. + existing = getattr(module, "wcscales", None) + if isinstance(existing, nn.Parameter): + with torch.no_grad(): + existing.data.copy_(wc.to(existing.data.dtype)) + num_wcscales += 1 + elif existing is not None: + setattr(module, "wcscales", wc) + num_wcscales += 1 + elif _patch_sglang_svdq_wcscales(module, wc, NunchakuSVDQLinearMethod): + num_wcscales += 1 + + if num_wtscale > 0: + logger.info("Patched wtscale for %d layers", num_wtscale) + if num_wcscales > 0: + logger.info("Patched wcscales for %d layers", num_wcscales) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/nunchaku_linear.py b/python/sglang/multimodal_gen/runtime/layers/quantization/nunchaku_linear.py new file mode 100644 index 000000000..516f76699 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/nunchaku_linear.py @@ -0,0 +1,291 @@ +# SPDX-License-Identifier: Apache-2.0 +from typing import List, Optional + +import torch +import torch.nn as nn +from torch.nn.parameter import Parameter + +from sglang.multimodal_gen.runtime.layers.linear import LinearMethodBase +from sglang.multimodal_gen.runtime.models.utils import set_weight_attrs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +try: + from nunchaku.ops.gemm import svdq_gemm_w4a4_cuda + from nunchaku.ops.gemv import awq_gemv_w4a16_cuda + from nunchaku.ops.quantize import svdq_quantize_w4a4_act_fuse_lora_cuda +except ImportError: + svdq_gemm_w4a4_cuda = None + awq_gemv_w4a16_cuda = None + svdq_quantize_w4a4_act_fuse_lora_cuda = None + + +class NunchakuSVDQLinearMethod(LinearMethodBase): + def __init__( + self, + precision: str = "int4", + rank: int = 32, + act_unsigned: bool = False, + ): + self.precision = precision + self.rank = rank + self.act_unsigned = act_unsigned + + if precision == "nvfp4": + self.group_size = 16 + else: + self.group_size = 64 + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + + qweight = Parameter( + torch.empty( + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.int8, + ), + requires_grad=False, + ) + set_weight_attrs(qweight, {"input_dim": 1, "output_dim": 0}) + + num_groups = input_size_per_partition // self.group_size + if self.precision == "nvfp4": + scale_dtype = torch.float8_e4m3fn + else: + scale_dtype = params_dtype + wscales = Parameter( + torch.empty(num_groups, output_size_per_partition, dtype=scale_dtype), + requires_grad=False, + ) + + smooth_factor = Parameter( + torch.empty(input_size_per_partition, dtype=params_dtype), + requires_grad=False, + ) + + smooth_factor_orig = Parameter( + torch.empty(input_size_per_partition, dtype=params_dtype), + requires_grad=False, + ) + + proj_down = Parameter( + torch.empty(input_size_per_partition, self.rank, dtype=params_dtype), + requires_grad=False, + ) + proj_up = Parameter( + torch.empty(output_size_per_partition, self.rank, dtype=params_dtype), + requires_grad=False, + ) + + if self.precision == "nvfp4": + wcscales = Parameter( + torch.empty( + output_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + wtscale = Parameter( + torch.empty(1, dtype=params_dtype), + requires_grad=False, + ) + else: + wcscales = None + wtscale = None + + layer.register_parameter("qweight", qweight) + layer.register_parameter("wscales", wscales) + layer.register_parameter("smooth_factor", smooth_factor) + layer.register_parameter("smooth_factor_orig", smooth_factor_orig) + layer.register_parameter("proj_down", proj_down) + layer.register_parameter("proj_up", proj_up) + if wcscales is not None: + layer.register_parameter("wcscales", wcscales) + if wtscale is not None: + layer.register_parameter("wtscale", wtscale) + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.precision = self.precision + layer.rank = self.rank + layer.group_size = self.group_size + layer.act_unsigned = self.act_unsigned + + weight_loader = extra_weight_attrs.get("weight_loader") + if weight_loader is not None: + set_weight_attrs(qweight, {"weight_loader": weight_loader}) + set_weight_attrs(wscales, {"weight_loader": weight_loader}) + set_weight_attrs(smooth_factor, {"weight_loader": weight_loader}) + set_weight_attrs(smooth_factor_orig, {"weight_loader": weight_loader}) + set_weight_attrs(proj_down, {"weight_loader": weight_loader}) + set_weight_attrs(proj_up, {"weight_loader": weight_loader}) + if wcscales is not None: + set_weight_attrs(wcscales, {"weight_loader": weight_loader}) + if wtscale is not None: + set_weight_attrs(wtscale, {"weight_loader": weight_loader}) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + layer.qweight = Parameter(layer.qweight.data, requires_grad=False) + layer.wscales = Parameter(layer.wscales.data, requires_grad=False) + layer.smooth_factor = Parameter(layer.smooth_factor.data, requires_grad=False) + layer.smooth_factor_orig = Parameter( + layer.smooth_factor_orig.data, requires_grad=False + ) + layer.proj_down = Parameter(layer.proj_down.data, requires_grad=False) + layer.proj_up = Parameter(layer.proj_up.data, requires_grad=False) + if hasattr(layer, "wcscales") and layer.wcscales is not None: + layer.wcscales = Parameter(layer.wcscales.data, requires_grad=False) + if hasattr(layer, "wtscale") and layer.wtscale is not None: + layer.wtscale = Parameter(layer.wtscale.data, requires_grad=False) + + alpha: float | None = None + wtscale = getattr(layer, "wtscale", None) + if wtscale is not None: + if isinstance(wtscale, Parameter): + wtscale = wtscale.data + if isinstance(wtscale, torch.Tensor): + alpha = float(wtscale.detach().cpu().item()) + else: + alpha = float(wtscale) + layer._nunchaku_alpha = alpha + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + orig_shape = x.shape + x_2d = x.reshape(-1, orig_shape[-1]) + quantized_x, ascales, lora_act_out = svdq_quantize_w4a4_act_fuse_lora_cuda( + x_2d, + lora_down=layer.proj_down, + smooth=layer.smooth_factor, + fp4=layer.precision == "nvfp4", + pad_size=256, + ) + out_2d = torch.empty( + x_2d.shape[0], + layer.output_size_per_partition, + dtype=x_2d.dtype, + device=x_2d.device, + ) + alpha: float | None = getattr(layer, "_nunchaku_alpha", None) + wcscales = getattr(layer, "wcscales", None) + + svdq_gemm_w4a4_cuda( + act=quantized_x, + wgt=layer.qweight, + out=out_2d, + ascales=ascales, + wscales=layer.wscales, + lora_act_in=lora_act_out, + lora_up=layer.proj_up, + bias=bias, + fp4=layer.precision == "nvfp4", + alpha=alpha, + wcscales=wcscales, + act_unsigned=getattr(layer, "act_unsigned", False), + ) + out = out_2d.reshape(*orig_shape[:-1], layer.output_size_per_partition) + return out + + +class NunchakuAWQLinearMethod(LinearMethodBase): + def __init__(self, group_size: int = 64): + self.group_size = group_size + self.pack_factor = 8 + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + + qweight = Parameter( + torch.empty( + output_size_per_partition // 4, + input_size_per_partition // 2, + dtype=torch.int32, + ), + requires_grad=False, + ) + set_weight_attrs(qweight, {"input_dim": 1, "output_dim": 0}) + + num_groups = input_size_per_partition // self.group_size + wscales = Parameter( + torch.empty(num_groups, output_size_per_partition, dtype=params_dtype), + requires_grad=False, + ) + + wzeros = Parameter( + torch.empty(num_groups, output_size_per_partition, dtype=params_dtype), + requires_grad=False, + ) + + layer.register_parameter("qweight", qweight) + layer.register_parameter("wscales", wscales) + layer.register_parameter("wzeros", wzeros) + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.group_size = self.group_size + layer.pack_factor = self.pack_factor + + weight_loader = extra_weight_attrs.get("weight_loader") + if weight_loader is not None: + set_weight_attrs(qweight, {"weight_loader": weight_loader}) + set_weight_attrs(wscales, {"weight_loader": weight_loader}) + set_weight_attrs(wzeros, {"weight_loader": weight_loader}) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + layer.qweight = Parameter(layer.qweight.data, requires_grad=False) + layer.wscales = Parameter(layer.wscales.data, requires_grad=False) + layer.wzeros = Parameter(layer.wzeros.data, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + orig_shape = x.shape + x_2d = x.reshape(-1, orig_shape[-1]) + + in_features = layer.input_size_per_partition + out_features = layer.output_size_per_partition + out_2d = awq_gemv_w4a16_cuda( + in_feats=x_2d, + kernel=layer.qweight, + scaling_factors=layer.wscales, + zeros=layer.wzeros, + m=x_2d.shape[0], + n=out_features, + k=in_features, + group_size=layer.group_size, + ) + if bias is not None: + view_shape = [1] * (out_2d.ndim - 1) + [-1] + out_2d.add_(bias.view(view_shape)) + + out = out_2d.reshape(*orig_shape[:-1], out_features) + return out diff --git a/python/sglang/multimodal_gen/runtime/layers/vocab_parallel_embedding.py b/python/sglang/multimodal_gen/runtime/layers/vocab_parallel_embedding.py index e32f662ac..fecb4245f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/vocab_parallel_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/vocab_parallel_embedding.py @@ -15,7 +15,7 @@ from sglang.multimodal_gen.runtime.distributed import ( get_tp_group, tensor_model_parallel_all_reduce, ) -from sglang.multimodal_gen.runtime.layers.quantization.base_config import ( +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding, 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 279b0b4ea..5db100304 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 @@ -1,9 +1,16 @@ +import inspect +import json +import logging import os from copy import deepcopy +from typing import Any import torch from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( + _patch_nunchaku_scales, +) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) @@ -16,8 +23,9 @@ 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, + get_metadata_from_safetensors_file, ) -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.logging_utils import get_log_level, init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) @@ -29,6 +37,34 @@ class TransformerLoader(ComponentLoader): component_names = ["transformer", "audio_dit", "video_dit"] expected_library = "diffusers" + def get_list_of_safetensors_to_load( + self, server_args: ServerArgs, component_model_path: str + ) -> list[str]: + """ + get list of safetensors to load. + + For some quantization framework, if --quantized-model-path is provided, load from this path instead of main model + """ + nunchaku_config = server_args.nunchaku_config + + 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] + else: + safetensors_list = _list_safetensors_files(weights_path) + else: + weights_path = component_model_path + safetensors_list = _list_safetensors_files(weights_path) + + if not safetensors_list: + raise ValueError(f"No safetensors files found in {weights_path}") + + return safetensors_list + def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): @@ -57,46 +93,50 @@ class TransformerLoader(ComponentLoader): model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) - # Find all safetensors files - safetensors_list = _list_safetensors_files(component_model_path) - if not safetensors_list: - raise ValueError(f"No safetensors files found in {component_model_path}") + nunchaku_config = server_args.nunchaku_config - # Check if we should use custom initialization weights - custom_weights_path = getattr( - server_args, "init_weights_from_safetensors", None + if nunchaku_config is not None: + # 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 ) - use_custom_weights = False - - if use_custom_weights: - logger.info( - "Using custom initialization weights from: %s", custom_weights_path - ) - assert ( - custom_weights_path is not None - ), "Custom initialization weights must be provided" - if os.path.isdir(custom_weights_path): - safetensors_list = _list_safetensors_files(custom_weights_path) - else: - assert custom_weights_path.endswith( - ".safetensors" - ), "Custom initialization weights must be a safetensors file" - safetensors_list = [custom_weights_path] - - default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] logger.info( - "Loading %s from %s safetensors files, default_dtype: %s", + "Loading %s from %s safetensors files %s, param_dtype: %s", cls_name, len(safetensors_list), - default_dtype, + f": {safetensors_list}" if get_log_level() == logging.DEBUG else "", + param_dtype, ) + init_params: dict[str, Any] = {"config": dit_config, "hf_config": hf_config} + if ( + nunchaku_config is not None + and "quant_config" in inspect.signature(model_cls.__init__).parameters + ): + init_params["quant_config"] = nunchaku_config + # Load the model using FSDP loader - assert server_args.hsdp_shard_dim is not None model = maybe_load_fsdp_model( model_cls=model_cls, - init_params={"config": dit_config, "hf_config": hf_config}, + init_params=init_params, weight_dir_list=safetensors_list, device=get_local_torch_device(), hsdp_replicate_dim=server_args.hsdp_replicate_dim, @@ -105,17 +145,22 @@ class TransformerLoader(ComponentLoader): pin_cpu_memory=server_args.pin_cpu_memory, fsdp_inference=server_args.use_fsdp_inference, # TODO(will): make these configurable - param_dtype=torch.bfloat16, + param_dtype=param_dtype, reduce_dtype=torch.float32, output_dtype=None, strict=False, ) + if nunchaku_config is not None: + _patch_nunchaku_scales(model, safetensors_list) + total_params = sum(p.numel() for p in model.parameters()) logger.info("Loaded model with %.2fB parameters", total_params / 1e9) - assert ( - next(model.parameters()).dtype == default_dtype - ), "Model dtype does not match default dtype" + # considering the existent of mixed-precision models (e.g., nunchaku) + if next(model.parameters()).dtype != param_dtype: + logger.warning( + f"Model dtype does not match expected param dtype, {next(model.parameters()).dtype} vs {param_dtype}" + ) return model diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index af10e8e31..c95be3ae8 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.loader.utils import ( from sglang.multimodal_gen.runtime.loader.weight_utils import ( safetensors_weights_iterator, ) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import set_mixed_precision_policy @@ -40,7 +41,12 @@ def _make_param_like( actual_param: torch.nn.Parameter, tensor: torch.Tensor ) -> torch.nn.Parameter: cls = actual_param.__class__ - new_param = cls.__new__(cls, tensor) + # nn.Parameter defaults to requires_grad=True, which is illegal for non-floating/complex dtypes (e.g., int8/FP8 + # quantized weights). + try: + new_param = cls.__new__(cls, tensor, requires_grad=False) + except TypeError: + new_param = cls.__new__(cls, tensor) new_param.__dict__.update(actual_param.__dict__) new_param.requires_grad = False return new_param @@ -74,26 +80,25 @@ def maybe_load_fsdp_model( """ # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are # manually casting the inputs to the model + default_torch_dtype = param_dtype if param_dtype else torch.bfloat16 mp_policy = MixedPrecisionPolicy( - param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False + default_torch_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False ) set_mixed_precision_policy( - param_dtype=param_dtype, + param_dtype=default_torch_dtype, reduce_dtype=reduce_dtype, output_dtype=output_dtype, mp_policy=mp_policy, ) - with set_default_torch_dtype(param_dtype), torch.device("meta"): + with set_default_torch_dtype(default_torch_dtype), torch.device("meta"): model = model_cls(**init_params) # Check if we should use FSDP use_fsdp = fsdp_inference # Disable FSDP for MPS as it's not compatible - from sglang.multimodal_gen.runtime.platforms import current_platform - if current_platform.is_mps(): use_fsdp = False logger.info("Disabling FSDP for MPS platform as it's not compatible") @@ -191,7 +196,7 @@ def shard_model( # TODO(will): don't reshard after forward for the last layer to save on the # all-gather that will immediately happen Shard the model with FSDP, for n, m in reversed(list(model.named_modules())): - if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]): + if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]): # type: ignore fully_shard(m, **fsdp_kwargs) num_layers_sharded += 1 @@ -209,7 +214,7 @@ def load_model_from_full_model_state_dict( model: FSDPModule | torch.nn.Module, full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None], device: torch.device, - param_dtype: torch.dtype, + param_dtype: torch.dtype | None, strict: bool = False, cpu_offload: bool = False, param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None, @@ -221,7 +226,7 @@ def load_model_from_full_model_state_dict( model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs device (torch.device): device used to move full state dict tensors - param_dtype (torch.dtype): dtype used to move full state dict tensors + param_dtype (torch.dtype): dtype used to move full state dict tensors. If none, respect original dtype from checkpoint strict (bool): flag to check if to load the model in strict mode cpu_offload (bool): flag to check if FSDP offload is enabled param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name @@ -245,6 +250,9 @@ def load_model_from_full_model_state_dict( # sort parameter names to ensure all ranks process parameters in the same order sorted_param_names = sorted(custom_param_sd.keys()) + requires_grad = False + + # shard from loaded state_dict, custom_param_sd -> sharded_sd for target_param_name in sorted_param_names: full_tensor = custom_param_sd[target_param_name] meta_sharded_param = meta_sd.get(target_param_name) @@ -259,8 +267,10 @@ def load_model_from_full_model_state_dict( f"Parameter '{target_param_name}' from checkpoint not found in model; skipping. This is expected for optional parameters." ) continue + + target_dtype = param_dtype if param_dtype else full_tensor.dtype if not hasattr(meta_sharded_param, "device_mesh"): - full_tensor = full_tensor.to(device=device, dtype=param_dtype) + full_tensor = full_tensor.to(device=device, dtype=target_dtype) actual_param = param_dict.get(target_param_name) weight_loader = ( getattr(actual_param, "weight_loader", None) @@ -270,12 +280,20 @@ def load_model_from_full_model_state_dict( if weight_loader is not None: assert actual_param is not None sharded_tensor = torch.empty_like( - meta_sharded_param, device=device, dtype=param_dtype + meta_sharded_param, device=device, dtype=target_dtype ) + # Preserve requires_grad flag to avoid errors with non-floating dtypes + requires_grad = getattr(meta_sharded_param, "requires_grad", False) temp_param = _make_param_like(actual_param, sharded_tensor) + if not ( + sharded_tensor.is_floating_point() or sharded_tensor.is_complex() + ): + requires_grad = False + temp_param.requires_grad = requires_grad weight_loader(temp_param, full_tensor) sharded_tensor = temp_param.data else: + # In cases where parts of the model aren't sharded, some parameters will be plain tensors sharded_tensor = full_tensor # Important: `cpu_offload` is intended for FSDP-managed parameter movement. @@ -290,7 +308,7 @@ def load_model_from_full_model_state_dict( if cpu_offload and not is_fsdp_model: sharded_tensor = sharded_tensor.cpu() else: - full_tensor = full_tensor.to(device=device, dtype=param_dtype) + full_tensor = full_tensor.to(device=device, dtype=target_dtype) sharded_tensor = distribute_tensor( full_tensor, meta_sharded_param.device_mesh, @@ -298,16 +316,27 @@ def load_model_from_full_model_state_dict( ) if cpu_offload: sharded_tensor = sharded_tensor.to("cpu") - sharded_sd[target_param_name] = nn.Parameter(sharded_tensor) + + requires_grad = False + sharded_sd[target_param_name] = nn.Parameter( + sharded_tensor, requires_grad=requires_grad + ) model.reverse_param_names_mapping = reverse_param_names_mapping + # parameters in nn.Module that doesn't exist in safetensor files unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys()) if unused_keys: logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys) - # List of allowed parameter name patterns - ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress"] # Can be extended as needed + # for nunchaku + ALLOWED_NEW_PARAM_PATTERNS = [ + "gate_compress", + "wcscales", + "wtscale", + "bias", + ] for new_param_name in unused_keys: + # check unallowed missing params if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS): logger.error( "Unsupported new parameter: %s. Allowed patterns: %s", @@ -318,17 +347,22 @@ def load_model_from_full_model_state_dict( f"New parameter '{new_param_name}' is not supported. " f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed." ) + meta_sharded_param = meta_sd.get(new_param_name) + + if "wcscales" in new_param_name or "wtscale" in new_param_name: + init_like = torch.ones_like + else: + init_like = torch.zeros_like + if not hasattr(meta_sharded_param, "device_mesh"): - # Initialize with zeros - sharded_tensor = torch.zeros_like( + sharded_tensor = init_like( meta_sharded_param, device=device, dtype=param_dtype ) if cpu_offload and not is_fsdp_model: sharded_tensor = sharded_tensor.cpu() else: - # Initialize with zeros and distribute - full_tensor = torch.zeros_like( + full_tensor = init_like( meta_sharded_param, device=device, dtype=param_dtype ) sharded_tensor = distribute_tensor( diff --git a/python/sglang/multimodal_gen/runtime/loader/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index 1b3cefa65..4be14ef41 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -86,6 +86,8 @@ def hf_to_custom_state_dict( target_param_name, merge_index, num_params_to_merge = param_names_mapping( source_param_name ) + if target_param_name == "" or target_param_name is None: # type: ignore[comparison-overlap] + continue reverse_param_names_mapping[target_param_name] = ( source_param_name, merge_index, diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 4003c4d3d..0b6995ced 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -352,6 +352,7 @@ OOM detected. Possible solutions: 2. Enable SP and/or TP 3. Opt for a sparse-attention backend 4. Enable FSDP by `--use-fsdp-inference` (in a multi-GPU setup) + 5. Enable quantization (e.g. nunchaku) Or, open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose """ diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 7809acaa0..2eb446aac 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -5,10 +5,12 @@ import functools from typing import Any, Dict, List, Optional, Tuple, Union +import diffusers import numpy as np import torch import torch.nn as nn import torch.nn.functional as F +from diffusers.models.attention import FeedForward from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous @@ -26,8 +28,16 @@ from sglang.multimodal_gen.runtime.layers.layernorm import ( ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, + MergedColumnParallelLinear, RowParallelLinear, ) +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, +) +from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( + NunchakuConfig, + is_nunchaku_available, +) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, ) @@ -45,19 +55,36 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name _is_cuda = current_platform.is_cuda() +try: + from nunchaku.models.attention import NunchakuFeedForward # type: ignore[import] +except Exception: + NunchakuFeedForward = None + def _get_qkv_projections( attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None ): - img_query, _ = attn.to_q(hidden_states) - img_key, _ = attn.to_k(hidden_states) - img_value, _ = attn.to_v(hidden_states) + if attn.use_fused_qkv: + img_qkv, _ = attn.to_qkv(hidden_states) + img_query, img_key, img_value = [ + x.contiguous() for x in img_qkv.chunk(3, dim=-1) + ] + else: + img_query, _ = attn.to_q(hidden_states) + img_key, _ = attn.to_k(hidden_states) + img_value, _ = attn.to_v(hidden_states) txt_query = txt_key = txt_value = None if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None: - txt_query, _ = attn.add_q_proj(encoder_hidden_states) - txt_key, _ = attn.add_k_proj(encoder_hidden_states) - txt_value, _ = attn.add_v_proj(encoder_hidden_states) + if attn.use_fused_added_qkv: + txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states) + txt_query, txt_key, txt_value = [ + x.contiguous() for x in txt_qkv.chunk(3, dim=-1) + ] + else: + txt_query, _ = attn.add_q_proj(encoder_hidden_states) + txt_key, _ = attn.add_k_proj(encoder_hidden_states) + txt_value, _ = attn.add_v_proj(encoder_hidden_states) return img_query, img_key, img_value, txt_query, txt_key, txt_value @@ -71,14 +98,27 @@ class GELU(nn.Module): dim_out (`int`): The number of channels in the output. approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation. bias (`bool`, defaults to True): Whether to use a bias in the linear layer. + quant_config: Quantization configure. + prefix: The name of the layer in the state dict. """ def __init__( - self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True + self, + dim_in: int, + dim_out: int, + approximate: str = "none", + bias: bool = True, + quant_config=None, + prefix: str = "", ): super().__init__() self.proj = ColumnParallelLinear( - dim_in, dim_out, bias=bias, gather_output=False + dim_in, + dim_out, + bias=bias, + gather_output=False, + quant_config=quant_config, + prefix=f"{prefix}.proj" if prefix else "", ) self.approximate = approximate @@ -95,9 +135,7 @@ class FeedForward(nn.Module): dim (`int`): The number of channels in the input. dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`. mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - final_dropout (`bool` *optional*, defaults to False): Apply a final dropout. bias (`bool`, defaults to True): Whether to use a bias in the linear layer. """ @@ -194,14 +232,6 @@ class QwenEmbedRope(nn.Module): dim=1, ) - # self.rope = NDRotaryEmbedding( - # rope_dim_list=axes_dim, - # rope_theta=theta, - # use_real=False, - # repeat_interleave_real=False, - # dtype=torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64, - # ) - # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART self.scale_rope = scale_rope @@ -546,6 +576,8 @@ class QwenImageCrossAttention(nn.Module): context_pre_only: bool = False, parallel_attention=False, out_dim: int = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", ) -> None: assert dim % num_heads == 0 super().__init__() @@ -557,38 +589,92 @@ class QwenImageCrossAttention(nn.Module): self.eps = eps self.parallel_attention = parallel_attention self.added_kv_proj_dim = added_kv_proj_dim + self.prefix = prefix + + self.use_fused_qkv = isinstance(quant_config, NunchakuConfig) - # Use separate Q/K/V projections self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads self.inner_kv_dim = self.inner_dim - self.to_q = ColumnParallelLinear( - dim, self.inner_dim, bias=True, gather_output=True - ) - self.to_k = ColumnParallelLinear( - dim, self.inner_dim, bias=True, gather_output=True - ) - self.to_v = ColumnParallelLinear( - dim, self.inner_dim, bias=True, gather_output=True - ) + + if self.use_fused_qkv: + # Use fused QKV projection for nunchaku quantization + self.to_qkv = MergedColumnParallelLinear( + dim, + [self.inner_dim] * 3, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_qkv", + ) + else: + # Use separate Q/K/V projections for non-quantized models + self.to_q = ColumnParallelLinear( + dim, + self.inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_q", + ) + self.to_k = ColumnParallelLinear( + dim, + self.inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_k", + ) + self.to_v = ColumnParallelLinear( + dim, + self.inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_v", + ) if self.qk_norm: self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() if added_kv_proj_dim is not None: - self.add_q_proj = ColumnParallelLinear( - added_kv_proj_dim, self.inner_dim, bias=True, gather_output=True - ) - self.add_k_proj = ColumnParallelLinear( - added_kv_proj_dim, self.inner_dim, bias=True, gather_output=True - ) - self.add_v_proj = ColumnParallelLinear( - added_kv_proj_dim, self.inner_dim, bias=True, gather_output=True - ) + self.use_fused_added_qkv = isinstance(quant_config, NunchakuConfig) + if self.use_fused_added_qkv: + self.to_added_qkv = MergedColumnParallelLinear( + added_kv_proj_dim, + [self.inner_dim] * 3, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.to_added_qkv", + ) + else: + # Use separate Q/K/V projections for non-quantized models + self.add_q_proj = ColumnParallelLinear( + added_kv_proj_dim, + self.inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.add_q_proj", + ) + self.add_k_proj = ColumnParallelLinear( + added_kv_proj_dim, + self.inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.add_k_proj", + ) + self.add_v_proj = ColumnParallelLinear( + added_kv_proj_dim, + self.inner_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.add_v_proj", + ) if context_pre_only is not None and not context_pre_only: self.to_add_out = ColumnParallelLinear( - self.inner_dim, self.dim, bias=out_bias, gather_output=True + self.inner_dim, + self.dim, + bias=out_bias, + gather_output=True, + quant_config=quant_config, + prefix=f"{prefix}.to_add_out", ) else: self.to_add_out = None @@ -597,7 +683,12 @@ class QwenImageCrossAttention(nn.Module): self.to_out = nn.ModuleList([]) self.to_out.append( ColumnParallelLinear( - self.inner_dim, self.dim, bias=out_bias, gather_output=True + self.inner_dim, + self.dim, + bias=out_bias, + gather_output=True, + quant_config=quant_config, + prefix=f"{prefix}.to_out.0", ) ) else: @@ -719,20 +810,29 @@ class QwenImageTransformerBlock(nn.Module): attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6, + quant_config: Optional[QuantizationConfig] | NunchakuConfig = None, + prefix: str = "", zero_cond_t: bool = False, ): super().__init__() + self.prefix = prefix self.dim = dim self.num_attention_heads = num_attention_heads self.attention_head_dim = attention_head_dim + self.quant_config = quant_config self.zero_cond_t = zero_cond_t # Image processing modules self.img_mod = nn.Sequential( nn.SiLU(), ColumnParallelLinear( - dim, 6 * dim, bias=True, gather_output=True + dim, + 6 * dim, + bias=True, + gather_output=True, + quant_config=quant_config, + prefix=f"{prefix}.img_mod", ), # For scale, shift, gate for norm1 and norm2 ) self.img_norm1 = LayerNormScaleShift( @@ -745,19 +845,23 @@ class QwenImageTransformerBlock(nn.Module): added_kv_proj_dim=dim, context_pre_only=False, head_dim=attention_head_dim, + quant_config=quant_config, + prefix=f"{prefix}.attn", ) self.img_norm2 = ScaleResidualLayerNormScaleShift( - hidden_size=dim, eps=eps, elementwise_affine=False - ) - self.img_mlp = FeedForward( - dim=dim, dim_out=dim, activation_fn="gelu-approximate" + dim, eps=eps, elementwise_affine=False ) # Text processing modules self.txt_mod = nn.Sequential( nn.SiLU(), ColumnParallelLinear( - dim, 6 * dim, bias=True, gather_output=True + dim, + 6 * dim, + bias=True, + gather_output=True, + quant_config=quant_config, + prefix=f"{prefix}.txt_mod", ), # For scale, shift, gate for norm1 and norm2 ) self.txt_norm1 = LayerNormScaleShift( @@ -767,12 +871,38 @@ class QwenImageTransformerBlock(nn.Module): self.txt_norm2 = ScaleResidualLayerNormScaleShift( hidden_size=dim, eps=eps, elementwise_affine=False ) - self.txt_mlp = FeedForward( - dim=dim, dim_out=dim, activation_fn="gelu-approximate" - ) # Utils self.fuse_mul_add = MulAdd() + nunchaku_enabled = ( + quant_config is not None + and hasattr(quant_config, "get_name") + and quant_config.get_name() == "svdquant" + and is_nunchaku_available() + ) + ff_class = ( + diffusers.models.attention.FeedForward if nunchaku_enabled else FeedForward + ) + self.img_mlp = ff_class( + dim=dim, + dim_out=dim, + activation_fn="gelu-approximate", + ) + self.txt_mlp = ff_class( + dim=dim, + dim_out=dim, + activation_fn="gelu-approximate", + ) + + if nunchaku_enabled: + nunchaku_kwargs = { + "precision": quant_config.precision, + "rank": quant_config.rank, + "act_unsigned": quant_config.act_unsigned, + } + self.img_mlp = NunchakuFeedForward(self.img_mlp, **nunchaku_kwargs) + self.txt_mlp = NunchakuFeedForward(self.txt_mlp, **nunchaku_kwargs) + def _modulate( self, x: torch.Tensor, @@ -880,6 +1010,24 @@ class QwenImageTransformerBlock(nn.Module): img_mod_params = self.img_mod[1](temb_img_silu)[0] # [B, 6*dim] txt_mod_params = self.txt_mod[1](temb_txt_silu)[0] # [B, 6*dim] + if ( + self.quant_config is not None + and hasattr(self.quant_config, "get_name") + and self.quant_config.get_name() == "svdquant" + ): + # When NOT using nunchaku, reshape mod_params from [B, 6*dim] to [B, dim*6] + # When using nunchaku (svdquant), keep original format + img_mod_params = ( + img_mod_params.view(img_mod_params.shape[0], -1, 6) + .transpose(1, 2) + .reshape(img_mod_params.shape[0], -1) + ) + txt_mod_params = ( + txt_mod_params.view(txt_mod_params.shape[0], -1, 6) + .transpose(1, 2) + .reshape(txt_mod_params.shape[0], -1) + ) + # Split modulation parameters for norm1 and norm2 img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] @@ -923,8 +1071,11 @@ class QwenImageTransformerBlock(nn.Module): gate_x=img_gate1, residual_x=hidden_states, ) - img_mlp_output = self.img_mlp(img_modulated2) - hidden_states = self.fuse_mul_add(img_mlp_output[0], img_gate2, hidden_states) + img_mlp_output = self.img_mlp(img_modulated2)[0] + + if img_mlp_output.dim() == 2: + img_mlp_output = img_mlp_output.unsqueeze(0) + hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states) # Process text stream - norm2 + MLP txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1) @@ -936,9 +1087,12 @@ class QwenImageTransformerBlock(nn.Module): scale=txt_scale2, ) txt_gate2 = txt_gate2_raw.unsqueeze(1) - txt_mlp_output = self.txt_mlp(txt_modulated2) + txt_mlp_output = self.txt_mlp(txt_modulated2)[0] + + if txt_mlp_output.dim() == 2: + txt_mlp_output = txt_mlp_output.unsqueeze(0) encoder_hidden_states = self.fuse_mul_add( - txt_mlp_output[0], txt_gate2, encoder_hidden_states + txt_mlp_output, txt_gate2, encoder_hidden_states ) # Clip to prevent overflow for fp16 @@ -973,6 +1127,7 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin): self, config: QwenImageDitConfig, hf_config: dict[str, Any], + quant_config: Optional[QuantizationConfig] = None, ): super().__init__(config=config, hf_config=hf_config) patch_size = config.arch_config.patch_size @@ -1019,9 +1174,11 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin): dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, + quant_config=quant_config, + prefix=f"transformer_blocks.{layer_idx}", zero_cond_t=self.zero_cond_t, ) - for _ in range(num_layers) + for layer_idx in range(num_layers) ] ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index 701b2a89d..0738c87f6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py @@ -74,8 +74,11 @@ class ImageEncodingStage(PipelineStage): self.move_to_device("cpu") def move_to_device(self, device): + if self.server_args.use_fsdp_inference: + return fields = [ "image_processor", + "image_encoder", ] for field in fields: processor = getattr(self, field, None) diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index cc998e169..edc1da2d9 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -21,6 +21,10 @@ import yaml from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig +from sglang.multimodal_gen.configs.quantization import NunchakuSVDQuantArgs +from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( + NunchakuConfig, +) from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, @@ -309,6 +313,11 @@ class ServerArgs: disable_autocast: bool | None = None + # Quantization / Nunchaku SVDQuant configuration + nunchaku_config: NunchakuSVDQuantArgs | NunchakuConfig | None = field( + default_factory=NunchakuSVDQuantArgs, repr=False + ) + # Master port for distributed inference # TODO: do not hard code master_port: int | None = None @@ -362,6 +371,23 @@ class ServerArgs: """ return self.host is None or self.port is None + def adjust_quant_config(self): + """validate and adjust""" + + # nunchaku + ncfg = self.nunchaku_config + ncfg.validate() + if not ncfg.enable_svdquant or not ncfg.quantized_model_path: + # if nunchaku is not applied + 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, + ) + def adjust_offload(self): if self.pipeline_config.task_type.is_image_gen(): logger.info( @@ -429,6 +455,7 @@ class ServerArgs: configure_logger(server_args=self) self.adjust_offload() + self.adjust_quant_config() if self.attention_backend in ["fa3", "fa4"]: self.attention_backend = "fa" @@ -684,6 +711,9 @@ class ServerArgs: help="Disable autocast for denoising loop and vae decoding in pipeline sampling", ) + # Nunchaku SVDQuant quantization parameters + NunchakuSVDQuantArgs.add_cli_args(parser) + # Master port for distributed inference parser.add_argument( "--master-port", @@ -846,6 +876,9 @@ class ServerArgs: pipeline_config = PipelineConfig.from_kwargs(kwargs) logger.debug(f"Using PipelineConfig: {type(pipeline_config)}") server_args_kwargs["pipeline_config"] = pipeline_config + elif attr == "nunchaku_config": + nunchaku_config = NunchakuSVDQuantArgs.from_dict(kwargs) + server_args_kwargs["nunchaku_config"] = nunchaku_config elif attr in kwargs: server_args_kwargs[attr] = kwargs[attr] @@ -977,6 +1010,7 @@ class ServerArgs: def check_server_args(self) -> None: """Validate inference arguments for consistency""" + # layerwise offload if current_platform.is_mps(): self.use_fsdp_inference = False 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 535860483..6ae87a1c4 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -38,6 +38,7 @@ from huggingface_hub.errors import ( ) from requests.exceptions import ConnectionError as RequestsConnectionError from requests.exceptions import RequestException +from safetensors import safe_open from transformers import AutoConfig, PretrainedConfig from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES @@ -836,3 +837,12 @@ def snapshot_download( } hf_kwargs.update(kwargs) return _hf_snapshot_download(**hf_kwargs) + + +def get_metadata_from_safetensors_file(file_path: str): + try: + with safe_open(file_path, framework="pt", device="cpu") as f: + metadata = f.metadata() + return metadata + except Exception as e: + logger.warning(e) diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index eb73abce4..b73fafabe 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -375,6 +375,12 @@ def configure_logger(server_args, prefix: str = ""): set_uvicorn_logging_configs() +@lru_cache(maxsize=1) +def get_log_level() -> int: + root = logging.getLogger() + return root.level + + def suppress_loggers(loggers_to_suppress: list[str], level: int = logging.WARNING): original_levels = {}