diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index 0d891df01..6cfa5cbcd 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -144,6 +144,14 @@ backend. None Serialized config stays quant_method=modelopt with quant_algo=FP8; dit_layerwise_offload is supported and dit_cpu_offload stays disabled + + auto-round W4A16 + Transformer component repo with a self-describing quantization_config and auto_round:auto_gptq packing + --transformer-path + Native dense DiTs with compatible component parameter mappings; MiniMax-H3 Diffusers components are supported + None + Auto-detected; reuses the SRT GPTQ/Marlin backend. No --quantization flag is needed; use TP/sequence parallelism rather than FSDP. + modelopt-nvfp4 Mixed transformer directory/repo with config.json, raw NVFP4 safetensors export/repo, or full ModelOpt Diffusers repo diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index da90353dd..7dd00c526 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -2,6 +2,9 @@ from typing import Literal, get_args +from sglang.multimodal_gen.runtime.layers.quantization.auto_round import ( + AutoRoundConfig, +) from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import ( BitsAndBytesConfig, ) @@ -27,6 +30,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.mxfp4_npu import ( from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config QuantizationMethods = Literal[ + "auto-round", "fp8", "modelopt", "modelopt_fp8", @@ -43,6 +47,7 @@ QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) # The customized quantization methods which will be added to this dict. _CUSTOMIZED_METHOD_TO_QUANT_CONFIG = { + "auto-round": AutoRoundConfig, "modelopt": ModelOptFp8DiffusionConfig, "modelopt_fp8": ModelOptFp8Config, "modelopt_fp4": ModelOptFp4Config, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/auto_round.py b/python/sglang/multimodal_gen/runtime/layers/quantization/auto_round.py new file mode 100644 index 000000000..b1f3cc6a5 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/auto_round.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sglang.multimodal_gen.runtime.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, +) +from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping +from sglang.srt.layers.quantization.auto_round import AutoRoundConfig as SRTConfig + + +class AutoRoundConfig(QuantizationConfig): + """Use SRT's serialized AutoRound kernels with diffusion linear layers.""" + + checkpoint_uses_native_qkv_layout = True + + def __init__(self, srt_config: SRTConfig) -> None: + super().__init__() + self.srt_config = srt_config + + @classmethod + def get_name(cls) -> str: + return "auto-round" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return SRTConfig.get_supported_act_dtypes() + + @classmethod + def get_min_capability(cls) -> int: + return SRTConfig.get_min_capability() + + @classmethod + def get_config_filenames(cls) -> list[str]: + return SRTConfig.get_config_filenames() + + @classmethod + def from_config(cls, config: dict) -> "AutoRoundConfig": + srt_config = SRTConfig.from_config(config) + if "gptq" not in srt_config.packing_format: + raise ValueError( + "SGLang diffusion currently supports AutoRound auto_gptq " + f"checkpoints, but got {srt_config.packing_format!r}." + ) + return cls(srt_config) + + def remap_checkpoint_prefixes(self, param_names_mapping: dict) -> None: + mapping = get_param_names_mapping(param_names_mapping) + remapped: dict[str, dict] = {} + for prefix, layer_config in (self.srt_config.extra_config or {}).items(): + target, _, _ = mapping(f"{prefix}.weight") + target = target.removesuffix(".weight") + previous = remapped.setdefault(target, layer_config) + if previous != layer_config: + raise ValueError( + f"AutoRound fused module {target!r} has inconsistent shard configs." + ) + + self.srt_config.extra_config = remapped + self.srt_config.block_name_to_quantize = None + self.srt_config.packed_modules_mapping = self.packed_modules_mapping + + def get_quant_method(self, layer: torch.nn.Module, prefix: str): + if not isinstance(layer, LinearBase): + return None + + weight_bits, _, _ = self.srt_config.get_layer_config(layer, prefix) + if not self.srt_config.check_quantized(weight_bits): + return UnquantizedLinearMethod() + + return self.srt_config.apply_gptq_quant_layer( + layer, + prefix, + self.srt_config.backend, + additional_linear_types=(LinearBase,), + ) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py index ffe08c181..ebb68a2e4 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py @@ -45,3 +45,7 @@ class QuantizationConfig(SRTQuantizationConfig): ) -> bool: """Whether a row-parallel shard preserves this format's input layout.""" return True + + def remap_checkpoint_prefixes(self, param_names_mapping: dict) -> None: + """Translate checkpoint module names to the native model namespace.""" + return 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 96e9e5b3f..46ebed128 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 @@ -290,6 +290,15 @@ class TransformerLoader(ComponentLoader): "Comfy quantized checkpoints do not support FSDP " "inference; use TP and/or sequence parallelism instead" ) + if ( + use_fsdp + and quant_spec.quant_config is not None + and quant_spec.quant_config.get_name() == "auto-round" + ): + raise ValueError( + "AutoRound checkpoints do not support diffusion FSDP inference; " + "use TP and/or sequence parallelism instead" + ) if quant_spec.gguf_file is not None: logger.info( diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index b5b7e0cb1..aec5354ff 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -789,6 +789,9 @@ def resolve_transformer_quant_load_spec( packed = getattr(model_cls, "packed_modules_mapping", None) if packed and hasattr(quant_config, "packed_modules_mapping"): quant_config.packed_modules_mapping = packed + quant_config.remap_checkpoint_prefixes( + vars(model_cls).get("param_names_mapping", {}) + ) nunchaku_config = server_args.nunchaku_config @@ -865,7 +868,7 @@ def _needs_device_weight_postprocess( ) -> bool: """Return whether post-load weight processing needs CUDA/NPU tensors.""" quant_name = _get_quant_config_name(quant_config) - if quant_name in ("modelopt_fp8", "comfy_fp8", "mxfp8"): + if quant_name in ("modelopt_fp8", "comfy_fp8", "auto-round", "mxfp8"): return True if quant_name == "kitchen_int8": assert isinstance(quant_config, KitchenInt8Config) diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 74cde2416..8218e9501 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -50,6 +50,9 @@ from sglang.multimodal_gen.runtime.layers.linear import ( ReplicatedLinear, UnquantizedLinearMethod, ) +from sglang.multimodal_gen.runtime.layers.quantization.auto_round import ( + AutoRoundConfig, +) from sglang.multimodal_gen.runtime.layers.quantization.comfy_fp8 import ( ComfyFp8Config, ComfyFullPrecisionFp8LinearMethod, @@ -95,6 +98,7 @@ from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( ) from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan from sglang.multimodal_gen.runtime.models.dits.flux import FluxSingleTransformerBlock +from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import MiniMaxH3DiTModel from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability from sglang.multimodal_gen.runtime.utils.quantization_utils import ( @@ -138,6 +142,41 @@ def _make_quant_config(name: str, **attrs): class TestTransformerQuantHelpers(unittest.TestCase): + def test_autoround_config_is_inferred_and_remapped_to_native_prefixes(self): + layer_config = { + "bits": 4, + "group_size": 128, + "sym": True, + "data_type": "int", + "act_bits": 16, + } + metadata = { + "quant_method": "auto-round", + "packing_format": "auto_round:auto_gptq", + **layer_config, + "block_name_to_quantize": "transformer_blocks", + "extra_config": { + "context_embedder": {**layer_config, "bits": 16}, + **{ + f"transformer_blocks.0.attn.to_{shard}": layer_config + for shard in ("q", "k", "v") + }, + }, + } + + config = get_quant_config( + {"quantization_config": metadata}, "/unused/component/path" + ) + self.assertIsInstance(config, AutoRoundConfig) + config.remap_checkpoint_prefixes(MiniMaxH3DiTModel.param_names_mapping) + self.assertEqual( + config.srt_config.get_layer_config(object(), "condition_proj")[0], 16 + ) + self.assertEqual( + config.srt_config.get_layer_config(object(), "blocks.0.attn.qkv_proj")[0], + 4, + ) + def test_mps_layerwise_load_uses_residency_api(self): server_args = SimpleNamespace( should_configure_layerwise_offload_for_lazy_component=lambda name: ( diff --git a/python/sglang/srt/layers/quantization/auto_round.py b/python/sglang/srt/layers/quantization/auto_round.py index 0924a0041..53caa88ad 100644 --- a/python/sglang/srt/layers/quantization/auto_round.py +++ b/python/sglang/srt/layers/quantization/auto_round.py @@ -434,7 +434,13 @@ class AutoRoundConfig(QuantizationConfig): return AWQLinearMethod(quant_args) return None - def apply_gptq_quant_layer(self, layer, prefix: str, backend: str = "auto"): + def apply_gptq_quant_layer( + self, + layer, + prefix: str, + backend: str = "auto", + additional_linear_types: tuple[type[torch.nn.Module], ...] = (), + ): from sglang.srt.layers.linear import LinearBase from sglang.srt.layers.moe.fused_moe_triton import FusedMoE from sglang.srt.layers.quantization.gptq import ( @@ -449,9 +455,12 @@ class AutoRoundConfig(QuantizationConfig): from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + linear_types = (LinearBase, ParallelLMHead, *additional_linear_types) + is_linear = isinstance(layer, linear_types) + weight_bits, group_size, sym = self.get_layer_config(layer, prefix) if not self.check_quantized(weight_bits): - if isinstance(layer, (LinearBase, ParallelLMHead)): + if is_linear: return UnquantizedLinearMethod() else: return None @@ -475,7 +484,7 @@ class AutoRoundConfig(QuantizationConfig): layer.scheme = quant_args.get_moe_scheme(layer) return GPTQMoEMethod(quant_args) - if isinstance(layer, (LinearBase, ParallelLMHead)): + if is_linear: layer.scheme = quant_args.get_linear_scheme(layer) return GPTQLinearMethod(quant_args) @@ -494,7 +503,7 @@ class AutoRoundConfig(QuantizationConfig): layer.scheme = quant_args.get_moe_scheme(layer) return GPTQMoEMethod(quant_args) - if isinstance(layer, (LinearBase, ParallelLMHead)): + if is_linear: layer.scheme = quant_args.get_linear_scheme(layer) return GPTQLinearMethod(quant_args) @@ -551,7 +560,7 @@ class AutoRoundConfig(QuantizationConfig): } return MoeWNA16Config.from_config(config).get_quant_method(layer, prefix) - if isinstance(layer, (LinearBase, ParallelLMHead)): + if is_linear: if use_marlin: return GPTQMarlinLinearMethod(quant_args_marlin) else: