utils/quant.py

import torch
import torch.nn as nn
from loguru import logger

try:
    from qtorch.quant import float_quantize
except Exception:
    logger.warning("qtorch not found. Please install qtorch (pip install qtorch).")
    float_quantize = None

try:
    import sgl_kernel
except ImportError:
    sgl_kernel = None


class BaseQuantizer(object):
    def __init__(self, bit, symmetric, granularity, **kwargs):
        self.bit = bit
        self.sym = symmetric
        self.granularity = granularity
        self.kwargs = kwargs
        if self.granularity == "per_group":
            self.group_size = self.kwargs["group_size"]
        self.calib_algo = self.kwargs.get("calib_algo", "minmax")

    def get_tensor_range(self, tensor):
        if self.calib_algo == "minmax":
            return self.get_minmax_range(tensor)
        elif self.calib_algo == "mse":
            return self.get_mse_range(tensor)
        else:
            raise ValueError(f"Unsupported calibration algorithm: {self.calib_algo}")

    def get_minmax_range(self, tensor):
        if self.granularity == "per_tensor":
            max_val = torch.max(tensor)
            min_val = torch.min(tensor)
        else:
            max_val = tensor.amax(dim=-1, keepdim=True)
            min_val = tensor.amin(dim=-1, keepdim=True)
        return (min_val, max_val)

    def get_mse_range(self, tensor):
        raise NotImplementedError

    def get_qparams(self, tensor_range, device):
        min_val, max_val = tensor_range[0], tensor_range[1]
        qmin = self.qmin.to(device)
        qmax = self.qmax.to(device)
        if self.sym:
            abs_max = torch.max(max_val.abs(), min_val.abs())
            abs_max = abs_max.clamp(min=1e-5)
            scales = abs_max / qmax
            zeros = torch.tensor(0.0)
        else:
            scales = (max_val - min_val).clamp(min=1e-5) / (qmax - qmin)
            zeros = (qmin - torch.round(min_val / scales)).clamp(qmin, qmax)
        return scales, zeros, qmax, qmin

    def reshape_tensor(self, tensor, allow_padding=False):
        if self.granularity == "per_group":
            t = tensor.reshape(-1, self.group_size)
        else:
            t = tensor
        return t

    def restore_tensor(self, tensor, shape):
        if tensor.shape == shape:
            t = tensor
        else:
            t = tensor.reshape(shape)
        return t

    def get_tensor_qparams(self, tensor):
        tensor = self.reshape_tensor(tensor)
        tensor_range = self.get_tensor_range(tensor)
        scales, zeros, qmax, qmin = self.get_qparams(tensor_range, tensor.device)
        return tensor, scales, zeros, qmax, qmin

    def fake_quant_tensor(self, tensor):
        org_shape = tensor.shape
        org_dtype = tensor.dtype
        tensor, scales, zeros, qmax, qmin = self.get_tensor_qparams(tensor)
        tensor = self.quant_dequant(tensor, scales, zeros, qmax, qmin)
        tensor = self.restore_tensor(tensor, org_shape).to(org_dtype)
        return tensor

    def real_quant_tensor(self, tensor):
        org_shape = tensor.shape
        tensor, scales, zeros, qmax, qmin = self.get_tensor_qparams(tensor)
        tensor = self.quant(tensor, scales, zeros, qmax, qmin)
        tensor = self.restore_tensor(tensor, org_shape)
        if self.sym:
            zeros = None
        return tensor, scales, zeros


class FloatQuantizer(BaseQuantizer):
    def __init__(self, bit, symmetric, granularity, **kwargs):
        super().__init__(bit, symmetric, granularity, **kwargs)
        assert self.bit in ["e4m3", "e5m2"], f"Unsupported bit configuration: {self.bit}"
        assert self.sym

        if self.bit == "e4m3":
            self.e_bits = 4
            self.m_bits = 3
            self.fp_dtype = torch.float8_e4m3fn
        elif self.bit == "e5m2":
            self.e_bits = 5
            self.m_bits = 2
            self.fp_dtype = torch.float8_e5m2
        else:
            raise ValueError(f"Unsupported bit configuration: {self.bit}")

        finfo = torch.finfo(self.fp_dtype)
        self.qmin, self.qmax = finfo.min, finfo.max

        self.qmax = torch.tensor(self.qmax)
        self.qmin = torch.tensor(self.qmin)

    def quant(self, tensor, scales, zeros, qmax, qmin):
        scaled_tensor = tensor / scales + zeros
        scaled_tensor = torch.clip(scaled_tensor, self.qmin.cuda(), self.qmax.cuda())
        org_dtype = scaled_tensor.dtype
        q_tensor = float_quantize(scaled_tensor.float(), self.e_bits, self.m_bits, rounding="nearest")
        q_tensor.to(org_dtype)
        return q_tensor

    def dequant(self, tensor, scales, zeros):
        tensor = (tensor - zeros) * scales
        return tensor

    def dequant(self, tensor, scales, out_dtype=torch.bfloat16):
        tensor_f = tensor.to(torch.float32)
        scales_f = scales.to(dtype=torch.float32, device=tensor.device)
        out = tensor_f * scales_f
        return out.to(out_dtype)

    def quant_dequant(self, tensor, scales, zeros, qmax, qmin):
        tensor = self.quant(tensor, scales, zeros, qmax, qmin)
        tensor = self.dequant(tensor, scales, zeros)
        return tensor


