tinynn/llm_quant/modules.py

import sys

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from transformers.models.llama.modeling_llama import LlamaAttention
from tinynn.llm_quant.llama import LlamaAttentionFused
from tinynn.util.util import get_logger

from .util import _init_patch_easyquant, get_submodule_with_parent_from_name

log = get_logger(__name__, 'INFO')
SPEEDUP = True

try:
    if sys.platform == "win32":
        _init_patch_easyquant()

    from easyquant import (
        decompress_int4,
        decompress_int8,
        quantize_per_token,
        gemm,
        dequantize_bias_per_token,
        dequantize_per_token,
    )
except (ImportError, OSError):
    log.warning('easyquant is not installed, the inference performance may be degraded')
    SPEEDUP = False


def compress_int(data_tensor, bit_width, per_channel=True, per_token=False):
    # use [-127, 127] as 8-bit quant range
    q_max = 2 ** (bit_width - 1) - 1
    q_min = -q_max

    assert (per_channel and per_token) is False
    if per_channel:
        # for weight, use w_max/quant_max as scale, and convert weight to int8 to save memory.
        scale = 2 * (data_tensor.abs().max(dim=-1).values.float() / (2**bit_width - 1))
        quantized_tensor = torch.clamp(torch.round(data_tensor.float() / scale[:, None]), q_min, q_max).to(torch.int8)
    elif per_token:
        # per-token quantization
        scales = data_tensor.abs().max(dim=-1).values.float() / q_max
        if len(data_tensor.shape) == 3:
            scales = scales[:, :, None]
        elif len(data_tensor.shape) == 2:
            scales = scales[:, None]
        else:
            assert False
        quantized_tensor = torch.clamp(torch.round(data_tensor.float() / scales.float()), q_min, q_max).to(torch.int8)
        scale = scales
    else:
        # per_tensor quantization
        scale = data_tensor.abs().max().float() / q_max
        quantized_tensor = torch.clamp(torch.round(data_tensor.float() / scale.float()), q_min, q_max).to(torch.int8)

    return scale, quantized_tensor


