283 lines
11 KiB
Python
283 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
|
||
import enum
|
||
from enum import Enum
|
||
from fractions import Fraction
|
||
from typing import Any, Dict, List, Optional, Union
|
||
import numpy as np
|
||
|
||
import torch
|
||
from torch.nn.parameter import Parameter
|
||
|
||
from vllm import _custom_ops as ops
|
||
from vllm.model_executor.layers.linear import LinearMethodBase
|
||
from vllm.model_executor.layers.quantization.base_config import (
|
||
QuantizationConfig)
|
||
from vllm.model_executor.layers.quantization.utils.gptq_utils import (
|
||
get_linear_quant_method)
|
||
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
|
||
GroupQuantScaleParameter,
|
||
PackedColumnParameter,
|
||
PackedvLLMParameter,
|
||
RowvLLMParameter)
|
||
from vllm.model_executor.layers.quantization.gptq import GPTQConfig as GPTQConfigOrig
|
||
from vllm.model_executor.layers.quantization.gptq import ExllamaState
|
||
from vllm_vacc.vllm.model_executor.models.vars import TRANSPOSE_GPTQ_WEIGHT
|
||
import math
|
||
|
||
def GPTQLinearMethod__init(self, quant_config: GPTQConfigOrig):
|
||
self.quant_config = quant_config
|
||
self.scale_k = 1
|
||
self.split_num = 4
|
||
|
||
def int32_to_int4(s0, axis = -2):
|
||
# 要先拉平 shape[1, n]
|
||
# 每个int32 拆成8个int4, 8个int32表示， 得到[8, n]
|
||
|
||
# x32(int32) => 32bit => 4bit x 8 x4[8] 4bit
|
||
|
||
# x32 31-28 => x4[7]
|
||
# x32 27-24 => x4[6]
|
||
# ...
|
||
# x32 3-0 => x4[0]
|
||
|
||
# x32[index=0] => x4[7,6,5,4,3,2,1,0]
|
||
|
||
# 4bit转真实数字：
|
||
# 不是按补码方式
|
||
|
||
# 1111 => 15 => 7
|
||
# 15-8 = 7
|
||
|
||
# 0101 => 6 =>-2
|
||
# 6-8 = -2
|
||
|
||
# 0x 6A CB 37 2B （内存中排列 2B 37 CB 6A） => B273BCA6 => (-8) => int4: 3, -6, -1, -5, 3, 4, 2, -2
|
||
|
||
# 内存中实际排布为小端模式：
|
||
# int32: 2B 37 CB 6A => 2,11,3,7,12,11,6,10 => (-8) => -6,3, -5,-1, 4,3, -2,2 => 同一字节所在的两个交换得到 3, -6, -1, -5, 3, 4, 2, -2
|
||
# int4: 3, -6, -1, -5, 3, 4, 2, -2
|
||
|
||
s = s0.view(torch.uint32)
|
||
all = []
|
||
for i in range(8):
|
||
x = 15 << (i*4)
|
||
# s2 = torch.bitwise_and(x,s)
|
||
s2 = torch.from_numpy(np.bitwise_and(x, s.numpy()))
|
||
s3 = s2 / (2 ** (i*4))
|
||
s4 = s3.to(torch.int32)
|
||
# 补码， 结果不对
|
||
# s4[s4 > 7] = s4[s4 > 7]-16
|
||
# 直接 - 8 结果正确， 范围： -8-7
|
||
s4 = s4 - 8
|
||
all.append(s4.reshape(1,*s4.shape))
|
||
all = torch.concatenate(all, 0)
|
||
if axis == -2 or axis == 0:
|
||
# 8,K//8,N => K//8,8,N => K,N
|
||
all = all.transpose(-2,0).reshape(-1,all.shape[-1]).contiguous()
|
||
else:
|
||
# 8,N,K//8 => N,K//8,8 => N,K
|
||
all = all.permute(1,2,0).reshape(all.shape[-2],-1).contiguous()
|
||
return all
|
||
|
||
|
||
def dequant_weight(qw, scales, group_size = 128):
|
||
N = qw.shape[1]
|
||
int4_to_int32_axis = -2
|
||
if TRANSPOSE_GPTQ_WEIGHT:
|
||
N = qw.shape[0]
|
||
int4_to_int32_axis = -1
|
||
qweight = int32_to_int4(qw,int4_to_int32_axis).to(torch.float16) #int32 => 8 int4 +> fp16
|
||
|
||
if TRANSPOSE_GPTQ_WEIGHT:
|
||
scales = scales.T.contiguous()
|
||
qweight = qweight.T.contiguous()
|
||
|
||
scales = torch.concatenate([scales] * group_size, 1).reshape(-1, N) # scale 按 group_size 扩展， 每 group_size 个数共用一个scale
|
||
|
||
# print('qweight', qweight.shape, qweight.dtype)
|
||
# print('scale', scales.shape, scales.dtype)
|
||
|
||
dequant_weight = qweight * scales #dequant
|
||
return dequant_weight
|
||
|
||
class GPTQConfig(QuantizationConfig):
|
||
"""Config class for GPTQ.
|
||
|
||
Reference: https://arxiv.org/abs/2210.17323
|
||
"""
|
||
@classmethod
|
||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||
return [torch.half, torch.bfloat16]
|
||
|
||
class GPTQLinearMethod(LinearMethodBase):
|
||
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,
|
||
):
|
||
del output_size # Unused.
|
||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||
# if input_size_per_partition % self.quant_config.group_size != 0:
|
||
# raise ValueError(
|
||
# "The input size is not aligned with the quantized "
|
||
# "weight shape. This can be caused by too large "
|
||
# "tensor parallel size.")
|
||
output_size_per_partition = sum(output_partition_sizes)
|
||
if (output_size_per_partition % self.quant_config.pack_factor.numerator
|
||
!= 0):
|
||
raise ValueError(
|
||
"The output size is not aligned with the quantized "
|
||
"weight shape. This can be caused by too large "
|
||
"tensor parallel size.")
|
||
|
||
if self.quant_config.group_size != -1:
|
||
group_size = self.quant_config.group_size
|
||
else:
|
||
group_size = input_size
|
||
exllama_state = ExllamaState.UNINITIALIZED
|
||
scale_and_zero_size = input_size // group_size
|
||
scale_and_zero_input_dim = None
|
||
if (input_size != input_size_per_partition
|
||
and self.quant_config.group_size != -1):
|
||
# For act-order models, we cannot use Exllama for row parallel layer
|
||
if self.quant_config.desc_act:
|
||
exllama_state = ExllamaState.UNUSED
|
||
else:
|
||
# we need to partition qzeros and scales for exllama kernel
|
||
scale_and_zero_size = input_size_per_partition // group_size
|
||
scale_and_zero_input_dim = 0
|
||
|
||
qweight = PackedvLLMParameter(
|
||
data=torch.empty(
|
||
input_size_per_partition // self.quant_config.pack_factor,
|
||
output_size_per_partition,
|
||
dtype=torch.int32,
|
||
),
|
||
input_dim=0,
|
||
output_dim=1,
|
||
packed_dim=0,
|
||
packed_factor=self.quant_config.pack_factor,
|
||
weight_loader=weight_loader)
|
||
|
||
g_idx = RowvLLMParameter(data=torch.tensor(
|
||
[
|
||
i // self.quant_config.group_size
|
||
for i in range(input_size_per_partition)
|
||
],
|
||
dtype=torch.int32,
|
||
),
|
||
input_dim=0,
|
||
weight_loader=weight_loader)
|
||
qzeros_args = {
|
||
"data":
|
||
torch.empty(
|
||
scale_and_zero_size,
|
||
output_size_per_partition // self.quant_config.pack_factor,
|
||
dtype=torch.int32,
|
||
),
|
||
"weight_loader":
|
||
weight_loader
|
||
}
|
||
weight_scale_args = {
|
||
"data":
|
||
torch.empty(
|
||
scale_and_zero_size,
|
||
output_size_per_partition,
|
||
dtype=params_dtype,
|
||
),
|
||
"weight_loader":
|
||
weight_loader
|
||
}
|
||
if scale_and_zero_input_dim is None:
|
||
scales = ChannelQuantScaleParameter(output_dim=1,
|
||
**weight_scale_args)
|
||
qzeros = PackedColumnParameter(
|
||
output_dim=1,
|
||
packed_dim=1,
|
||
packed_factor=self.quant_config.pack_factor,
|
||
**qzeros_args)
|
||
|
||
else:
|
||
scales = GroupQuantScaleParameter(output_dim=1,
|
||
input_dim=0,
|
||
**weight_scale_args)
|
||
qzeros = PackedvLLMParameter(
|
||
input_dim=0,
|
||
output_dim=1,
|
||
packed_dim=1,
|
||
packed_factor=self.quant_config.pack_factor,
|
||
**qzeros_args)
|
||
|
||
layer.register_parameter("qweight", qweight)
|
||
layer.register_parameter("g_idx", g_idx)
|
||
layer.register_parameter("qzeros", qzeros)
|
||
layer.register_parameter("scales", scales)
|
||
|
||
layer.exllama_state = exllama_state
|
||
|
||
|
||
|
||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||
# for torch.compile
|
||
# self.quant_config.weight_bits == 4
|
||
if TRANSPOSE_GPTQ_WEIGHT:
|
||
layer.qzeros = Parameter(layer.qzeros.data.T.contiguous(), requires_grad=False)
|
||
layer.qweight = Parameter(layer.qweight.data.T.contiguous(), requires_grad=False)
|
||
layer.g_idx = Parameter(layer.g_idx.data, requires_grad=False)
|
||
layer.scales = Parameter(layer.scales.data.T.contiguous(), requires_grad=False)
|
||
else:
|
||
layer.qzeros = Parameter(layer.qzeros.data, requires_grad=False)
|
||
layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
|
||
layer.g_idx = Parameter(layer.g_idx.data, requires_grad=False)
|
||
layer.scales = Parameter(layer.scales.data, requires_grad=False)
|
||
|
||
# exllama needs to shuffle the weight after the weight is loaded
|
||
# here we do the shuffle on first forward pass
|
||
if layer.exllama_state == ExllamaState.UNINITIALIZED:
|
||
if self.quant_config.desc_act:
|
||
layer.g_idx.data = torch.argsort(layer.g_idx).to(torch.int)
|
||
layer.exllama_state = ExllamaState.READY
|
||
ops.gptq_shuffle(layer.qweight, layer.g_idx,
|
||
self.quant_config.weight_bits)
|
||
else:
|
||
layer.g_idx.data = torch.empty((0, ),
|
||
dtype=torch.int,
|
||
device=layer.g_idx.device)
|
||
|
||
def apply(self,
|
||
layer: torch.nn.Module,
|
||
x: torch.Tensor,
|
||
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||
out_shape = x.shape[:-1] + (layer.qweight.shape[-2 if TRANSPOSE_GPTQ_WEIGHT else -1], ) # M,N
|
||
reshaped_x = x.reshape(-1, x.shape[-1])
|
||
|
||
# print(f"~~~~ start dequant")
|
||
# import time
|
||
# start_quant_time = time.time()
|
||
# weight = dequant_weight(layer.qweight.cpu(), layer.scales.cpu(), self.quant_config.group_size // self.scale_k).to(layer.qweight.device)
|
||
# end_quant_time = time.time()
|
||
# print(f"~~~~ dequant time: {end_quant_time - start_quant_time}")
|
||
# if torch.distributed.get_rank() == 0:
|
||
# print(f"~~~~ weight shape: {weight.shape}, dtype: {weight.dtype}")
|
||
# output = torch.matmul(reshaped_x, weight)
|
||
# print("entering GPTQLinearMethod apply, reshaped_x shape:", reshaped_x.shape, "reshaped_x stride", reshaped_x.stride(), "input_tensor", x.shape, "qweight shape:", layer.qweight.shape, "scales shape:", layer.scales.shape)
|
||
output = torch.vacc.w4a8_block_int4_matmul(
|
||
reshaped_x,
|
||
layer.qweight.transpose(-1, -2),
|
||
layer.scales.transpose(-1, -2),
|
||
[1, self.quant_config.group_size // self.scale_k],
|
||
)
|
||
# print("exiting GPTQLinearMethod apply, output shape:", output.shape)
|
||
# end_gemm_time = time.time()
|
||
# if torch.distributed.get_rank() == 0:
|
||
# print(f"~~~~ gemm time: {end_gemm_time - end_quant_time}")
|
||
if bias is not None:
|
||
output.add_(bias)
|
||
return output.reshape(out_shape)
|