class SglQuantLinearFp8(nn.Module):
    def __init__(self, myweight, mybias, bias=True, dtype=torch.bfloat16):
        super().__init__()
        w_quantizer = FloatQuantizer("e4m3", True, "per_channel")
        weight, weight_scale, _ = w_quantizer.real_quant_tensor(myweight)
        self.register_buffer("weight", weight.to(torch.float8_e4m3fn))
        self.register_buffer("weight_scale", weight_scale.to(torch.float32))
        if bias:
            self.register_buffer("bias", mybias)
        else:
            self.register_buffer("bias", None)

    def act_quant_func(self, x):
        m, k = x.shape
        input_tensor_quant = torch.empty((m, k), dtype=torch.float8_e4m3fn, device="cuda", requires_grad=False)
        input_tensor_scale = torch.empty((m, 1), dtype=torch.float32, device="cuda", requires_grad=False)
        sgl_kernel.sgl_per_token_quant_fp8(x, input_tensor_quant, input_tensor_scale)
        return input_tensor_quant, input_tensor_scale

    def forward(self, input_tensor):
        input_tensor = input_tensor.squeeze(0)
        shape = (input_tensor.shape[0], self.weight.shape[0])
        dtype = input_tensor.dtype
        device = input_tensor.device
        output_tensor = torch.empty(shape, dtype=dtype, device=device, requires_grad=False)
        input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor)
        output_tensor = sgl_kernel.fp8_scaled_mm(
            input_tensor_quant,
            self.weight.t(),
            input_tensor_scale,
            self.weight_scale.float(),
            dtype,
            bias=self.bias,
        )

        return output_tensor.unsqueeze(0)

    def _apply(self, fn):
        for module in self.children():
            module._apply(fn)

        def maybe_cast(t):
            if t is not None and t.device != fn(t).device:
                return fn(t)
            return t

        self.weight = maybe_cast(self.weight)
        self.weight_scale = maybe_cast(self.weight_scale)
        self.bias = maybe_cast(self.bias)
        return self


def replace_blocks_linear_only(
    model: nn.Module,
    group_size: int = 16,
    verbose: bool = True,
) -> nn.Module:
    replaced_count = 0

    for block_idx, block in enumerate(model.blocks):
        if verbose:
            print(f"\nProcessing block {block_idx}:")
        
        if hasattr(block, 'self_attn'):
            self_attn = block.self_attn
            for attr_name in ['q', 'k', 'v', 'o']:
                if hasattr(self_attn, attr_name):
                    linear = getattr(self_attn, attr_name)
                    if isinstance(linear, nn.Linear):
                        print(f"  Replacing self_attn.{attr_name}")
                        quant_linear = SglQuantLinearFp8(linear.weight, linear.bias)
                        setattr(self_attn, attr_name, quant_linear)
                        replaced_count += 1
        
        if hasattr(block, 'cross_attn'):
            cross_attn = block.cross_attn
            for attr_name in ['q', 'k', 'v', 'o']:
                if hasattr(cross_attn, attr_name):
                    linear = getattr(cross_attn, attr_name)
                    if isinstance(linear, nn.Linear):
                        print(f"  Replacing cross_attn.{attr_name}")
                        quant_linear = SglQuantLinearFp8(linear.weight, linear.bias)
                        setattr(cross_attn, attr_name, quant_linear)
                        replaced_count += 1
        
        if hasattr(block, 'ffn'):
            ffn = block.ffn
            linear_0 = ffn[0]
            if isinstance(linear_0, nn.Linear):
                print(f"  Replacing ffn[0]")
                quant_linear = SglQuantLinearFp8(linear_0.weight, linear_0.bias)
                ffn[0] = quant_linear
                replaced_count += 1
            linear_2 = ffn[2]
            if isinstance(linear_2, nn.Linear):
                print(f"  Replacing ffn[2]")
                quant_linear = SglQuantLinearFp8(linear_2.weight, linear_2.bias)
                ffn[2] = quant_linear
                replaced_count += 1