diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index 2cae2aa4b..c4cdb687c 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -6,18 +6,24 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor QuantizationConfig, ) from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_fp8 import ( + ModelOptFp8Config as ModelOptFp8DiffusionConfig, +) from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( ModelOptFp4Config, ModelOptFp8Config, ) from sglang.multimodal_gen.runtime.layers.quantization.modelslim import ModelSlimConfig -QuantizationMethods = Literal["fp8", "modelopt_fp8", "modelopt_fp4", "modelslim"] +QuantizationMethods = Literal[ + "fp8", "modelopt", "modelopt_fp8", "modelopt_fp4", "modelslim" +] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) # The customized quantization methods which will be added to this dict. _CUSTOMIZED_METHOD_TO_QUANT_CONFIG = { + "modelopt": ModelOptFp8DiffusionConfig, "modelopt_fp8": ModelOptFp8Config, "modelopt_fp4": ModelOptFp4Config, "modelslim": ModelSlimConfig, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py new file mode 100644 index 000000000..7a4bcd6e5 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py @@ -0,0 +1,204 @@ +"""ModelOpt FP8 quantization support for diffusion models. + +Handles checkpoints produced by NVIDIA Model Optimizer (ModelOpt) with +``quant_algo: "FP8"`` and ``quant_method: "modelopt"``. + +Per quantized linear layer the checkpoint contains: + .weight float8_e4m3fn [out, in] FP8 quantized weight + .weight_scale float32 scalar per-tensor weight scale + .input_scale float32 scalar per-tensor static activation scale + .bias bfloat16 [out] bias (unquantized) + ._amax (ignored) calibration artifact + +Layers listed in the ``ignore`` field of the quantization config remain in +bfloat16 and use the standard unquantized linear method. +""" + +from __future__ import annotations + +import fnmatch +import logging +from typing import Any, Dict, List, Optional + +import torch + +from sglang.multimodal_gen.runtime.layers.linear import ( + LinearMethodBase, + UnquantizedLinearMethod, +) +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.multimodal_gen.runtime.models.parameter import ( + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + cutlass_fp8_supported, +) +from sglang.srt.layers.quantization.utils import convert_to_channelwise + +logger = logging.getLogger(__name__) + + +class ModelOptFp8Config(QuantizationConfig): + """Config for ModelOpt static per-tensor FP8 quantization.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = True, + ignore: Optional[List[str]] = None, + ) -> None: + super().__init__() + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + self.ignore = ignore or [] + + # -- QuantizationConfig interface ---------------------------------------- + + @classmethod + def get_name(cls) -> str: + return "modelopt" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 89 + + @staticmethod + def get_config_filenames() -> list[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "ModelOptFp8Config": + quant_algo = config.get("quant_algo") + if quant_algo is None: + raise ValueError( + "ModelOptFp8Config requires 'quant_algo' in the quantization config." + ) + if "FP8" not in quant_algo: + raise ValueError( + f"ModelOptFp8Config only supports FP8, got quant_algo={quant_algo!r}." + ) + ignore = config.get("ignore", []) + return cls(is_checkpoint_fp8_serialized=True, ignore=ignore) + + def _is_layer_ignored(self, prefix: str) -> bool: + """Check whether *prefix* matches any pattern in the ignore list. + + ModelOpt ignore patterns are matched against the full prefix as a glob + (e.g. ``"norm_out*"`` matches ``"norm_out.linear"``) **and** against the + first path component (e.g. ``"proj_out"`` matches only the top-level + ``proj_out``, not ``single_transformer_blocks.0.proj_out``). + """ + first_component = prefix.split(".")[0] + for pattern in self.ignore: + if fnmatch.fnmatch(prefix, pattern): + return True + if fnmatch.fnmatch(first_component, pattern): + return True + return False + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.multimodal_gen.runtime.layers.linear import LinearBase + + if isinstance(layer, LinearBase): + if self._is_layer_ignored(prefix): + return UnquantizedLinearMethod() + return ModelOptFp8LinearMethod(self) + return None + + def get_scaled_act_names(self) -> list[str]: + return [] + + +class ModelOptFp8LinearMethod(LinearMethodBase): + """Linear method for ModelOpt static per-tensor FP8 quantization. + + Uses ``torch._scaled_mm`` (or CUTLASS FP8 GEMM when available) for + the FP8 matrix multiply — the same kernels used by the LLM runtime. + """ + + def __init__(self, quant_config: ModelOptFp8Config): + self.quant_config = quant_config + self.cutlass_fp8_supported = cutlass_fp8_supported() + + 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) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=torch.float8_e4m3fn, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + for scale_name in ("weight_scale", "input_scale"): + scale = PerTensorScaleParameter( + data=torch.full( + (len(output_partition_sizes),), + torch.finfo(torch.float32).min, + dtype=torch.float32, + ), + weight_loader=weight_loader, + ) + layer.register_parameter(scale_name, scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Diffusion models use single-partition layers (no TP, no fused QKV), + # so we just take the max scale directly without the + # dequantize-requantize round-trip that the LLM path does (which + # requires CUDA kernels that are unavailable during CPU-phase loading). + max_w_scale = layer.weight_scale.max() + + # Transpose weight to [in, out] column-major layout for + # apply_fp8_linear / CUTLASS fp8_scaled_mm. Do NOT call + # .contiguous() — the kernel requires column-major stride. + layer.weight = torch.nn.Parameter(layer.weight.data.t(), requires_grad=False) + + if self.cutlass_fp8_supported: + max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths) + layer.weight_scale = torch.nn.Parameter(max_w_scale, requires_grad=False) + layer.input_scale = torch.nn.Parameter( + layer.input_scale.max(), requires_grad=False + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_fp8_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + ) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 385931b19..2941f5e24 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -79,6 +79,36 @@ def _make_param_like( return new_param +def _maybe_dequantize_fp8( + full_tensor: torch.Tensor, + target_dtype: torch.dtype, + target_param_name: str, + param_sd: dict[str, torch.Tensor], +) -> torch.Tensor: + """Auto-dequantize an FP8 checkpoint weight when the model parameter expects a higher-precision type. + + Some modules (e.g. AdaLayerNormZero) don't accept quant_config, so their + parameters remain in higher precision even when the checkpoint stores FP8 + weights. In that case we multiply by the per-tensor weight_scale to + recover the original unquantized value. + """ + if not ( + full_tensor.dtype == torch.float8_e4m3fn and target_dtype != torch.float8_e4m3fn + ): + return full_tensor + + scale_key = target_param_name.rsplit(".", 1)[0] + ".weight_scale" + scale_tensor = param_sd.get(scale_key) + if scale_tensor is not None: + full_tensor = full_tensor.to(torch.float32) * scale_tensor.float() + logger.debug( + "Auto-dequantized FP8 weight %s using %s", + target_param_name, + scale_key, + ) + return full_tensor + + # TODO(PY): add compile option def maybe_load_fsdp_model( model_cls: type[nn.Module], @@ -331,6 +361,10 @@ def load_model_from_full_model_state_dict( else: target_dtype = meta_sharded_param.dtype + full_tensor = _maybe_dequantize_fp8( + full_tensor, target_dtype, target_param_name, custom_param_sd + ) + if full_tensor.dtype != target_dtype: mismatch_key = (full_tensor.dtype, target_dtype) if (