[diffusion] refactor: reuse srt quantization contracts and mxfp8 kernels (#36063)
This commit is contained in:
Mick
@@ -13,6 +13,7 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig):
|
||||
# H3 fuses Q/K/V, so split projections are stacked for the fused LoRA layer
|
||||
param_names_mapping: dict = field(
|
||||
default_factory=lambda: {
|
||||
r"^(.*)\.weight_scale$": r"\1.weight_scale_inv",
|
||||
r"^(.*\.lora_[AB])\.[^.]+$": r"\1",
|
||||
r"^base_model\.model\.(.*\.lora_[AB])$": r"\1",
|
||||
r"^transformer\.(.*\.lora_[AB])$": r"\1",
|
||||
|
||||
@@ -24,7 +24,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.mxfp4 import Mxfp4Config
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.mxfp4_npu import (
|
||||
NPUMXFP4Config,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.mxfp8_npu import MXFP8Config
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config
|
||||
|
||||
QuantizationMethods = Literal[
|
||||
"fp8",
|
||||
|
||||
@@ -6,7 +6,6 @@ from typing import Any, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from packaging import version
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.linear import (
|
||||
LinearBase,
|
||||
@@ -18,39 +17,17 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
|
||||
from sglang.srt.layers.quantization.bitsandbytes import (
|
||||
BitsAndBytesConfig as SRTBitsAndBytesConfig,
|
||||
)
|
||||
from sglang.srt.layers.quantization.bitsandbytes import (
|
||||
calculate_quant_ratio,
|
||||
is_layer_skipped_bnb,
|
||||
require_bitsandbytes,
|
||||
)
|
||||
|
||||
|
||||
def _require_bitsandbytes() -> None:
|
||||
try:
|
||||
import bitsandbytes
|
||||
|
||||
if version.parse(bitsandbytes.__version__) < version.parse("0.46.1"):
|
||||
raise ImportError(
|
||||
"bitsandbytes version is wrong. Please install bitsandbytes>=0.46.1."
|
||||
)
|
||||
except ImportError as err:
|
||||
raise ImportError(
|
||||
"Please install bitsandbytes>=0.46.1 via "
|
||||
"`pip install bitsandbytes>=0.46.1` to use bitsandbytes quantizer."
|
||||
) from err
|
||||
|
||||
|
||||
def _calculate_quant_ratio(dtype: torch.dtype) -> int:
|
||||
if dtype.is_floating_point:
|
||||
return torch.finfo(dtype).bits // torch.iinfo(torch.uint8).bits
|
||||
return torch.iinfo(dtype).bits // torch.iinfo(torch.uint8).bits
|
||||
|
||||
|
||||
def _is_layer_skipped(prefix: str, skipped_modules: list[str]) -> bool:
|
||||
components = prefix.split(".")
|
||||
if any(module_name in components for module_name in skipped_modules):
|
||||
return True
|
||||
|
||||
prefixes = {".".join(components[: i + 1]) for i in range(len(components))}
|
||||
return bool(set(skipped_modules) & prefixes)
|
||||
|
||||
|
||||
class BitsAndBytesConfig(QuantizationConfig):
|
||||
class BitsAndBytesConfig(SRTBitsAndBytesConfig, QuantizationConfig):
|
||||
"""Config class for pre-quantized bitsandbytes 4-bit checkpoints."""
|
||||
|
||||
def __init__(
|
||||
@@ -66,79 +43,26 @@ class BitsAndBytesConfig(QuantizationConfig):
|
||||
llm_int8_skip_modules: list[str] | None = None,
|
||||
llm_int8_threshold: float = 6.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.load_in_8bit = load_in_8bit
|
||||
self.load_in_4bit = load_in_4bit
|
||||
self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype
|
||||
self.bnb_4bit_quant_storage = bnb_4bit_quant_storage
|
||||
self.bnb_4bit_quant_type = bnb_4bit_quant_type
|
||||
self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant
|
||||
self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload
|
||||
self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight
|
||||
self.llm_int8_skip_modules = llm_int8_skip_modules or []
|
||||
self.llm_int8_threshold = llm_int8_threshold
|
||||
|
||||
super().__init__(
|
||||
load_in_8bit=load_in_8bit,
|
||||
load_in_4bit=load_in_4bit,
|
||||
bnb_4bit_compute_dtype=bnb_4bit_compute_dtype,
|
||||
bnb_4bit_quant_storage=bnb_4bit_quant_storage,
|
||||
bnb_4bit_quant_type=bnb_4bit_quant_type,
|
||||
bnb_4bit_use_double_quant=bnb_4bit_use_double_quant,
|
||||
llm_int8_enable_fp32_cpu_offload=llm_int8_enable_fp32_cpu_offload,
|
||||
llm_int8_has_fp16_weight=llm_int8_has_fp16_weight,
|
||||
llm_int8_skip_modules=llm_int8_skip_modules,
|
||||
llm_int8_threshold=llm_int8_threshold,
|
||||
)
|
||||
if self.load_in_8bit or not self.load_in_4bit:
|
||||
raise ValueError("SGLang diffusion only supports bitsandbytes 4-bit.")
|
||||
if self.bnb_4bit_quant_storage != "uint8":
|
||||
raise ValueError(
|
||||
f"Unsupported bnb_4bit_quant_storage: {self.bnb_4bit_quant_storage}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
return "bitsandbytes"
|
||||
|
||||
def get_scaled_act_names(self) -> list[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.float32, torch.float16, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 70
|
||||
|
||||
@staticmethod
|
||||
def get_config_filenames() -> list[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> BitsAndBytesConfig:
|
||||
def get_safe_value(keys, default_value=None):
|
||||
try:
|
||||
value = QuantizationConfig.get_from_keys(config, keys)
|
||||
return value if value is not None else default_value
|
||||
except ValueError:
|
||||
return default_value
|
||||
|
||||
return cls(
|
||||
load_in_8bit=get_safe_value(["load_in_8bit"], False),
|
||||
load_in_4bit=get_safe_value(["load_in_4bit"], True),
|
||||
bnb_4bit_compute_dtype=get_safe_value(
|
||||
["bnb_4bit_compute_dtype"], "float32"
|
||||
),
|
||||
bnb_4bit_quant_storage=get_safe_value(["bnb_4bit_quant_storage"], "uint8"),
|
||||
bnb_4bit_quant_type=get_safe_value(["bnb_4bit_quant_type"], "fp4"),
|
||||
bnb_4bit_use_double_quant=get_safe_value(
|
||||
["bnb_4bit_use_double_quant"], False
|
||||
),
|
||||
llm_int8_enable_fp32_cpu_offload=get_safe_value(
|
||||
["llm_int8_enable_fp32_cpu_offload"], False
|
||||
),
|
||||
llm_int8_has_fp16_weight=get_safe_value(
|
||||
["llm_int8_has_fp16_weight"], False
|
||||
),
|
||||
llm_int8_skip_modules=get_safe_value(["llm_int8_skip_modules"], []),
|
||||
llm_int8_threshold=get_safe_value(["llm_int8_threshold"], 6.0),
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> Optional[QuantizeMethodBase]:
|
||||
if isinstance(layer, LinearBase):
|
||||
if _is_layer_skipped(prefix, self.llm_int8_skip_modules):
|
||||
if is_layer_skipped_bnb(prefix, self.llm_int8_skip_modules):
|
||||
return UnquantizedLinearMethod()
|
||||
return BitsAndBytesLinearMethod(self)
|
||||
return None
|
||||
@@ -148,7 +72,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase):
|
||||
"""Linear method for pre-quantized bitsandbytes 4-bit weights."""
|
||||
|
||||
def __init__(self, quant_config: BitsAndBytesConfig):
|
||||
_require_bitsandbytes()
|
||||
require_bitsandbytes()
|
||||
self.quant_config = quant_config
|
||||
|
||||
def create_weights(
|
||||
@@ -161,7 +85,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase):
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
) -> None:
|
||||
quant_ratio = _calculate_quant_ratio(params_dtype)
|
||||
quant_ratio = calculate_quant_ratio(params_dtype)
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
total_size = input_size_per_partition * output_size_per_partition
|
||||
if total_size % quant_ratio != 0:
|
||||
@@ -255,11 +179,11 @@ class BitsAndBytes4BitLinear(nn.Module):
|
||||
compute_dtype: torch.dtype | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
_require_bitsandbytes()
|
||||
require_bitsandbytes()
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.compute_dtype = compute_dtype
|
||||
quant_ratio = _calculate_quant_ratio(compute_dtype or torch.get_default_dtype())
|
||||
quant_ratio = calculate_quant_ratio(compute_dtype or torch.get_default_dtype())
|
||||
total_size = in_features * out_features
|
||||
if total_size % quant_ratio != 0:
|
||||
raise ValueError(
|
||||
|
||||
@@ -4,51 +4,21 @@
|
||||
# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/quantization/base_config.py
|
||||
|
||||
import inspect
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.multimodal_gen.runtime.layers.quantization import QuantizationMethods
|
||||
else:
|
||||
QuantizationMethods = str
|
||||
from sglang.srt.layers.quantization.base_config import (
|
||||
QuantizationConfig as SRTQuantizationConfig,
|
||||
)
|
||||
from sglang.srt.layers.quantization.base_config import (
|
||||
QuantizeMethodBase as SRTQuantizeMethodBase,
|
||||
)
|
||||
|
||||
|
||||
class QuantizeMethodBase(ABC):
|
||||
"""Base class for different quantized methods."""
|
||||
|
||||
@abstractmethod
|
||||
def create_weights(
|
||||
self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
|
||||
):
|
||||
"""Create weights for a layer.
|
||||
|
||||
The weights will be set as attributes of the layer."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
|
||||
"""Apply the weights in layer to the input tensor.
|
||||
|
||||
Expects create_weights to have been called before on the layer."""
|
||||
raise NotImplementedError
|
||||
|
||||
# Not required functions
|
||||
class QuantizeMethodBase(SRTQuantizeMethodBase):
|
||||
def embedding(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
|
||||
"""Gather embeddings in the layer based on indices in the input tensor.
|
||||
|
||||
Expects create_weights to have been called before on the layer."""
|
||||
raise NotImplementedError
|
||||
|
||||
def process_weights_after_loading(self, layer: nn.Module) -> None:
|
||||
"""Process the weight after loading.
|
||||
|
||||
This can be used for example, to transpose weights for computation.
|
||||
"""
|
||||
return
|
||||
|
||||
|
||||
def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool:
|
||||
"""
|
||||
@@ -62,98 +32,13 @@ def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) ->
|
||||
return class_embedding is not None and class_embedding is not base_embedding
|
||||
|
||||
|
||||
class QuantizationConfig(ABC):
|
||||
"""Base class for quantization configs."""
|
||||
|
||||
class QuantizationConfig(SRTQuantizationConfig):
|
||||
# for quantization frameworks with a separate quantized model provided, e.g. Nunchaku
|
||||
quantized_model_path: str | None = None
|
||||
checkpoint_uses_native_qkv_layout: bool = False
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# mapping is updated by models as they initialize
|
||||
self.packed_modules_mapping: dict[str, list[str]] = dict()
|
||||
|
||||
@abstractmethod
|
||||
def get_name(self) -> QuantizationMethods:
|
||||
"""Name of the quantization method."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def get_supported_act_dtypes(self) -> list[torch.dtype]:
|
||||
"""List of supported activation dtypes."""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
"""Minimum GPU capability to support the quantization method.
|
||||
|
||||
E.g., 70 for Volta, 75 for Turing, 80 for Ampere.
|
||||
This requirement is due to the custom CUDA kernels used by the
|
||||
quantization method.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def get_config_filenames() -> list[str]:
|
||||
"""List of filenames to search for in the model directory."""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig":
|
||||
"""Create a config class from the model's quantization config."""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(
|
||||
cls, hf_quant_cfg, user_quant
|
||||
) -> QuantizationMethods | None:
|
||||
"""
|
||||
Detects if this quantization method can support a given checkpoint
|
||||
format by overriding the user specified quantization method --
|
||||
this method should only be overwritten by subclasses in exceptional
|
||||
circumstances
|
||||
"""
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any:
|
||||
"""Get a value from the model's quantization config."""
|
||||
for key in keys:
|
||||
if key in config:
|
||||
return config[key]
|
||||
raise ValueError(
|
||||
f"Cannot find any of {keys} in the model's " "quantization config."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any:
|
||||
"""Get a optional value from the model's quantization config."""
|
||||
try:
|
||||
return QuantizationConfig.get_from_keys(config, keys)
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
@abstractmethod
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> QuantizeMethodBase | None:
|
||||
"""Get the quantize method to use for the quantized layer.
|
||||
|
||||
Args:
|
||||
layer: The layer for the quant method.
|
||||
prefix: The full name of the layer in the state dict
|
||||
Returns:
|
||||
The quantize method. None if the given layer doesn't support quant
|
||||
method.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_cache_scale(self, name: str) -> str | None:
|
||||
return None
|
||||
def get_scaled_act_names(self) -> list[str]:
|
||||
return []
|
||||
|
||||
def supports_input_partition(
|
||||
self, prefix: str, input_size_per_partition: int
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
|
||||
from typing import TYPE_CHECKING, List, Optional, Union
|
||||
|
||||
import torch
|
||||
from torch.nn import Module
|
||||
@@ -15,6 +15,7 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.linear import (
|
||||
LinearBase,
|
||||
LinearMethodBase,
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
@@ -35,6 +36,7 @@ from sglang.multimodal_gen.runtime.utils.common import (
|
||||
use_intel_amx_backend,
|
||||
)
|
||||
from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config as SRTFp8Config
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
apply_fp8_linear,
|
||||
can_auto_enable_marlin_fp8,
|
||||
@@ -69,89 +71,23 @@ if USE_AITER or _use_hip_int4:
|
||||
pass
|
||||
|
||||
|
||||
ACTIVATION_SCHEMES = ["static", "dynamic"]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Fp8Config(QuantizationConfig):
|
||||
class Fp8Config(SRTFp8Config, QuantizationConfig):
|
||||
"""Config class for FP8.
|
||||
|
||||
No-arg ``Fp8Config()`` selects online (post-load) weight quantization:
|
||||
``is_checkpoint_fp8_serialized=False`` with ``activation_scheme="dynamic"``.
|
||||
"""
|
||||
|
||||
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,
|
||||
packed_modules_mapping: Optional[Dict[str, List[str]]] = 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 []
|
||||
self.packed_modules_mapping = packed_modules_mapping or {}
|
||||
if weight_block_size is not None:
|
||||
if not is_checkpoint_fp8_serialized:
|
||||
raise ValueError(
|
||||
"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,
|
||||
@@ -162,9 +98,6 @@ class Fp8Config(QuantizationConfig):
|
||||
return Fp8LinearMethod(self)
|
||||
return None
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
|
||||
|
||||
class Fp8LinearMethod(LinearMethodBase):
|
||||
"""Linear method for FP8.
|
||||
|
||||
@@ -53,7 +53,8 @@ class ModelSlimMXFP8Scheme(ModelSlimLinearScheme):
|
||||
|
||||
# msmodelslim exports weight_scale as uint8, shape [out, in/32].
|
||||
# NOTE: This parameter is intentionally named "weight_scale" (not
|
||||
# "weight_scale_inv" as used in mxfp8_npu.py) because the weight loader
|
||||
# "weight_scale_inv" as used by SRT's NPUMXFP8LinearMethod) because the
|
||||
# weight loader
|
||||
# matches parameter names to checkpoint keys, and msmodelslim checkpoints
|
||||
# store this tensor under the key "<layer>.weight_scale".
|
||||
scale_dim = input_size_per_partition // MXFP8_BLOCK_SIZE
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MXFP8 diffusion adapter backed by SRT's dense linear kernels."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.linear import (
|
||||
LinearBase,
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
|
||||
NPUMXFP8LinearMethod,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config as SRTFp8Config
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod as SRTFp8LinearMethod
|
||||
|
||||
|
||||
class MXFP8Config(SRTFp8Config, QuantizationConfig):
|
||||
"""Route diffusion linears through SRT's MXFP8 implementation."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
is_checkpoint_fp8_serialized: bool = False,
|
||||
layer_markers: dict[str, dict[str, Any]] | None = None,
|
||||
ignored_layers: list[str] | None = None,
|
||||
) -> None:
|
||||
if current_platform.is_mps():
|
||||
raise ValueError("MXFP8 is not supported on MPS")
|
||||
super().__init__(
|
||||
is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
|
||||
activation_scheme="dynamic",
|
||||
ignored_layers=ignored_layers,
|
||||
weight_block_size=[1, 32] if is_checkpoint_fp8_serialized else None,
|
||||
use_mxfp8=True,
|
||||
)
|
||||
self.layer_markers = layer_markers
|
||||
self.checkpoint_uses_native_qkv_layout = layer_markers is not None
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 90
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> MXFP8Config:
|
||||
quant_method = str(cls.get_from_keys(config, ["quant_method"])).lower()
|
||||
if "mxfp8" not in quant_method:
|
||||
raise ValueError(f"Expected an MXFP8 checkpoint, got {quant_method!r}")
|
||||
activation_scheme = cls.get_from_keys_or(
|
||||
config, ["activation_scheme"], "dynamic"
|
||||
)
|
||||
if activation_scheme != "dynamic":
|
||||
raise ValueError("MXFP8 only supports dynamic activation scaling")
|
||||
ignored_layers = cls.get_from_keys_or(
|
||||
config, ["ignored_layers", "modules_to_not_convert"], None
|
||||
)
|
||||
return cls(
|
||||
is_checkpoint_fp8_serialized=True,
|
||||
ignored_layers=ignored_layers,
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> QuantizeMethodBase | None:
|
||||
if not isinstance(layer, LinearBase):
|
||||
return None
|
||||
if self.layer_markers is not None and prefix not in self.layer_markers:
|
||||
return UnquantizedLinearMethod()
|
||||
if current_platform.is_npu():
|
||||
return NPUMXFP8LinearMethod(self)
|
||||
return SRTFp8LinearMethod(self)
|
||||
|
||||
|
||||
__all__ = ["MXFP8Config"]
|
||||
@@ -1,176 +0,0 @@
|
||||
"""Online MXFP8 quantization for Diffusion models on Ascend NPU.
|
||||
|
||||
Provides ``MXFP8Config`` (registered as ``"mxfp8"``) and
|
||||
``NPUMXFP8DiffusionLinearMethod`` which quantise FP16/BF16 weights to MXFP8
|
||||
at load time and use ``npu_dynamic_mx_quant`` + ``npu_quant_matmul`` for
|
||||
inference, mirroring the LLM-side ``NPUMXFP8LinearMethod``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
_is_npu = current_platform.is_npu()
|
||||
|
||||
if _is_npu:
|
||||
import torch_npu
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.linear import LinearBase, LinearMethodBase
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.parameter import ModelWeightParameter
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
MXFP8_BLOCK_SIZE = 32
|
||||
|
||||
|
||||
class MXFP8Config(QuantizationConfig):
|
||||
"""Config for online MXFP8 quantization on Ascend NPU (Diffusion)."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
return "mxfp8"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
|
||||
return [torch.bfloat16, torch.float16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 0 # NPU, not CUDA
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> List[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> MXFP8Config:
|
||||
return cls()
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> Optional[QuantizeMethodBase]:
|
||||
if isinstance(layer, LinearBase):
|
||||
return NPUMXFP8DiffusionLinearMethod(self)
|
||||
return None
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
|
||||
|
||||
class NPUMXFP8DiffusionLinearMethod(LinearMethodBase):
|
||||
"""Ascend NPU MXFP8 linear method for Diffusion models.
|
||||
|
||||
Online mode: loads FP16/BF16 weights → quantises to MXFP8 at load time.
|
||||
Inference: dynamic MXFP8 activation quant + MXFP8 matmul (block_size=32).
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: MXFP8Config):
|
||||
self.quant_config = quant_config
|
||||
|
||||
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")
|
||||
|
||||
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
|
||||
|
||||
# Load weights in original dtype; quantise later in process_weights_after_loading
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
|
||||
weight_fp = layer.weight.data
|
||||
if weight_fp.dtype not in (torch.float16, torch.bfloat16):
|
||||
weight_fp = weight_fp.to(torch.bfloat16)
|
||||
|
||||
# Move weight to NPU if needed. We intentionally use a conditional
|
||||
# move rather than an assert because `dit_cpu_offload` defaults to
|
||||
# True in ServerArgs, which causes fsdp_load to move every parameter
|
||||
# back to CPU after loading (even when the target device is NPU).
|
||||
# npu_dynamic_mx_quant requires an NPU tensor, so we must transfer
|
||||
# here. The quantized fp8 weights produced below will remain on NPU
|
||||
# for inference; if the model still needs to be offloaded after
|
||||
# quantization (e.g. very large model on a small NPU), a higher-level
|
||||
# offload pass can move them back afterwards.
|
||||
if not weight_fp.is_npu:
|
||||
weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")
|
||||
|
||||
# Online MXFP8 quantisation of weights (block_size=32)
|
||||
qw, w_scale = torch_npu.npu_dynamic_mx_quant(
|
||||
weight_fp, dst_type=torch_npu.float8_e4m3fn
|
||||
)
|
||||
layer.weight = Parameter(qw, requires_grad=False)
|
||||
layer.weight_scale_inv = Parameter(w_scale, requires_grad=False)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
original_dtype = x.dtype
|
||||
if original_dtype not in (torch.float16, torch.bfloat16):
|
||||
x = x.to(torch.bfloat16)
|
||||
original_dtype = torch.bfloat16
|
||||
|
||||
# Flatten to 2D [tokens, hidden] so npu_dynamic_mx_quant returns 3D scale
|
||||
input_shape = x.shape
|
||||
x_2d = x.reshape(-1, x.shape[-1])
|
||||
|
||||
# Dynamic MXFP8 activation quantisation
|
||||
qx, input_scale = torch_npu.npu_dynamic_mx_quant(
|
||||
x_2d, dst_type=torch_npu.float8_e4m3fn
|
||||
)
|
||||
|
||||
# MXFP8 matmul
|
||||
output = torch_npu.npu_quant_matmul(
|
||||
qx,
|
||||
layer.weight.transpose(0, 1),
|
||||
layer.weight_scale_inv.transpose(0, 1),
|
||||
scale_dtype=torch_npu.float8_e8m0fnu,
|
||||
pertoken_scale=input_scale,
|
||||
pertoken_scale_dtype=torch_npu.float8_e8m0fnu,
|
||||
bias=bias.to(torch.float32) if bias is not None else None,
|
||||
output_dtype=original_dtype,
|
||||
group_sizes=[1, 1, MXFP8_BLOCK_SIZE],
|
||||
)
|
||||
|
||||
# Restore original shape (replace last dim with output features)
|
||||
output_shape = list(input_shape[:-1]) + [output.shape[-1]]
|
||||
output = output.reshape(output_shape)
|
||||
|
||||
return output
|
||||
@@ -63,8 +63,7 @@ _DTYPE_MISMATCH_EXAMPLE_LIMIT = 3
|
||||
def _is_bitsandbytes_quant_config(quant_config: Any | None) -> bool:
|
||||
if quant_config is None:
|
||||
return False
|
||||
quant_name_getter = getattr(type(quant_config), "get_name", None)
|
||||
return bool(callable(quant_name_getter) and quant_name_getter() == "bitsandbytes")
|
||||
return quant_config.get_name() == "bitsandbytes"
|
||||
|
||||
|
||||
def _format_dtype_mismatch_summary(
|
||||
|
||||
@@ -75,8 +75,7 @@ _HF_SAFETENSORS_URL_RE = re.compile(
|
||||
def _get_quant_config_name(config: Optional[QuantizationConfig]) -> Optional[str]:
|
||||
if config is None:
|
||||
return None
|
||||
quant_name_getter = getattr(type(config), "get_name", None)
|
||||
return quant_name_getter() if callable(quant_name_getter) else None
|
||||
return config.get_name()
|
||||
|
||||
|
||||
def _merge_modelopt_fp4_configs(
|
||||
@@ -176,7 +175,14 @@ class TransformerQuantLoadSpec:
|
||||
|
||||
@property
|
||||
def uses_comfy_layer_markers(self) -> bool:
|
||||
return self.is_comfy_fp8 or self.is_serialized_kitchen_int8
|
||||
return (
|
||||
self.is_comfy_fp8
|
||||
or self.is_serialized_kitchen_int8
|
||||
or (
|
||||
_get_quant_config_name(self.quant_config) == "mxfp8"
|
||||
and self.quant_config.layer_markers is not None
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class _TransformerQuantAdapter:
|
||||
@@ -285,9 +291,7 @@ class _Flux2Nvfp4FallbackAdapter(_TransformerQuantAdapter):
|
||||
if cls_name != "Flux2Transformer2DModel" or quant_config is None:
|
||||
return
|
||||
|
||||
quant_name_getter = getattr(type(quant_config), "get_name", None)
|
||||
quant_name = quant_name_getter() if callable(quant_name_getter) else None
|
||||
if quant_name != "modelopt_fp4":
|
||||
if _get_quant_config_name(quant_config) != "modelopt_fp4":
|
||||
return
|
||||
|
||||
weights_path = os.path.basename(server_args.transformer_weights_path or "")
|
||||
@@ -370,10 +374,7 @@ class _ModelOptFp8OffloadAdapter(_TransformerQuantAdapter):
|
||||
if quant_config is None:
|
||||
return
|
||||
|
||||
quant_name_getter = getattr(type(quant_config), "get_name", None)
|
||||
quant_name = quant_name_getter() if callable(quant_name_getter) else None
|
||||
|
||||
if quant_name != "modelopt_fp8":
|
||||
if _get_quant_config_name(quant_config) != "modelopt_fp8":
|
||||
return
|
||||
|
||||
component_offload = _uses_component_offload(
|
||||
@@ -856,7 +857,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"):
|
||||
if quant_name in ("modelopt_fp8", "comfy_fp8", "mxfp8"):
|
||||
return True
|
||||
if quant_name == "kitchen_int8":
|
||||
assert isinstance(quant_config, KitchenInt8Config)
|
||||
@@ -864,7 +865,6 @@ def _needs_device_weight_postprocess(
|
||||
|
||||
serialized_flag_by_quant_name = {
|
||||
"fp8": "is_checkpoint_fp8_serialized",
|
||||
"mxfp8": "is_checkpoint_fp8_serialized",
|
||||
"mxfp4": "is_checkpoint_mxfp4_serialized",
|
||||
"mxfp4_npu": "is_checkpoint_mxfp4_npu_serialized",
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.comfy_fp8 import ComfyFp8
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import (
|
||||
KitchenInt8Config,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
from sglang.srt.layers.modelopt_utils import canonicalize_modelopt_quant_algo
|
||||
from sglang.srt.model_loader.checkpoint_quantization import (
|
||||
@@ -33,9 +34,13 @@ def inspect_comfy_quant_markers(
|
||||
checkpoint_meta: dict[str, tuple[str, tuple[int, ...]]] = {}
|
||||
raw_markers: dict[str, dict[str, Any]] = {}
|
||||
marked_dtype_weight_prefixes: set[str] = set()
|
||||
global_quant_formats: set[str] = set()
|
||||
|
||||
for path in safetensors_list:
|
||||
with safe_open(path, framework="pt", device="cpu") as checkpoint:
|
||||
metadata = checkpoint.metadata() or {}
|
||||
if quant_format := metadata.get("quant_format"):
|
||||
global_quant_formats.add(quant_format.lower())
|
||||
for key in checkpoint.keys():
|
||||
tensor_slice = checkpoint.get_slice(key)
|
||||
checkpoint_meta[key] = (
|
||||
@@ -67,6 +72,23 @@ def inspect_comfy_quant_markers(
|
||||
)
|
||||
raw_markers[prefix] = marker
|
||||
|
||||
if global_quant_formats == {"mxfp8"}:
|
||||
for prefix in marked_dtype_weight_prefixes:
|
||||
weight_meta = checkpoint_meta[f"{prefix}.weight"]
|
||||
scale_meta = checkpoint_meta.get(f"{prefix}.weight_scale")
|
||||
weight_dtype, weight_shape = weight_meta
|
||||
if (
|
||||
weight_dtype != "F8_E4M3"
|
||||
or len(weight_shape) != 2
|
||||
or weight_shape[1] % 32 != 0
|
||||
or scale_meta != ("U8", (weight_shape[0], weight_shape[1] // 32))
|
||||
):
|
||||
raise ValueError(
|
||||
f"MXFP8 layer {prefix!r} has incompatible weight/scale metadata: "
|
||||
f"{weight_meta} and {scale_meta}"
|
||||
)
|
||||
raw_markers.setdefault(prefix, {"format": "mxfp8"})
|
||||
|
||||
missing_markers = marked_dtype_weight_prefixes - raw_markers.keys()
|
||||
if missing_markers:
|
||||
raise ValueError(
|
||||
@@ -124,6 +146,11 @@ def resolve_comfy_checkpoint_quantization(
|
||||
return KitchenInt8Config(layer_markers=layer_markers)
|
||||
if formats == ["float8_e4m3fn"]:
|
||||
return ComfyFp8Config(layer_markers)
|
||||
if formats == ["mxfp8"]:
|
||||
return MXFP8Config(
|
||||
is_checkpoint_fp8_serialized=True,
|
||||
layer_markers=layer_markers,
|
||||
)
|
||||
raise NotImplementedError(
|
||||
"Unsupported Comfy quantization format(s): " + ", ".join(formats)
|
||||
)
|
||||
|
||||
@@ -10,8 +10,12 @@ from safetensors.torch import safe_open, save_file
|
||||
from torch import nn
|
||||
|
||||
from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
|
||||
BitsAndBytesConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.loader import fsdp_load, rank_local_checkpoint
|
||||
from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config
|
||||
|
||||
|
||||
class _UniformDtypeModel(nn.Module):
|
||||
@@ -45,6 +49,10 @@ class _CustomEntrypointModel(_UniformDtypeModel):
|
||||
|
||||
|
||||
class TestFSDPMixedPrecisionPolicy(unittest.TestCase):
|
||||
def test_quant_config_detection_uses_the_runtime_instance(self):
|
||||
self.assertTrue(fsdp_load._is_bitsandbytes_quant_config(BitsAndBytesConfig()))
|
||||
self.assertFalse(fsdp_load._is_bitsandbytes_quant_config(Fp8Config()))
|
||||
|
||||
def _load_and_capture_policy(
|
||||
self,
|
||||
model_cls: type[nn.Module],
|
||||
|
||||
@@ -69,6 +69,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp8Config,
|
||||
_prepare_nvfp4_weight_bytes,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config
|
||||
from sglang.multimodal_gen.runtime.loader.component_loaders import transformer_loader
|
||||
from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
|
||||
TransformerLoader,
|
||||
@@ -101,6 +102,14 @@ from sglang.multimodal_gen.runtime.utils.quantization_utils import (
|
||||
from sglang.multimodal_gen.tools.build_modelopt_nvfp4_transformer import (
|
||||
_updated_quant_config,
|
||||
)
|
||||
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
|
||||
NPUMXFP8LinearMethod,
|
||||
)
|
||||
from sglang.srt.layers.quantization.bitsandbytes import (
|
||||
BitsAndBytesConfig as SRTBitsAndBytesConfig,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config as SRTFp8Config
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod as SRTFp8LinearMethod
|
||||
|
||||
|
||||
class _FakeFluxTransformer:
|
||||
@@ -407,6 +416,57 @@ class TestTransformerQuantHelpers(unittest.TestCase):
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
|
||||
def test_minimax_h3_global_mxfp8_metadata_selects_srt_per_layer(self):
|
||||
with tempfile.NamedTemporaryFile(suffix=".safetensors") as checkpoint:
|
||||
save_file(
|
||||
{
|
||||
"blocks.0.attn.out_proj.weight": torch.ones(
|
||||
(32, 64), dtype=torch.float8_e4m3fn
|
||||
),
|
||||
"blocks.0.attn.out_proj.weight_scale": torch.ones(
|
||||
(32, 2), dtype=torch.uint8
|
||||
),
|
||||
"token_refiner.blocks.0.attn.out_proj.weight": torch.ones(
|
||||
(32, 64), dtype=torch.bfloat16
|
||||
),
|
||||
},
|
||||
checkpoint.name,
|
||||
metadata={"quant_format": "mxfp8"},
|
||||
)
|
||||
_, layer_markers = inspect_minimax_h3_safetensors([checkpoint.name])
|
||||
config = resolve_minimax_h3_checkpoint_quantization(layer_markers)
|
||||
|
||||
self.assertIsInstance(config, MXFP8Config)
|
||||
self.assertIsInstance(config, SRTFp8Config)
|
||||
layer = LinearBase(input_size=64, output_size=32)
|
||||
self.assertIsInstance(
|
||||
config.get_quant_method(layer, "blocks.0.attn.out_proj"),
|
||||
SRTFp8LinearMethod,
|
||||
)
|
||||
self.assertIsInstance(
|
||||
config.get_quant_method(layer, "token_refiner.blocks.0.attn.out_proj"),
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
|
||||
def test_mxfp8_npu_selects_srt_linear_method(self):
|
||||
with (
|
||||
patch(
|
||||
"sglang.multimodal_gen.runtime.layers.quantization.mxfp8.current_platform.is_mps",
|
||||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"sglang.multimodal_gen.runtime.layers.quantization.mxfp8.current_platform.is_npu",
|
||||
return_value=True,
|
||||
),
|
||||
):
|
||||
config = MXFP8Config()
|
||||
method = config.get_quant_method(
|
||||
LinearBase(input_size=64, output_size=32),
|
||||
"blocks.0.attn.out_proj",
|
||||
)
|
||||
|
||||
self.assertIsInstance(method, NPUMXFP8LinearMethod)
|
||||
|
||||
def test_comfy_full_precision_fp8_dequantizes_before_linear(self):
|
||||
layer = torch.nn.Module()
|
||||
layer.weight = torch.nn.Parameter(
|
||||
@@ -626,7 +686,7 @@ class TestTransformerQuantHelpers(unittest.TestCase):
|
||||
)
|
||||
)
|
||||
self.assertTrue(_needs_device_weight_postprocess(_make_quant_config("mxfp8")))
|
||||
self.assertFalse(
|
||||
self.assertTrue(
|
||||
_needs_device_weight_postprocess(
|
||||
_make_quant_config("mxfp8", is_checkpoint_fp8_serialized=True)
|
||||
)
|
||||
@@ -788,6 +848,7 @@ class TestTransformerQuantHelpers(unittest.TestCase):
|
||||
)
|
||||
|
||||
self.assertEqual(config.get_name(), "bitsandbytes")
|
||||
self.assertIsInstance(config, SRTBitsAndBytesConfig)
|
||||
self.assertTrue(config.load_in_4bit)
|
||||
self.assertEqual(config.bnb_4bit_quant_type, "nf4")
|
||||
|
||||
@@ -805,6 +866,7 @@ class TestTransformerQuantHelpers(unittest.TestCase):
|
||||
)
|
||||
|
||||
self.assertIsInstance(config, Fp8Config)
|
||||
self.assertIsInstance(config, SRTFp8Config)
|
||||
self.assertTrue(config.is_checkpoint_fp8_serialized)
|
||||
|
||||
def test_bitsandbytes_quant_config_resolves_from_compression_config(self):
|
||||
|
||||
@@ -26,6 +26,21 @@ if TYPE_CHECKING:
|
||||
)
|
||||
|
||||
|
||||
def require_bitsandbytes() -> None:
|
||||
try:
|
||||
import bitsandbytes
|
||||
|
||||
if version.parse(bitsandbytes.__version__) < version.parse("0.46.1"):
|
||||
raise ImportError(
|
||||
"bitsandbytes version is wrong. Please install bitsandbytes>=0.46.1."
|
||||
)
|
||||
except ImportError as err:
|
||||
raise ImportError(
|
||||
"Please install bitsandbytes>=0.46.1 via "
|
||||
"`pip install bitsandbytes>=0.46.1` to use bitsandbytes quantizer."
|
||||
) from err
|
||||
|
||||
|
||||
class BitsAndBytesConfig(QuantizationConfig):
|
||||
"""Config class for BitsAndBytes Quantization.
|
||||
|
||||
@@ -184,21 +199,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase):
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: BitsAndBytesConfig):
|
||||
try:
|
||||
import bitsandbytes
|
||||
|
||||
if version.parse(bitsandbytes.__version__) < version.parse("0.46.1"):
|
||||
raise ImportError(
|
||||
"bitsandbytes version is wrong. Please "
|
||||
"install bitsandbytes>=0.46.1."
|
||||
)
|
||||
except ImportError as err:
|
||||
raise ImportError(
|
||||
"Please install bitsandbytes>=0.46.1 via "
|
||||
"`pip install bitsandbytes>=0.46.1` to use "
|
||||
"bitsandbytes quantizer."
|
||||
) from err
|
||||
|
||||
require_bitsandbytes()
|
||||
self.quant_config = quant_config
|
||||
|
||||
def create_weights(
|
||||
|
||||
Reference in New Issue
Block a user