sgl-project--sglang/python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py

# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

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,
    LinearMethodBase,
    UnquantizedLinearMethod,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
    QuantizationConfig,
    QuantizeMethodBase,
)
from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs


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):
    """Config class for pre-quantized bitsandbytes 4-bit checkpoints."""

    def __init__(
        self,
        load_in_8bit: bool = False,
        load_in_4bit: bool = True,
        bnb_4bit_compute_dtype: str = "float32",
        bnb_4bit_quant_storage: str = "uint8",
        bnb_4bit_quant_type: str = "fp4",
        bnb_4bit_use_double_quant: bool = False,
        llm_int8_enable_fp32_cpu_offload: bool = False,
        llm_int8_has_fp16_weight: bool = False,
        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

        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):
                return UnquantizedLinearMethod()
            return BitsAndBytesLinearMethod(self)
        return None


class BitsAndBytesLinearMethod(LinearMethodBase):
    """Linear method for pre-quantized bitsandbytes 4-bit weights."""

    def __init__(self, quant_config: BitsAndBytesConfig):
        _require_bitsandbytes()
        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,
    ) -> None:
        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:
            raise ValueError(
                "The input size is not aligned with the quantized weight shape."
            )

        qweight = nn.Parameter(
            torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8),
            requires_grad=False,
        )
        set_weight_attrs(
            qweight,
            {
                "input_dim": 0,
                "output_dim": 0,
                "pack_factor": quant_ratio,
                "use_bitsandbytes_4bit": True,
                "bnb_full_shape": (output_size, input_size),
                "bnb_local_shape": (
                    output_size_per_partition,
                    input_size_per_partition,
                ),
                "bnb_output_shard_start": getattr(layer, "tp_rank", 0)
                * output_size_per_partition,
                "bnb_input_shard_start": (
                    0
                    if input_size_per_partition == input_size
                    else getattr(layer, "tp_rank", 0) * input_size_per_partition
                ),
            },
        )
        layer.register_parameter("weight", qweight)
        set_weight_attrs(qweight, extra_weight_attrs)

    def apply(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        original_type = x.dtype
        original_shape = x.shape
        if x.ndim > 2:
            x = x.reshape(-1, x.size(-1))

        out_dim = sum(
            quant_state.shape[0]
            for quant_state in layer.weight.bnb_quant_state.values()
        )
        out = torch.empty(x.shape[0], out_dim, dtype=torch.bfloat16, device=x.device)
        apply_bnb_4bit(x.to(torch.bfloat16), layer.weight, out)
        out = out.to(original_type)

        if len(original_shape) > 2:
            out = out.view(*original_shape[:-1], out.size(-1))

        if bias is not None:
            out = out + bias
        return out


def apply_bnb_4bit(
    x: torch.Tensor,
    weight: torch.Tensor,
    out: torch.Tensor,
) -> None:
    from bitsandbytes import matmul_4bit

    offsets = weight.bnb_shard_offsets
    quant_states = weight.bnb_quant_state
    current_index = 0
    for i in range(len(quant_states)):
        output_size = quant_states[i].shape[0]
        out[:, current_index : current_index + output_size] = matmul_4bit(
            x,
            weight[offsets[i] : offsets[i + 1]].t(),
            quant_states[i],
        )
        current_index += output_size


class BitsAndBytes4BitLinear(nn.Module):
    """Storage-only bitsandbytes 4-bit linear for nn.Linear-based encoders."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        compute_dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        _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())
        total_size = in_features * out_features
        if total_size % quant_ratio != 0:
            raise ValueError(
                "The input size is not aligned with the quantized weight shape."
            )

        self.weight = nn.Parameter(
            torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8),
            requires_grad=False,
        )
        set_weight_attrs(
            self.weight,
            {
                "pack_factor": quant_ratio,
                "use_bitsandbytes_4bit": True,
            },
        )
        if bias:
            self.bias = nn.Parameter(
                torch.empty(
                    out_features, dtype=compute_dtype or torch.get_default_dtype()
                ),
                requires_grad=False,
            )
        else:
            self.register_parameter("bias", None)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        original_type = x.dtype
        original_shape = x.shape
        if x.ndim > 2:
            x = x.reshape(-1, x.size(-1))

        out = torch.empty(
            x.shape[0], self.out_features, dtype=torch.bfloat16, device=x.device
        )
        apply_bnb_4bit(x.to(torch.bfloat16), self.weight, out)
        out = out.to(original_type)

        if len(original_shape) > 2:
            out = out.view(*original_shape[:-1], out.size(-1))

        if self.bias is not None:
            out = out + self.bias
        return out


def swap_linears_to_bitsandbytes_4bit(module: nn.Module) -> None:
    for name, child in list(module.named_children()):
        if isinstance(child, nn.Linear):
            replacement = BitsAndBytes4BitLinear(
                child.in_features,
                child.out_features,
                bias=child.bias is not None,
                compute_dtype=child.weight.dtype,
            )
            setattr(module, name, replacement)
        else:
            swap_linears_to_bitsandbytes_4bit(child)


_BNB_4BIT_STATE_SUFFIXES = {
    "absmax",
    "quant_map",
    "nested_absmax",
    "nested_quant_map",
    "bitsandbytes",
}


def is_bitsandbytes_4bit_state_name(weight_name: str) -> bool:
    suffix = weight_name.split(".")[-1]
    return any(state_suffix in suffix for state_suffix in _BNB_4BIT_STATE_SUFFIXES)


def split_bitsandbytes_4bit_state(
    weights: Any,
) -> tuple[list[tuple[str, torch.Tensor]], dict[str, torch.Tensor]]:
    normal_weights: list[tuple[str, torch.Tensor]] = []
    quant_state_dict: dict[str, torch.Tensor] = {}
    for name, tensor in weights:
        if is_bitsandbytes_4bit_state_name(name):
            if "quant_state.bitsandbytes" in name:
                tensor = tensor.cpu().data
            quant_state_dict[name] = tensor
            continue
        normal_weights.append((name, tensor))
    return normal_weights, quant_state_dict


def build_bitsandbytes_4bit_quant_states(
    normal_weight_names: list[str],
    quant_state_dict: dict[str, torch.Tensor],
    device: torch.device,
    param_names_mapping=None,
) -> dict[str, Any]:
    from bitsandbytes.functional import QuantState

    quant_states: dict[str, Any] = {}
    device_str = str(device)
    for source_name in normal_weight_names:
        if (
            f"{source_name}.quant_state.bitsandbytes__nf4" not in quant_state_dict
            and f"{source_name}.quant_state.bitsandbytes__fp4" not in quant_state_dict
        ):
            continue
        target_name = source_name
        if param_names_mapping is not None:
            target_name, _, _ = param_names_mapping(source_name)
        state_tensors = {
            name: tensor
            for name, tensor in quant_state_dict.items()
            if name.startswith(f"{source_name}.")
        }
        quant_states[target_name] = QuantState.from_dict(
            state_tensors, device=device_str
        )
    return quant_states


def attach_bitsandbytes_4bit_quant_states(
    params_dict: dict[str, torch.nn.Parameter],
    quant_states: dict[str, Any],
) -> None:
    for param_name, quant_state in quant_states.items():
        param = params_dict.get(param_name)
        if param is None:
            raise ValueError(f"Parameter {param_name} not found in the model.")

        quant_state = _maybe_shard_bitsandbytes_4bit_quant_state(param, quant_state)
        state_by_shard = {0: quant_state}
        set_weight_attrs(param, {"bnb_quant_state": state_by_shard})
        offsets = torch.tensor([0, param.numel()]).cpu()
        set_weight_attrs(param, {"bnb_shard_offsets": offsets})


def _maybe_shard_bitsandbytes_4bit_quant_state(
    param: torch.nn.Parameter,
    quant_state: Any,
) -> Any:
    full_shape = tuple(getattr(param, "bnb_full_shape", tuple(quant_state.shape or ())))
    local_shape = tuple(getattr(param, "bnb_local_shape", full_shape))
    if not full_shape or local_shape == full_shape:
        return quant_state

    output_start = getattr(param, "bnb_output_shard_start", 0)
    input_start = getattr(param, "bnb_input_shard_start", 0)
    if input_start != 0 or local_shape[1] != full_shape[1]:
        raise NotImplementedError(
            "bitsandbytes 4-bit TP only supports column-parallel output shards."
        )
    if getattr(quant_state, "nested", False):
        raise NotImplementedError(
            "bitsandbytes 4-bit TP does not support nested quant states."
        )

    blocksize = quant_state.blocksize
    start_elem = output_start * full_shape[1]
    local_numel = local_shape[0] * local_shape[1]
    if start_elem % blocksize != 0 or local_numel % blocksize != 0:
        raise ValueError(
            "bitsandbytes 4-bit TP shard is not aligned to quantization blocks."
        )
    start_block = start_elem // blocksize
    num_blocks = local_numel // blocksize
    return type(quant_state)(
        absmax=quant_state.absmax.narrow(0, start_block, num_blocks).contiguous(),
        shape=torch.Size(local_shape),
        code=quant_state.code,
        blocksize=quant_state.blocksize,
        quant_type=quant_state.quant_type,
        dtype=quant_state.dtype,
        offset=None,
        state2=None,
    )