wehub-resource-sync
405 lines
15 KiB
Python
405 lines
15 KiB
Python
"""The FasterTransformer quantization config"""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Callable, List, Literal, Optional, Tuple # noqa: UP035
|
|
|
|
import tvm
|
|
from tvm import DataType, DataTypeCode, IRModule, relax, te, tirx
|
|
from tvm.relax.frontend import nn
|
|
from tvm.runtime import Tensor
|
|
from tvm.s_tir import dlight as dl
|
|
from tvm.target import Target
|
|
|
|
from ..loader import QuantizeMapping
|
|
from ..op import faster_transformer_dequantize_gemm
|
|
from ..support import logging
|
|
from ..support.auto_target import detect_cuda_arch_list
|
|
from ..support.style import bold
|
|
from .group_quantization import (
|
|
GroupQuantize,
|
|
GroupQuantizeEmbedding,
|
|
GroupQuantizeLinear,
|
|
)
|
|
from .utils import is_final_fc, is_moe_gate
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class FTQuantize:
|
|
"""Configuration for FasterTransformer quantization"""
|
|
|
|
name: str
|
|
kind: str
|
|
quantize_dtype: Literal["int4", "int8"]
|
|
storage_dtype: Literal["int8"]
|
|
model_dtype: Literal["float16"]
|
|
group_size: Optional[int] = None
|
|
|
|
num_elem_per_storage: int = 0
|
|
max_int_value: int = 0
|
|
|
|
def fallback_group_quantize(self) -> GroupQuantize:
|
|
"""
|
|
The fallback group quantization config for other parameters.
|
|
|
|
Returns
|
|
------
|
|
quantize: GroupQuantize
|
|
The group quantization config to fallback.
|
|
"""
|
|
return GroupQuantize(
|
|
name=self.name,
|
|
kind="group-quant",
|
|
group_size=32, # hardcoded to 32 as only supporting int4 quantization
|
|
quantize_dtype=self.quantize_dtype,
|
|
storage_dtype="uint32",
|
|
model_dtype=self.model_dtype,
|
|
linear_weight_layout="NK",
|
|
)
|
|
|
|
def __post_init__(self):
|
|
assert self.kind == "ft-quant"
|
|
quantize_dtype = DataType(self.quantize_dtype)
|
|
storage_dtype = DataType(self.storage_dtype)
|
|
assert self.quantize_dtype in ["int4", "int8"]
|
|
assert storage_dtype.type_code == DataTypeCode.INT
|
|
assert self.model_dtype == "float16"
|
|
assert self.group_size in [None, 64, 128]
|
|
if storage_dtype.bits < quantize_dtype.bits:
|
|
raise ValueError("Storage unit should be greater or equal to quantized element")
|
|
|
|
self.num_elem_per_storage = storage_dtype.bits // quantize_dtype.bits
|
|
self.max_int_value = (2 ** (quantize_dtype.bits - 1)) - 1
|
|
self._quantize_func_cache = {}
|
|
|
|
def quantize_model(
|
|
self,
|
|
model: nn.Module,
|
|
quant_map: QuantizeMapping,
|
|
name_prefix: str,
|
|
) -> nn.Module:
|
|
"""
|
|
Quantize model with FasterTransformer quantization
|
|
|
|
Parameters
|
|
----------
|
|
model : nn.Module
|
|
The non-quantized nn.Module.
|
|
|
|
quant_map : QuantizeMapping
|
|
The quantize mapping with name mapping and func mapping.
|
|
|
|
name_prefix : str
|
|
The name prefix for visited weight.
|
|
|
|
Returns
|
|
-------
|
|
ret : nn.Module
|
|
The quantized nn.Module.
|
|
"""
|
|
|
|
class _Mutator(nn.Mutator):
|
|
def __init__(self, config: FTQuantize, quant_map: QuantizeMapping) -> None:
|
|
super().__init__()
|
|
self.config = config
|
|
self.quant_map = quant_map
|
|
|
|
def visit_module(self, name: str, node: nn.Module) -> Any:
|
|
"""
|
|
The visiting method for FasterTransformer quantization of nn.Module nodes.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the current node.
|
|
|
|
node : nn.Module
|
|
The current node of nn.Module to mutate.
|
|
|
|
Returns
|
|
------
|
|
ret_node: Any
|
|
The new node to replace current node.
|
|
"""
|
|
if isinstance(node, nn.Linear):
|
|
weight_name = f"{name}.weight"
|
|
self.quant_map.param_map[weight_name] = [
|
|
f"{name}.q_weight",
|
|
f"{name}.q_scale",
|
|
]
|
|
if (
|
|
is_final_fc(name)
|
|
or node.out_dtype == "float32"
|
|
or (self.config.quantize_dtype == "int4" and node.out_features % 8 != 0)
|
|
or (self.config.quantize_dtype == "int8" and node.out_features % 4 != 0)
|
|
):
|
|
# Under any of the conditions we fall back to GroupQuantize
|
|
# For `is_final_fc()` see https://github.com/mlc-ai/mlc-llm/issues/1723
|
|
# If simply skipping lm_head quantization degrades performance
|
|
# Other requirements are from CUTLASS
|
|
logger.info(
|
|
'Fallback to GroupQuantize for nn.Linear: "%s", '
|
|
+ "weight.shape: %s, out_dtype: %s",
|
|
bold(name),
|
|
node.weight.shape,
|
|
node.out_dtype,
|
|
)
|
|
group_quantize = self.config.fallback_group_quantize()
|
|
self.quant_map.map_func[weight_name] = group_quantize.quantize_weight
|
|
return GroupQuantizeLinear.from_linear(node, group_quantize)
|
|
if not is_moe_gate(name, node):
|
|
self.quant_map.map_func[weight_name] = self.config.quantize_weight
|
|
return FTQuantizeLinear.from_linear(node, self.config)
|
|
if isinstance(node, nn.Embedding):
|
|
weight_name = f"{name}.weight"
|
|
self.quant_map.param_map[weight_name] = [
|
|
f"{name}.q_weight",
|
|
f"{name}.q_scale",
|
|
]
|
|
group_quantize = self.config.fallback_group_quantize()
|
|
self.quant_map.map_func[weight_name] = group_quantize.quantize_weight
|
|
return GroupQuantizeEmbedding.from_embedding(node, group_quantize)
|
|
return self.visit(name, node)
|
|
|
|
model.to(dtype=self.model_dtype)
|
|
mutator = _Mutator(self, quant_map)
|
|
model = mutator.visit(name_prefix, model)
|
|
return model
|
|
|
|
def quantize_weight(self, weight: Tensor) -> List[Tensor]: # noqa: UP006
|
|
"""
|
|
Quantize weight with FasterTransformer quantization
|
|
|
|
Parameters
|
|
----------
|
|
weight : Tensor
|
|
The original weight.
|
|
|
|
Returns
|
|
------
|
|
ret: List[Tensor]
|
|
The list of FasterTransformer quantized weights.
|
|
"""
|
|
assert tvm.get_global_func("relax.ext.cutlass", True), (
|
|
"Cutlass should be enabled in TVM runtime to quantize weight, "
|
|
"but not enabled in current TVM runtime environment. "
|
|
"To enable Cutlass in TVM runtime, please `set(USE_CUTLASS ON)` "
|
|
"in config.cmake when compiling TVM from source"
|
|
)
|
|
assert len(weight.shape) == 2
|
|
device = weight.device
|
|
device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
|
|
if device_type == "cuda":
|
|
target = Target.current()
|
|
if target is None:
|
|
target = Target.from_device(device)
|
|
with target:
|
|
|
|
def _create_quantize_func() -> IRModule:
|
|
bb = relax.BlockBuilder()
|
|
weight_var = relax.Var("weight", relax.TensorType(weight.shape, weight.dtype))
|
|
with bb.function(name="main", params=[weight_var]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit_te(self._quantize, weight_var)
|
|
lv1 = bb.normalize(lv0[0])
|
|
lv2 = bb.emit(
|
|
relax.call_pure_packed(
|
|
"cutlass.ft_preprocess_weight",
|
|
lv1,
|
|
detect_cuda_arch_list(target=target)[0],
|
|
DataType(self.quantize_dtype).bits == 4,
|
|
ty_args=lv1.ty,
|
|
)
|
|
)
|
|
gv = bb.emit_output(relax.Tuple([lv2, lv0[1]]))
|
|
bb.emit_func_output(gv)
|
|
return bb.finalize()
|
|
|
|
def _compile_quantize_func(mod: IRModule) -> Callable:
|
|
mod = dl.ApplyDefaultSchedule(
|
|
dl.gpu.Reduction(),
|
|
dl.gpu.GeneralReduction(),
|
|
dl.gpu.Fallback(),
|
|
)(mod)
|
|
ex = relax.build(mod, target=target)
|
|
vm = relax.VirtualMachine(ex, device)
|
|
return vm["main"]
|
|
|
|
key = str(
|
|
(
|
|
int(weight.shape[0]),
|
|
int(weight.shape[1]),
|
|
weight.dtype,
|
|
device_type,
|
|
)
|
|
)
|
|
quantize_func = self._quantize_func_cache.get(key, None)
|
|
if quantize_func is None:
|
|
logger.info("Compiling quantize function for key: %s", key)
|
|
quantize_func = _compile_quantize_func(_create_quantize_func())
|
|
self._quantize_func_cache[key] = quantize_func
|
|
data = quantize_func(weight)
|
|
return data
|
|
else:
|
|
raise NotImplementedError(f"Device type {device_type} is not supported")
|
|
|
|
def _quantize(
|
|
self,
|
|
weight: te.Tensor,
|
|
) -> Tuple[te.Tensor, te.Tensor]: # noqa: UP006
|
|
"""FasterTransformer quantization for weight tensor, defined in tensor expression."""
|
|
assert len(weight.shape) == 2
|
|
n, k = weight.shape
|
|
|
|
cur_group_size = k if not self.group_size else self.group_size
|
|
scale_shape = (tirx.ceildiv(k, cur_group_size), n)
|
|
r = te.reduce_axis((0, cur_group_size), name="r")
|
|
|
|
max_abs = te.compute(
|
|
shape=scale_shape,
|
|
fcompute=lambda j, i: te.max(
|
|
tirx.if_then_else(
|
|
j * cur_group_size + r < k,
|
|
te.abs(weight[i, j * cur_group_size + r]),
|
|
te.min_value(self.model_dtype),
|
|
),
|
|
axis=r,
|
|
),
|
|
name="max_abs_value",
|
|
)
|
|
max_int = tirx.const(self.max_int_value, self.model_dtype)
|
|
scale = te.compute(
|
|
scale_shape,
|
|
lambda i, j: max_abs[i, j].astype(self.model_dtype) / max_int,
|
|
name="scale",
|
|
)
|
|
# compute scaled weight
|
|
quantize_dtype = DataType(self.quantize_dtype)
|
|
bin_mask = tirx.const((1 << quantize_dtype.bits) - 1, self.storage_dtype)
|
|
scaled_weight = te.compute(
|
|
shape=weight.shape,
|
|
fcompute=lambda i, j: (
|
|
tirx.min(
|
|
tirx.max(
|
|
tirx.round(weight[i, j] / scale[j // cur_group_size, i]),
|
|
-max_int - 1,
|
|
),
|
|
max_int,
|
|
).astype(self.storage_dtype)
|
|
& bin_mask
|
|
),
|
|
)
|
|
|
|
quantized_weight_shape = (k, tirx.ceildiv(n, self.num_elem_per_storage))
|
|
r = te.reduce_axis((0, self.num_elem_per_storage), name="r")
|
|
quantized_weight = te.compute(
|
|
shape=quantized_weight_shape,
|
|
fcompute=lambda j, i: tirx.sum(
|
|
tirx.if_then_else(
|
|
i * self.num_elem_per_storage + r < n,
|
|
scaled_weight[i * self.num_elem_per_storage + r, j]
|
|
<< (
|
|
r.astype(self.storage_dtype)
|
|
* tirx.const(quantize_dtype.bits, self.storage_dtype)
|
|
),
|
|
tirx.const(0, self.storage_dtype),
|
|
),
|
|
axis=r,
|
|
),
|
|
name="weight",
|
|
)
|
|
|
|
return quantized_weight, scale
|
|
|
|
|
|
class FTQuantizeLinear(nn.Module):
|
|
"""An nn.Linear module with FasterTransformer quantization"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_features: int,
|
|
out_features: int,
|
|
config: FTQuantize,
|
|
bias: bool = True,
|
|
out_dtype: Optional[str] = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.out_dtype = out_dtype
|
|
self.config = config
|
|
cur_group_size = in_features if not config.group_size else config.group_size
|
|
self.q_weight = nn.Parameter(
|
|
(in_features, tirx.ceildiv(out_features, config.num_elem_per_storage)),
|
|
config.storage_dtype,
|
|
)
|
|
self.q_scale = nn.Parameter(
|
|
(tirx.ceildiv(in_features, cur_group_size), out_features), config.model_dtype
|
|
)
|
|
if bias:
|
|
self.bias = nn.Parameter(
|
|
(out_features,), config.model_dtype if out_dtype is None else out_dtype
|
|
)
|
|
else:
|
|
self.bias = None
|
|
|
|
@staticmethod
|
|
def from_linear(src: nn.Linear, config: FTQuantize) -> "FTQuantizeLinear":
|
|
"""
|
|
Converts a non-quantized nn.Linear to a FasterTransformer quantized FTQuantizeLinear
|
|
|
|
Parameters
|
|
----------
|
|
src : nn.Linear
|
|
The non-quantized nn.Linear.
|
|
|
|
config : FTQuantize
|
|
The FasterTransformer quantization config.
|
|
|
|
Returns
|
|
-------
|
|
ret : FTQuantizeLinear
|
|
The FasterTransformer quantized FTQuantizeLinear layer.
|
|
"""
|
|
quantized_linear = FTQuantizeLinear(
|
|
in_features=src.in_features,
|
|
out_features=src.out_features,
|
|
config=config,
|
|
bias=getattr(src, "bias", None) is not None,
|
|
out_dtype=src.out_dtype,
|
|
)
|
|
if quantized_linear.bias is not None:
|
|
quantized_linear.bias.attrs = src.bias.attrs
|
|
return quantized_linear
|
|
|
|
def forward(self, x: nn.Tensor) -> nn.Tensor:
|
|
"""
|
|
Forward method for FasterTransformer quantized linear layer.
|
|
|
|
Parameters
|
|
----------
|
|
x : nn.Tensor
|
|
The input tensor.
|
|
|
|
Returns
|
|
-------
|
|
ret : nn.Tensor
|
|
The output tensor for the FasterTransformer quantized linear layer.
|
|
"""
|
|
return faster_transformer_dequantize_gemm(
|
|
x, self.q_weight, self.q_scale, self.bias, group_size=self.config.group_size
|
|
)
|
|
|
|
def to(self, dtype: Optional[str] = None) -> None:
|
|
"""
|
|
Override to() such that we do not convert bias if there is an out_dtype.
|
|
Otherwise, we might run into dtype mismatch when computing x + self.bias.
|
|
"""
|
|
self.q_weight.to(dtype=dtype)
|
|
self.q_scale.to(dtype=dtype)
|
|
if self.bias is not None and self.out_dtype is None:
|
|
self.bias.to(dtype=dtype)
|
|
if dtype is not None and isinstance(getattr(self, "dtype", None), str):
|
|
self.dtype = dtype
|