diff --git a/python/sglang/multimodal_gen/docs/quantization.md b/python/sglang/multimodal_gen/docs/quantization.md index e9c687d35..ffb6822cc 100644 --- a/python/sglang/multimodal_gen/docs/quantization.md +++ b/python/sglang/multimodal_gen/docs/quantization.md @@ -121,3 +121,44 @@ Choose the appropriate configuration based on your hardware and requirements: ### 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. + +## FP8 Quantization + +### Usage + +#### Option 1: Use Pre-quantized Models (Recommended) + +If available, you can directly use pre-quantized FP8 models from Hugging Face or other sources. Simply load them with SGLang: + +```bash +sglang generate \ + --model-path /path/to/FLUX.1-dev-FP8/ \ + --prompt "A Logo With Bold Large Text: SGL Diffusion" \ + --save-output +``` + +#### Option 2: Convert Your Own Models + +If you need to convert a model to FP8 format, use the provided conversion script: + +**Step 1: Convert the Model** + +```bash +# convert transformer to FP8 with block quantization +python -m sglang.multimodal_gen.tools.convert_hf_to_fp8 \ + --model-dir /path/to/FLUX.1-dev/transformer \ + --save-dir /path/to/FLUX.1-dev/transformer-FP8 \ + --strategy block \ + --block-size 128 128 +``` + +**Step 2: Run Inference** + +```bash +sglang generate \ + --model-path /path/to/FLUX.1-dev/ + # override transformer component with path to converted model + --transformer-path /path/to/FLUX.1-dev/transformer-FP8 + --prompt "A Logo With Bold Large Text: SGL Diffusion" \ + --save-output +``` diff --git a/python/sglang/multimodal_gen/runtime/layers/mlp.py b/python/sglang/multimodal_gen/runtime/layers/mlp.py index fc1de1bfb..4c1ed0dce 100644 --- a/python/sglang/multimodal_gen/runtime/layers/mlp.py +++ b/python/sglang/multimodal_gen/runtime/layers/mlp.py @@ -2,14 +2,24 @@ # SPDX-License-Identifier: Apache-2.0 +from typing import Optional + import torch import torch.nn as nn +from diffusers.models.activations import ( + GEGLU, + GELU, + ApproximateGELU, + LinearActivation, + SwiGLU, +) from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ) +from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig class MLP(nn.Module): @@ -26,6 +36,7 @@ class MLP(nn.Module): act_type: str = "gelu_pytorch_tanh", dtype: torch.dtype | None = None, prefix: str = "", + quant_config: QuantizationConfig = None, ): super().__init__() self.fc_in = ColumnParallelLinear( @@ -33,6 +44,7 @@ class MLP(nn.Module): mlp_hidden_dim, bias=True, gather_output=False, + quant_config=quant_config, ) self.act = get_act_fn(act_type) @@ -43,6 +55,7 @@ class MLP(nn.Module): output_dim, bias=True, input_is_parallel=True, + quant_config=quant_config, ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -50,3 +63,56 @@ class MLP(nn.Module): x = self.act(x) x, _ = self.fc_out(x) return x + + +class FeedForward(nn.Module): + r""" + A feed-forward layer. + + Parameters: + 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. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + bias (`bool`, defaults to True): Whether to use a bias in the linear layer. + """ + + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + activation_fn: str = "geglu", + inner_dim=None, + bias: bool = True, + ): + super().__init__() + if inner_dim is None: + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + if activation_fn == "gelu": + act_fn = GELU(dim, inner_dim, bias=bias) + if activation_fn == "gelu-approximate": + act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias) + elif activation_fn == "geglu": + act_fn = GEGLU(dim, inner_dim, bias=bias) + elif activation_fn == "geglu-approximate": + act_fn = ApproximateGELU(dim, inner_dim, bias=bias) + elif activation_fn == "swiglu": + act_fn = SwiGLU(dim, inner_dim, bias=bias) + elif activation_fn == "linear-silu": + act_fn = LinearActivation(dim, inner_dim, bias=bias, activation="silu") + + self.net = nn.ModuleList([]) + # project in + self.net.append(act_fn) + # dummy dropout layer to match with checkpoints compatible with diffusers + self.net.append(nn.Dropout(0.0)) + # project out + self.net.append(nn.Linear(inner_dim, dim_out, bias=bias)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index 0324606ae..c83f3b17f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -5,13 +5,16 @@ from typing import Literal, get_args from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) +from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config -QuantizationMethods = Literal[None] +QuantizationMethods = Literal["fp8"] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) # The customized quantization methods which will be added to this dict. -_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {} +_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = { + "fp8": Fp8Config, +} def register_quantization_config(quantization: str): diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py index 67c77f8a3..839692a84 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py @@ -22,6 +22,7 @@ def is_nunchaku_available() -> bool: try: import nunchaku # noqa + logger.debug("Nunchaku package detected") return True except Exception: return False diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py new file mode 100644 index 000000000..fcde0ab88 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py @@ -0,0 +1,498 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import torch +from torch.nn import Module +from torch.nn.parameter import Parameter + +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_tensor_model_parallel_world_size, +) +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 ( + BlockQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.utils.common import ( + cpu_has_amx_support, + get_bool_env_var, + use_intel_amx_backend, +) +from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading +from sglang.srt.layers.quantization.fp8_kernel import ( + is_fp8_fnuz, + per_token_group_quant_fp8, +) +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + can_auto_enable_marlin_fp8, + cutlass_fp8_supported, + dispatch_w8a8_block_fp8_linear, + input_to_float8, + normalize_e4m3fn_to_e4m3fnuz, + requant_weight_ue8m0_inplace, +) +from sglang.srt.layers.quantization.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + prepare_fp8_layer_for_marlin, +) +from sglang.srt.layers.quantization.utils import ( + convert_to_channelwise, + is_layer_skipped, + requantize_with_max_scale, +) + +if TYPE_CHECKING: + from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config + +_is_hip = current_platform.is_hip() +_is_cuda = current_platform.is_cuda() +_is_npu = current_platform.is_npu() +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = current_platform.is_cpu() +_is_fp8_fnuz = is_fp8_fnuz() +_use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") and _is_hip +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _use_aiter or _use_hip_int4: + pass + + +ACTIVATION_SCHEMES = ["static", "dynamic"] + +logger = logging.getLogger(__name__) + + +class Fp8Config(QuantizationConfig): + """Config class for FP8.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = False, + activation_scheme: str = "dynamic", + ignored_layers: Optional[List[str]] = None, + weight_block_size: List[int] = None, + ) -> None: + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + if is_checkpoint_fp8_serialized: + logger.info("Detected fp8 checkpoint.") + if activation_scheme not in ACTIVATION_SCHEMES: + raise ValueError(f"Unsupported activation scheme {activation_scheme}") + self.activation_scheme = activation_scheme + self.ignored_layers = ignored_layers or [] + if weight_block_size is not None: + if not is_checkpoint_fp8_serialized: + raise ValueError( + f"The block-wise quantization only supports fp8-serialized checkpoint for now." + ) + if len(weight_block_size) != 2: + raise ValueError( + f"The quantization block size of weight must have 2 dimensions, but got {len(weight_block_size)} dimensions." + ) + if activation_scheme != "dynamic": + raise ValueError( + f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme." + ) + self.weight_block_size = weight_block_size + + @classmethod + def get_name(cls) -> str: + return "fp8" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> Fp8Config: + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_fp8_serialized = "fp8" in quant_method + activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) + ignored_layers = cls.get_from_keys_or( + config, ["ignored_layers", "modules_to_not_convert"], None + ) + if ignored_layers: + # hacking ministral + ignored_layers = [layer.replace("model.", "") for layer in ignored_layers] + weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) + return cls( + is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + weight_block_size=weight_block_size, + ) + + 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 is_layer_skipped(prefix, self.ignored_layers): + return UnquantizedLinearMethod() + return Fp8LinearMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class Fp8LinearMethod(LinearMethodBase): + """Linear method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + + Limitations: + 1. Only support per-tensor quantization due to torch._scaled_mm support. + 2. Only support float8_e4m3fn data type due to the limitation of + torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856) + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Union[Fp8Config, W4AFp8Config]): + self.quant_config = quant_config + self.cutlass_fp8_supported = cutlass_fp8_supported() + + # For GPUs that lack FP8 hardware support, we can leverage the Marlin + # kernel for fast weight-only FP8 quantization + self.use_marlin = False + if _is_cuda: + force_marlin = get_bool_env_var("SGLANG_FORCE_FP8_MARLIN") + auto_enable = can_auto_enable_marlin_fp8() + self.use_marlin = force_marlin or auto_enable + + self.block_quant = self.quant_config.weight_block_size is not None + + self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear() + + 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, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + tp_size = get_tensor_model_parallel_world_size() + if self.block_quant: + block_n, block_k = ( + self.quant_config.weight_block_size[0], + self.quant_config.weight_block_size[1], + ) + # Required by row parallel + if tp_size > 1 and input_size // input_size_per_partition == tp_size: + if input_size_per_partition % block_k != 0: + raise ValueError( + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + # Required by column parallel or enabling merged weights + if ( + tp_size > 1 and output_size // output_size_per_partition == tp_size + ) or len(output_partition_sizes) > 1: + for output_partition_size in output_partition_sizes: + if output_partition_size % block_n != 0: + raise ValueError( + f"Weight output_partition_size = " + f"{output_partition_size} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # WEIGHT + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_fp8_serialized + else params_dtype + ) + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, input_size_per_partition, dtype=weight_dtype + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + # If checkpoint is serialized fp8, load them. + # Otherwise, wait until process_weights_after_loading. + if self.quant_config.is_checkpoint_fp8_serialized: + # WEIGHT SCALE + if self.block_quant: + if hasattr(self.quant_config, "activation_scheme"): + assert self.quant_config.activation_scheme == "dynamic" + elif hasattr(self.quant_config, "linear_activation_scheme"): + assert self.quant_config.linear_activation_scheme == "dynamic" + scale = BlockQuantScaleParameter( + data=torch.empty( + (output_size_per_partition + block_n - 1) // block_n, + (input_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + scale.format_ue8m0 = False + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale_inv", scale) + else: + scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("weight_scale", scale) + + # INPUT ACTIVATION SCALE + if ( + hasattr(self.quant_config, "activation_scheme") + and self.quant_config.activation_scheme == "static" + ) or ( + hasattr(self.quant_config, "linear_activation_scheme") + and self.quant_config.linear_activation_scheme == "static" + ): + scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + + scale[:] = torch.finfo(torch.float32).min + layer.register_parameter("input_scale", scale) + else: + layer.register_parameter("input_scale", None) + + def process_weights_after_loading(self, layer: Module) -> None: + if self.block_quant: + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_fp8_fnuz: + # activation_scheme: dynamic + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=layer.weight, + weight_scale=layer.weight_scale_inv, + input_scale=None, + ) + layer.input_scale = None + elif _is_cpu: + assert ( + _is_cpu_amx_available + ), "Fp8LinearMethod on CPU requires that CPU has AMX support" + _amx_process_weight_after_loading(layer, ["weight"]) + layer.weight_scale_inv = torch.nn.Parameter( + layer.weight_scale_inv.data, requires_grad=False + ) + return + else: + # For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0 + from sglang.srt.layers.quantization.fp8_utils import ( + deepgemm_w8a8_block_fp8_linear_with_fallback, + ) + from sglang.srt.model_loader.utils import ( + should_deepgemm_weight_requant_ue8m0, + ) + + if ( + should_deepgemm_weight_requant_ue8m0( + weight_block_size=getattr( + self.quant_config, "weight_block_size", None + ), + ) + and ( + self.w8a8_block_fp8_linear + is deepgemm_w8a8_block_fp8_linear_with_fallback + ) + and (not layer.weight_scale_inv.format_ue8m0) + ): + requant_weight_ue8m0_inplace( + layer.weight, + layer.weight_scale_inv, + self.quant_config.weight_block_size, + ) + layer.weight_scale_inv.format_ue8m0 = True + weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data + + layer.weight.data = weight.data + layer.weight_scale_inv.data = weight_scale.data + else: + layer.weight = Parameter(layer.weight.data, requires_grad=False) + + # If checkpoint not serialized fp8, quantize the weights. + if not self.quant_config.is_checkpoint_fp8_serialized: + if self.cutlass_fp8_supported or self.use_marlin: + # apply per-channel quantization default as + # cutlass sgl-kernel and marlin only support per-channel scale + qweight, weight_scale = per_token_group_quant_fp8( + layer.weight, layer.weight.shape[-1] + ) + weight_scale = weight_scale.t().contiguous() + else: + # per-tensor quantization + qweight, weight_scale = input_to_float8(layer.weight) + + # Update the layer with the new values. + layer.weight = Parameter(qweight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + layer.input_scale = None + + # If checkpoint is fp8, handle that there are N scales for N + # shards in a fused module + else: + layer.weight_scale = Parameter( + layer.weight_scale.data, requires_grad=False + ) + if ( + hasattr(self.quant_config, "activation_scheme") + and self.quant_config.activation_scheme == "static" + ) or ( + hasattr(self.quant_config, "linear_activation_scheme") + and self.quant_config.linear_activation_scheme == "static" + ): + layer.input_scale = Parameter( + layer.input_scale.data, requires_grad=False + ) + + # cutlass sgl-kernel and marlin only support per-channel scale + if self.cutlass_fp8_supported or self.use_marlin: + weight = layer.weight + weight_scale = convert_to_channelwise( + layer.weight_scale, layer.logical_widths + ) + else: + # Dequant -> Quant with max scale so we can run per tensor. + weight = layer.weight + weight_scale = layer.weight_scale + # If ROCm, normalize the weights and scales to e4m3fnuz + if _is_fp8_fnuz: + weight, weight_scale, input_scale = ( + normalize_e4m3fn_to_e4m3fnuz( + weight=weight, + weight_scale=weight_scale, + input_scale=layer.input_scale, + ) + ) + if input_scale is not None: + layer.input_scale = Parameter( + input_scale, requires_grad=False + ) + + weight_scale, weight = requantize_with_max_scale( + weight=weight, + weight_scale=weight_scale, + logical_widths=layer.logical_widths, + ) + + # Update layer with new values. + layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + if ( + hasattr(self.quant_config, "activation_scheme") + and self.quant_config.activation_scheme == "static" + ) or ( + hasattr(self.quant_config, "linear_activation_scheme") + and self.quant_config.linear_activation_scheme == "static" + ): + layer.input_scale = Parameter( + layer.input_scale.max(), requires_grad=False + ) + + if self.use_marlin: + if self.block_quant: + layer.weight_block_size = self.quant_config.weight_block_size + prepare_fp8_layer_for_marlin(layer, not self.block_quant) + # Activations not quantized for marlin. + del layer.input_scale + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if self.use_marlin: + return apply_fp8_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) + + if self.block_quant: + if use_intel_amx_backend(layer): + return torch.ops.sgl_kernel.fp8_scaled_mm_cpu( + x, + layer.weight, + layer.weight_scale_inv, + self.quant_config.weight_block_size, + bias, + x.dtype, + True, # is_vnni + ) + + if isinstance(x, tuple): + return self.w8a8_block_fp8_linear( + input=x[0], + weight=layer.weight, + block_size=self.quant_config.weight_block_size, + weight_scale=layer.weight_scale_inv, + input_scale=x[1], + bias=bias, + ) + + return self.w8a8_block_fp8_linear( + input=x, + weight=layer.weight, + block_size=self.quant_config.weight_block_size, + weight_scale=layer.weight_scale_inv, + input_scale=None, + bias=bias, + ) + + 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, + use_per_token_if_dynamic=False, + ) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py new file mode 100644 index 000000000..e69de29bb 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 67978e72b..20078b0f4 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 @@ -24,6 +24,7 @@ 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, + get_quant_config, ) from sglang.multimodal_gen.runtime.utils.logging_utils import get_log_level, init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE @@ -70,6 +71,7 @@ class TransformerLoader(ComponentLoader): ): """Load the transformer based on the model path, and inference args.""" config = get_diffusers_component_config(component_path=component_model_path) + hf_config = deepcopy(config) cls_name = config.pop("_class_name") if cls_name is None: @@ -94,6 +96,7 @@ class TransformerLoader(ComponentLoader): model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) nunchaku_config = server_args.nunchaku_config + if nunchaku_config is not None: nunchaku_config.model_cls = model_cls @@ -120,7 +123,7 @@ class TransformerLoader(ComponentLoader): ) logger.info( - "Loading %s from %s safetensors files %s, param_dtype: %s", + "Loading %s from %s safetensors file(s) %s, param_dtype: %s", cls_name, len(safetensors_list), f": {safetensors_list}" if get_log_level() == logging.DEBUG else "", @@ -128,11 +131,12 @@ class TransformerLoader(ComponentLoader): ) 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 + + if "quant_config" in inspect.signature(model_cls.__init__).parameters: + quant_config = get_quant_config(config) + init_params["quant_config"] = ( + quant_config if quant_config else nunchaku_config + ) # Load the model using FSDP loader model = maybe_load_fsdp_model( diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 76c2f0312..0ecb2c955 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -257,6 +257,7 @@ def load_model_from_full_model_state_dict( 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) + if meta_sharded_param is None: # For FSDP models, ensure all ranks process parameters consistently if strict or is_fsdp_model: @@ -264,12 +265,16 @@ def load_model_from_full_model_state_dict( f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." ) else: - logger.warning( - 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 + # use meta param dtype so quantized params (e.g. FP8) keep their dtype; + # for non-quantized models meta dtype equals param_dtype anyway + if meta_sharded_param is None: + # for nunchaku, some scales are patched later + target_dtype = full_tensor.dtype + else: + target_dtype = meta_sharded_param.dtype + if not hasattr(meta_sharded_param, "device_mesh"): full_tensor = full_tensor.to(device=device, dtype=target_dtype) actual_param = param_dict.get(target_param_name) @@ -350,6 +355,7 @@ def load_model_from_full_model_state_dict( ) meta_sharded_param = meta_sd.get(new_param_name) + meta_sharded_param_dtype = meta_sharded_param.dtype if "wcscales" in new_param_name or "wtscale" in new_param_name: init_like = torch.ones_like @@ -358,13 +364,13 @@ def load_model_from_full_model_state_dict( if not hasattr(meta_sharded_param, "device_mesh"): sharded_tensor = init_like( - meta_sharded_param, device=device, dtype=param_dtype + meta_sharded_param, device=device, dtype=meta_sharded_param_dtype ) if cpu_offload and not is_fsdp_model: sharded_tensor = sharded_tensor.cpu() else: full_tensor = init_like( - meta_sharded_param, device=device, dtype=param_dtype + meta_sharded_param, device=device, dtype=meta_sharded_param_dtype ) sharded_tensor = distribute_tensor( full_tensor, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index e8928eedc..1f2c4b806 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -18,7 +18,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn -from diffusers.models.attention import AttentionModuleMixin, FeedForward +from diffusers.models.attention import AttentionModuleMixin from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import ( AdaLayerNormContinuous, @@ -29,13 +29,12 @@ from torch.nn import LayerNorm as LayerNorm from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig from sglang.multimodal_gen.runtime.layers.attention import USPAttention - -# from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm as LayerNorm from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, ) +from sglang.multimodal_gen.runtime.layers.mlp import FeedForward from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( QuantizationConfig, ) @@ -255,13 +254,25 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): ) else: self.to_q = ColumnParallelLinear( - query_dim, self.inner_dim, bias=bias, gather_output=True + query_dim, + self.inner_dim, + bias=bias, + gather_output=True, + quant_config=quant_config, ) self.to_k = ColumnParallelLinear( - query_dim, self.inner_dim, bias=bias, gather_output=True + query_dim, + self.inner_dim, + bias=bias, + gather_output=True, + quant_config=quant_config, ) self.to_v = ColumnParallelLinear( - query_dim, self.inner_dim, bias=bias, gather_output=True + query_dim, + self.inner_dim, + bias=bias, + gather_output=True, + quant_config=quant_config, ) if not self.pre_only: self.to_out = torch.nn.ModuleList([]) @@ -296,18 +307,21 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): self.inner_dim, bias=added_proj_bias, gather_output=True, + quant_config=quant_config, ) self.add_k_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=True, + quant_config=quant_config, ) self.add_v_proj = ColumnParallelLinear( added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, gather_output=True, + quant_config=quant_config, ) self.to_add_out = ColumnParallelLinear( self.inner_dim, @@ -453,7 +467,7 @@ class FluxSingleTransformerBlock(nn.Module): quant_config=quant_config, prefix=f"{prefix}.attn" if prefix else "attn", ) - if NunchakuAdaLayerNormZeroSingle is not None: + if is_nunchaku_available(): self.norm = NunchakuAdaLayerNormZeroSingle(self.norm, scale_shift=0) else: self.proj_mlp = ColumnParallelLinear( @@ -461,6 +475,7 @@ class FluxSingleTransformerBlock(nn.Module): self.mlp_hidden_dim, bias=True, gather_output=True, + quant_config=quant_config, ) self.act_mlp = nn.GELU(approximate="tanh") self.proj_out = ColumnParallelLinear( @@ -468,6 +483,7 @@ class FluxSingleTransformerBlock(nn.Module): dim, bias=True, gather_output=True, + quant_config=quant_config, ) self.attn = FluxAttention( query_dim=dim, @@ -477,6 +493,7 @@ class FluxSingleTransformerBlock(nn.Module): bias=True, eps=1e-6, pre_only=True, + quant_config=quant_config, ) def forward( @@ -579,12 +596,13 @@ class FluxTransformerBlock(nn.Module): and hasattr(quant_config, "get_name") and quant_config.get_name() == "svdquant" and is_nunchaku_available() - and NunchakuFeedForward is not None ) self.use_nunchaku_structure = nunchaku_enabled self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") self.ff_context = FeedForward( - dim=dim, dim_out=dim, activation_fn="gelu-approximate" + dim=dim, + dim_out=dim, + activation_fn="gelu-approximate", ) if nunchaku_enabled: nunchaku_kwargs = { @@ -594,11 +612,10 @@ class FluxTransformerBlock(nn.Module): } self.ff = NunchakuFeedForward(self.ff, **nunchaku_kwargs) self.ff_context = NunchakuFeedForward(self.ff_context, **nunchaku_kwargs) - if NunchakuAdaLayerNormZero is not None: - self.norm1 = NunchakuAdaLayerNormZero(self.norm1, scale_shift=0) - self.norm1_context = NunchakuAdaLayerNormZero( - self.norm1_context, scale_shift=0 - ) + self.norm1 = NunchakuAdaLayerNormZero(self.norm1, scale_shift=0) + self.norm1_context = NunchakuAdaLayerNormZero( + self.norm1_context, scale_shift=0 + ) def forward( self, diff --git a/python/sglang/multimodal_gen/runtime/utils/common.py b/python/sglang/multimodal_gen/runtime/utils/common.py index 134609b2a..f9f5df924 100644 --- a/python/sglang/multimodal_gen/runtime/utils/common.py +++ b/python/sglang/multimodal_gen/runtime/utils/common.py @@ -281,3 +281,28 @@ def get_bool_env_var(name: str, default: str = "false") -> bool: _warned_bool_env_var_keys.add(value) return value in truthy_values + + +try: + import sgl_kernel # noqa: F401 + + is_intel_amx_backend_available = hasattr( + torch.ops.sgl_kernel, "convert_weight_packed" + ) +except: + is_intel_amx_backend_available = False + +try: + # move torch._C._cpu._is_amx_tile_supported() from cpu_has_amx_support + # to support torch compile + is_amx_tile_supported = torch._C._cpu._is_amx_tile_supported() +except: + is_amx_tile_supported = False + + +def cpu_has_amx_support(): + return is_amx_tile_supported and is_intel_amx_backend_available + + +def use_intel_amx_backend(layer): + return getattr(layer, "use_intel_amx_backend", 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 6f49ed0b1..f348f90f7 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -26,7 +26,7 @@ import shutil import time from functools import reduce from pathlib import Path -from typing import Any, Optional, Union, cast +from typing import Any, Dict, List, Optional, Union, cast from diffusers.loaders.lora_base import ( _best_guess_weight_name, # watch out for potetential removal from diffusers @@ -42,6 +42,10 @@ from safetensors import safe_open from transformers import AutoConfig, PretrainedConfig from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES +from sglang.multimodal_gen.runtime.layers.quantization import ( + QuantizationConfig, + get_quantization_config, +) from sglang.multimodal_gen.runtime.loader.weight_utils import get_lock from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -320,6 +324,80 @@ def get_diffusers_component_config( return combined_config +def replace_prefix(key: str, prefix_mapping: dict[str, str]) -> str: + for prefix, new_prefix in prefix_mapping.items(): + if key.startswith(prefix): + key = key.replace(prefix, new_prefix, 1) + return key + + +def get_quant_config( + model_config, + packed_modules_mapping: Dict[str, List[str]] = {}, + remap_prefix: Dict[str, str] | None = None, +) -> QuantizationConfig: + if "quantization_config" not in model_config: + return None + quant_cls = get_quantization_config( + model_config["quantization_config"]["quant_method"] + ) + + # GGUF doesn't have config file + if model_config["quantization_config"]["quant_method"] == "gguf": + return quant_cls.from_config({}) + + # Read the quantization config from the HF model config, if available. + hf_quant_config = model_config["quantization_config"] + # some vision model may keep quantization_config in their text_config + hf_text_config = getattr(model_config, "text_config", None) + if hf_quant_config is None and hf_text_config is not None: + hf_quant_config = getattr(hf_text_config, "quantization_config", None) + if hf_quant_config is None: + # compressed-tensors uses a compressions_config + hf_quant_config = getattr(model_config, "compression_config", None) + if hf_quant_config is not None: + hf_quant_config["packed_modules_mapping"] = packed_modules_mapping + return quant_cls.from_config(hf_quant_config) + # In case of bitsandbytes/QLoRA, get quant config from the adapter model. + else: + model_name_or_path = model_config["model_path"] + is_local = os.path.isdir(model_name_or_path) + hf_folder = model_name_or_path + + possible_config_filenames = quant_cls.get_config_filenames() + + # If the quantization config is not found, use the default config. + if not possible_config_filenames: + return quant_cls() + + config_files = glob.glob(os.path.join(hf_folder, "*.json")) + + quant_config_files = [ + f for f in config_files if any(f.endswith(x) for x in possible_config_filenames) + ] + if len(quant_config_files) == 0: + raise ValueError( + f"Cannot find the config file for {model_config['quantization_config']['quant_method']}" + ) + if len(quant_config_files) > 1: + raise ValueError( + f"Found multiple config files for {model_config['quantization_config']['quant_method']}: " + f"{quant_config_files}" + ) + + quant_config_file = quant_config_files[0] + with open(quant_config_file) as f: + config = json.load(f) + if remap_prefix is not None: + exclude_modules = [ + replace_prefix(key, remap_prefix) + for key in config["quantization"]["exclude_modules"] + ] + config["quantization"]["exclude_modules"] = exclude_modules + config["packed_modules_mapping"] = packed_modules_mapping + return quant_cls.from_config(config) + + # Models don't use the same configuration key for determining the maximum # context length. Store them here so we can sanely check them. # NOTE: The ordering here is important. Some models have two of these and we diff --git a/python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py b/python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py new file mode 100644 index 000000000..2759364a6 --- /dev/null +++ b/python/sglang/multimodal_gen/tools/convert_hf_to_fp8.py @@ -0,0 +1,310 @@ +# copied and adapted from Slime +""" +Convert HuggingFace safetensors model to FP8 format for efficient inference. + +Example usage: + # convert FLUX.1-dev transformer to FP8 + python -m sglang.multimodal_gen.tools.convert_hf_to_fp8 \ + --model-dir /path/to/FLUX.1-dev/transformer \ + --save-dir /path/to/FLUX.1-dev/transformer-FP8 \ + --strategy block \ + --block-size 128 128 + +Options: + --model-dir MODEL_DIR + path to the directory of the HF safetensors model (e.g., transformer subfolder) + --save-dir SAVE_DIR + path to the directory to save the converted FP8 model + --strategy {block,channel,tensor} + quantization strategy (default: block) + --block-size [BLOCK_SIZE ...] + block size for block quantization, e.g., --block-size 128 128 + --max-workers MAX_WORKERS + number of worker threads for parallel processing (default: 1) +""" + +import argparse +import gc +import json +import os +import shutil +import threading +from concurrent.futures import ThreadPoolExecutor + +import safetensors +import safetensors.torch +import torch +import torch.nn.functional as F +from tqdm import tqdm + +FP8_INFO = torch.finfo(torch.float8_e4m3fn) +FP8_MAX, FP8_MIN = FP8_INFO.max, FP8_INFO.min + + +def ceildiv(a, b): + return -(-a // b) + + +def block_fp8(weight, block_size): + + # per block quant + block_n, block_k = block_size[0], block_size[1] + + shape_0, shape_1 = weight.shape + + n_tiles = ceildiv(shape_0, block_n) + k_tiles = ceildiv(shape_1, block_k) + + q_weight = F.pad( + weight, + (0, k_tiles * block_k - shape_1, 0, n_tiles * block_n - shape_0), + mode="constant", + value=0.0, + ) + + qweight = q_weight.reshape(n_tiles, block_n, k_tiles, block_k) + block_max = torch.max(torch.abs(qweight), dim=1, keepdim=True)[0] + block_max = torch.max(block_max, dim=3, keepdim=True)[0] + + scale = block_max.to(torch.float32) / FP8_MAX + qweight = ( + (qweight / scale) + .clamp(min=FP8_MIN, max=FP8_MAX) + .reshape((n_tiles * block_n, k_tiles * block_k)) + .to(torch.float8_e4m3fn) + ) + qweight = qweight[:shape_0, :shape_1].clone().detach() + scale = scale.squeeze() + + return qweight, scale + + +def channel_fp8(weight): + channel_max = torch.max(weight.abs(), dim=-1, keepdim=True)[0] + scale = channel_max.clamp(min=1e-12).to(torch.float32) / FP8_MAX + qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX) + qweight = qweight.to(torch.float8_e4m3fn) + return qweight, scale + + +def tensor_fp8(weight): + scale = weight.abs().max().clamp(min=1e-12).to(torch.float32) / FP8_MAX + qweight = (weight / scale).clamp(min=FP8_MIN, max=FP8_MAX) + qweight = qweight.to(torch.float8_e4m3fn) + scale = scale.view(1) + return qweight, scale + + +def quant_fp8(weight, strategy, block_size=None): + if strategy == "tensor": + return tensor_fp8(weight) + elif strategy == "channel": + return channel_fp8(weight) + else: + return block_fp8(weight, block_size) + + +class ConversionResult: + def __init__(self): + self.lock = threading.Lock() + self.weight_map = {} + self.param_count = 0 + self.modules_to_not_convert = [] + + def add_result(self, filename, q_weights, module_names): + with self.lock: + for k, v in q_weights.items(): + self.weight_map[k] = filename + self.param_count += len(v) + self.modules_to_not_convert.extend(module_names) + + +def process_file( + input_path, output_path, filename, strategy, block_size, result_collector +): + if not filename.endswith(".safetensors"): + return + + print(f"Processing {filename}, memory usage: {torch.cuda.memory_allocated()}") + weights = {} + q_weights = {} + + with safetensors.safe_open( + os.path.join(input_path, filename), framework="pt", device="cuda" + ) as f: + for k in f.keys(): + weights[k] = f.get_tensor(k) + + modules_to_not_convert = [] + for key in weights.keys(): + if ( + "weight" in key + and "layernorm" not in key + and "embed" not in key + and "router" not in key + and "mlp.gate." not in key + and "norm" not in key + and "lm_head" not in key + and "eh_proj" not in key + and "net" not in key + and "proj_out.weight" != key + ): + qw, s = quant_fp8(weights[key], strategy, block_size) + q_weights[key] = qw + if block_size: + scale_name = key.replace(".weight", ".weight_scale_inv") + else: + scale_name = key.replace(".weight", ".weight_scale") + q_weights[scale_name] = s + else: + modules_to_not_convert.append(key.replace(".weight", "")) + q_weights[key] = weights[key] + + safetensors.torch.save_file( + q_weights, os.path.join(output_path, filename), metadata={"format": "pt"} + ) + + result_collector.add_result(filename, q_weights, modules_to_not_convert) + + +def convert_fp8(input_path, output_path, strategy, block_size=None, max_workers=4): + input_path = os.path.abspath(input_path) + os.makedirs(output_path, exist_ok=True) + + for filename in os.listdir(input_path): + if not filename.endswith(".safetensors") and not os.path.isdir( + os.path.join(input_path, filename) + ): + shutil.copyfile( + os.path.join(input_path, filename), os.path.join(output_path, filename) + ) + + safetensors_files = [ + f for f in os.listdir(input_path) if f.endswith(".safetensors") + ] + + result_collector = ConversionResult() + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [] + for filename in safetensors_files: + future = executor.submit( + process_file, + input_path, + output_path, + filename, + strategy, + block_size, + result_collector, + ) + futures.append(future) + + for future in tqdm(futures, desc="Processing files"): + future.result() + + if strategy == "block" or strategy == "tensor": + quantization_config = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + } + if block_size: + quantization_config["weight_block_size"] = block_size + if len(result_collector.modules_to_not_convert) > 0: + quantization_config["modules_to_not_convert"] = list( + set(result_collector.modules_to_not_convert) + ) + else: + quant_group = { + "group_0": { + "input_activations": { + "actorder": None, + "block_structure": None, + "dynamic": True, + "group_size": None, + "num_bits": 8, + "observer": None, + "observer_kwargs": {}, + "strategy": "token", + "symmetric": True, + "type": "float", + }, + "output_activations": None, + "targets": ["Linear"], + "weights": { + "actorder": None, + "block_structure": None, + "dynamic": False, + "group_size": None, + "num_bits": 8, + "observer": "minmax", + "observer_kwargs": {}, + "strategy": strategy, + "symmetric": True, + "type": "float", + }, + }, + } + quantization_config = { + "config_groups": quant_group, + "format": "float-quantized", + "ignore": list(set(result_collector.modules_to_not_convert)), + "quant_method": "compressed-tensors", + "quantization_status": "compressed", + } + + config_path = os.path.join(input_path, "config.json") + if os.path.exists(config_path): + cfg = json.load(open(config_path)) + cfg["quantization_config"] = quantization_config + json.dump(cfg, open(os.path.join(output_path, "config.json"), "w"), indent=2) + + index_dict = { + "weight_map": result_collector.weight_map, + "metadata": {"total_size": result_collector.param_count}, + } + json.dump( + index_dict, + open(os.path.join(output_path, "model.safetensors.index.json"), "w"), + indent=2, + ) + + gc.collect() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--model-dir", + type=str, + help="Path to the directory of the HF safetensors model.", + ) + parser.add_argument( + "--save-dir", + type=str, + help="Path to the directory to save the converted model.", + ) + parser.add_argument( + "--strategy", type=str, default="block", choices=["block", "channel", "tensor"] + ) + parser.add_argument( + "--block-size", type=int, nargs="*", default=None, help="eg. --block-size 32 32" + ) + parser.add_argument( + "--max-workers", + type=int, + default=1, + help="Number of worker threads for parallel processing", + ) + args = parser.parse_args() + + if not os.path.exists(args.save_dir): + print(f"Creating directory {args.save_dir}") + os.makedirs(args.save_dir) + elif not os.path.isdir(args.save_dir): + raise ValueError("The save_dir should be a directory.") + + convert_fp8( + args.model_dir, args.save_dir, args.strategy, args.block_size, args.max_workers + )