class QLinear(nn.Module):
    def __init__(self, fc: nn.Linear, quant_mode: str):
        super().__init__()
        assert quant_mode in ("weight4", "weight8", "dynamic")
        if quant_mode == 'weight4':
            weight_bit_width = 4
        else:
            weight_bit_width = 8

        self.weight_bit_width = weight_bit_width
        self.quant_mod = quant_mode
        self.in_features = fc.in_features
        self.out_features = fc.out_features

        bias = None if fc.bias is None else fc.bias.data
        # compress weight by given bit, use per-channel and [-127,127]/[-7,7] to clamp
        scale, weight_q = compress_int(fc.weight.data, weight_bit_width)
        if self.in_features % 4 != 0 and quant_mode == 'dynamic':
            weight_q = F.pad(weight_q, (0, 4 - self.in_features % 4))

        if self.weight_bit_width == 4:
            weight_shape = weight_q.shape
            assert len(weight_shape) == 2
            assert weight_shape[1] % 2 == 0
            pre_packed = weight_q.view(weight_shape[0], weight_shape[1] // 2, 2)
            weight_q = ((pre_packed[..., 0] & 0b00001111) << 4) | (pre_packed[..., 1] & 0b00001111)

        self.weight = nn.Parameter(weight_q, requires_grad=False)
        self.weight_scale = nn.Parameter(scale, requires_grad=False)
        self.bias = nn.Parameter(bias, requires_grad=False) if bias is not None else None

        fc.weight = None
        fc.bias = None

    def forward(self, input: Tensor) -> Tensor:
        input_device = input.device
        input_dtype = input.dtype
        input_shape = input.shape
        if self.quant_mod == 'static':
            assert False, f'{self.quant_mod} not supported'
        else:
            if self.quant_mod == 'weight4':
                if SPEEDUP:
                    weight_fp = torch.empty(
                        (self.out_features, self.in_features), dtype=torch.float16, device=input.device
                    )
                    decompress_int4(weight_fp, self.weight, self.weight_scale)
                else:
                    weight_fp = (
                        torch.stack((self.weight >> 4, self.weight << 4 >> 4), -1)
                        .view(self.out_features, self.in_features)
                        .to(dtype=torch.float32)
                        * self.weight_scale[:, None]
                    ).to(dtype=torch.half)
            elif self.quant_mod == 'weight8':
                if SPEEDUP:
                    weight_fp = torch.empty_like(self.weight.data, dtype=input_dtype, device=input_device)
                    decompress_int8(weight_fp, self.weight, self.weight_scale)
                else:
                    weight_fp = (self.weight.to(dtype=torch.float32) * self.weight_scale[:, None]).to(dtype=torch.half)

            if 'dynamic' in self.quant_mod:
                if SPEEDUP:
                    # the real dynamic quantization process, first quantize input to int8, then do int8Gemm calculation,
                    # and finally dequantize the output to float
                    input_viewed = input.view(-1, input_shape[-1])

                    # pad self.weight to 4x
                    padding_num = 4 - self.in_features % 4 if self.in_features % 4 != 0 else 0

                    # init easyquant kernels' output
                    input_q = torch.empty(
                        (input_viewed.shape[0], input_viewed.shape[1] + padding_num),
                        dtype=torch.int8,
                        device=input_device,
                    )
                    scale_shape = input_viewed.shape[0] if 'token' in self.quant_mod else 1
                    input_scale = torch.zeros(scale_shape, device=input_device)
                    out_q = torch.empty(
                        (int(input_viewed.shape[0]), self.out_features), dtype=torch.int32, device=input_device
                    )
                    output = torch.empty_like(out_q, dtype=torch.float16, device=input_device)

                    # use easyquant kernels to accelerate computation
                    quantize_per_token(input_q, input_viewed, input_scale)
                    gemm(out_q, input_q, self.weight)

                    if self.bias is not None:
                        dequantize_bias_per_token(output, out_q, input_scale, self.weight_scale, self.bias)
                    else:
                        dequantize_per_token(output, out_q, input_scale, self.weight_scale)

                    output = output.view(input_shape[:-1] + (output.shape[-1],))
                else:
                    # simulate quantization
                    input_scale, input_q = compress_int(input, 8, per_channel=False, per_token=True)
                    if self.in_features % 4 != 0:
                        output = F.linear(
                            input_q.float(), self.weight[:, : self.in_features % 4 - 4].float(), self.bias
                        )
                    else:
                        output = F.linear(input_q.float(), self.weight.float(), self.bias)
                    output = (output.float() * (self.weight_scale * input_scale.view(-1, 1))).half()
            else:
                input_fq = input
                output = F.linear(input_fq, weight_fp, self.bias)

        return output


class TDQLinear_noinit(QLinear):
    def forward(self, input: Tensor) -> Tensor:
        input_shape = input.shape
        bs, seq, _ = input_shape
        input_device = input.device
        input_viewed = input.view(-1, self.in_features)
        # pad self.weight to 4x
        padding_num = 4 - self.in_features % 4 if self.in_features % 4 != 0 else 0

        input_q = torch.empty(
            (input_viewed.shape[0], self.in_features + padding_num), dtype=torch.int8, device=input_device
        )
        input_scale = torch.empty(bs * seq, device=input_device)
        out_q = torch.empty((bs * seq, self.out_features), dtype=torch.int32, device=input_device)
        output = torch.empty_like(out_q, dtype=torch.float16, device=input_device)

        quantize_per_token(input_q, input_viewed, input_scale)
        gemm(out_q, input_q, self.weight)
        dequantize_per_token(output, out_q, input_scale, self.weight_scale)

        output = output.view(input_shape[:-1] + (output.shape[-1],))
        return output


@torch.no_grad()
def fuse_atten(model: nn.Module):
    """fuse qkv linear, fuse scaled_dot_product_attention if torch>=1.13"""
    for name, mod in model.named_modules():
        if isinstance(mod, LlamaAttention):
            _, parent_mod, last_name = get_submodule_with_parent_from_name(model, name)
            fused_attn = LlamaAttentionFused(mod)
            setattr(parent_mod, last_name, fused_attn)


@torch.no_grad()
def quant_fc(model: nn.Module, quant_mod='weight8', fuse_qkv=False):
    """convert all fcs of LLM model to quantized linear inplace.

    Args:
        model: the Given LLM model.
        quant_mod: the working quantization mode. Default to be 'weight8', Optional:['weight4', 'dynamic_token'].
            The 'dynamic_token' quantization use easyquant lib to do Int8Gemm accelerate.
        fuse_qkv: whether to fuse qkv linear of attention to speedup inference,
            the scaled-dot-product-attention will be fusedif the PyTorch version >= 1.13.
    """
    model.cpu()
    log.info(f'use quant mod {quant_mod} speedup={SPEEDUP}')
    if fuse_qkv:
        fuse_atten(model)
        log.info('qkv has been fused')

    for name, mod in model.named_modules():
        if 'lm_head' in name:
            continue
        if isinstance(mod, nn.Linear):
            _, parent_mod, last_name = get_submodule_with_parent_from_name(model, name)
            if quant_mod == 'dynamic' and SPEEDUP:
                quantized_fc_cls = TDQLinear_noinit
            else:
                quantized_fc_cls = QLinear
            quantized_fc = quantized_fc_cls(
                mod,
                quant_mod,
            )
            setattr(parent_mod, last_name, quantized_fc)