chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,660 @@
|
||||
"""The group quantization config"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from typing import Any, List, Literal, Optional, Tuple, Union # noqa: UP035
|
||||
|
||||
from tvm import DataType, DataTypeCode, IRModule, relax, te, tirx, topi
|
||||
from tvm.relax.frontend import nn
|
||||
from tvm.runtime import Tensor
|
||||
|
||||
from mlc_llm.loader import QuantizeMapping
|
||||
from mlc_llm.nn import MixtralExperts
|
||||
from mlc_llm.support import logging
|
||||
|
||||
from .utils import (
|
||||
apply_sharding,
|
||||
compile_quantize_func,
|
||||
convert_uint_to_float,
|
||||
is_final_fc,
|
||||
is_moe_gate,
|
||||
pack_weight,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GroupQuantize:
|
||||
"""Configuration for group quantization"""
|
||||
|
||||
name: str
|
||||
kind: str
|
||||
group_size: int
|
||||
quantize_dtype: Literal["int3", "int4", "int8"]
|
||||
storage_dtype: Literal["uint32"]
|
||||
model_dtype: Literal["float16", "float32", "bfloat16"]
|
||||
linear_weight_layout: Literal["KN", "NK"]
|
||||
quantize_embedding: bool = True
|
||||
quantize_final_fc: bool = True
|
||||
|
||||
num_elem_per_storage: int = 0
|
||||
num_storage_per_group: int = 0
|
||||
max_int_value: int = 0
|
||||
tensor_parallel_shards: int = 0
|
||||
|
||||
def __post_init__(self):
|
||||
assert self.kind == "group-quant"
|
||||
quantize_dtype = DataType(self.quantize_dtype)
|
||||
storage_dtype = DataType(self.storage_dtype)
|
||||
model_dtype = DataType(self.model_dtype)
|
||||
assert quantize_dtype.type_code == DataTypeCode.INT
|
||||
assert storage_dtype.type_code == DataTypeCode.UINT
|
||||
assert model_dtype.type_code in (DataTypeCode.FLOAT, DataTypeCode.BFLOAT)
|
||||
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
|
||||
if self.group_size % self.num_elem_per_storage != 0:
|
||||
raise ValueError("Group size should be divisible by numbers of elements per storage")
|
||||
self.num_storage_per_group = self.group_size // self.num_elem_per_storage
|
||||
self.max_int_value = (2 ** (quantize_dtype.bits - 1)) - 1
|
||||
self.linear_quant_axis = 0 if self.linear_weight_layout == "KN" else 1
|
||||
self._quantize_func_cache = {}
|
||||
|
||||
def quantize_model(
|
||||
self,
|
||||
model: nn.Module,
|
||||
quant_map: QuantizeMapping,
|
||||
name_prefix: str,
|
||||
) -> nn.Module:
|
||||
"""
|
||||
Quantize model with group 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: GroupQuantize, 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 group 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 getattr(node, "no_quantization", False):
|
||||
return node
|
||||
|
||||
if (
|
||||
isinstance(node, nn.Linear)
|
||||
and (not is_final_fc(name) or self.config.quantize_final_fc)
|
||||
and not is_moe_gate(name, node)
|
||||
):
|
||||
weight_name = f"{name}.weight"
|
||||
self.quant_map.param_map[weight_name] = [
|
||||
f"{name}.q_weight",
|
||||
f"{name}.q_scale",
|
||||
]
|
||||
self.quant_map.map_func[weight_name] = partial(
|
||||
self.config.quantize_weight,
|
||||
output_transpose=self.config.linear_weight_layout == "KN",
|
||||
)
|
||||
return GroupQuantizeLinear.from_linear(node, self.config)
|
||||
if isinstance(node, nn.Embedding) and self.config.quantize_embedding:
|
||||
weight_name = f"{name}.weight"
|
||||
self.quant_map.param_map[weight_name] = [
|
||||
f"{name}.q_weight",
|
||||
f"{name}.q_scale",
|
||||
]
|
||||
self.quant_map.map_func[weight_name] = self.config.quantize_weight
|
||||
return GroupQuantizeEmbedding.from_embedding(node, self.config)
|
||||
if isinstance(node, MixtralExperts):
|
||||
weight_name = f"{name}.weight"
|
||||
self.quant_map.param_map[weight_name] = [
|
||||
f"{name}.q_weight",
|
||||
f"{name}.q_scale",
|
||||
]
|
||||
self.quant_map.map_func[weight_name] = self.config.quantize_weight
|
||||
return GroupQuantizeMixtralExperts.from_mixtral_experts(node, self.config)
|
||||
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 _dequantize(
|
||||
self,
|
||||
weight: te.Tensor,
|
||||
scale: te.Tensor,
|
||||
axis: int,
|
||||
out_shape: Optional[List[tirx.Expr]] = None, # noqa: UP006
|
||||
):
|
||||
tir_max_int = tirx.const(self.max_int_value, self.model_dtype)
|
||||
float_weight = convert_uint_to_float(
|
||||
weight,
|
||||
DataType(self.quantize_dtype).bits,
|
||||
self.num_elem_per_storage,
|
||||
self.storage_dtype,
|
||||
self.model_dtype,
|
||||
axis=axis,
|
||||
out_shape=out_shape,
|
||||
)
|
||||
if out_shape is None:
|
||||
out_shape = weight.shape
|
||||
out_shape[axis] *= self.num_elem_per_storage
|
||||
axis = axis if axis >= 0 else len(out_shape) + axis
|
||||
return te.compute(
|
||||
shape=out_shape,
|
||||
fcompute=lambda *idx: tirx.Mul(
|
||||
tirx.Sub(
|
||||
float_weight(*idx),
|
||||
tir_max_int,
|
||||
),
|
||||
scale(*idx[:axis], idx[axis] // self.group_size, *idx[axis + 1 :]),
|
||||
),
|
||||
name="dequantize",
|
||||
)
|
||||
|
||||
def quantize_weight(
|
||||
self, weight: Tensor, axis: int = -1, output_transpose: bool = False
|
||||
) -> List[Tensor]: # noqa: UP006
|
||||
"""
|
||||
Quantize weight with group quantization
|
||||
|
||||
Parameters
|
||||
----------
|
||||
weight : Tensor
|
||||
The original weight.
|
||||
|
||||
axis : int
|
||||
The group axis.
|
||||
|
||||
output_transpose : bool
|
||||
Whether to transpose the output quantized weight. Only 2D weight is supported.
|
||||
|
||||
Returns
|
||||
------
|
||||
ret: List[Tensor]
|
||||
The list of group quantized weights.
|
||||
"""
|
||||
device = weight.device
|
||||
device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
|
||||
axis = axis if axis >= 0 else len(weight.shape) + axis
|
||||
|
||||
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():
|
||||
lv = bb.emit_te(self._quantize, weight_var, axis, output_transpose)
|
||||
gv = bb.emit_output(lv)
|
||||
bb.emit_func_output(gv)
|
||||
return bb.finalize()
|
||||
|
||||
key = (
|
||||
f"({weight.shape}, {weight.dtype}, {device_type}, "
|
||||
f"axis={axis}, output_transpose={output_transpose})"
|
||||
)
|
||||
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(), device=device)
|
||||
self._quantize_func_cache[key] = quantize_func
|
||||
return quantize_func(weight)
|
||||
|
||||
def _quantize(
|
||||
self,
|
||||
weight: te.Tensor,
|
||||
axis: int = -1,
|
||||
output_transpose: bool = False,
|
||||
) -> Tuple[te.Tensor, te.Tensor]: # noqa: UP006
|
||||
"""Group quantization for weight tensor, defined in tensor expression."""
|
||||
max_int = tirx.const(self.max_int_value, self.model_dtype)
|
||||
shape = weight.shape
|
||||
axis = axis if axis >= 0 else len(shape) + axis
|
||||
k = shape[axis]
|
||||
# compute scale per group
|
||||
r = te.reduce_axis((0, self.group_size), name="r")
|
||||
num_group = tirx.ceildiv(k, self.group_size)
|
||||
scale_shape = (*shape[:axis], num_group, *shape[axis + 1 :])
|
||||
max_abs = te.compute(
|
||||
shape=scale_shape,
|
||||
fcompute=lambda *idx: te.max(
|
||||
tirx.if_then_else(
|
||||
idx[axis] * self.group_size + r < k,
|
||||
te.abs(
|
||||
weight(
|
||||
*idx[:axis],
|
||||
idx[axis] * self.group_size + r,
|
||||
*idx[axis + 1 :],
|
||||
)
|
||||
),
|
||||
te.min_value(self.model_dtype),
|
||||
),
|
||||
axis=r,
|
||||
),
|
||||
name="max_abs_value",
|
||||
)
|
||||
scale = te.compute(
|
||||
scale_shape,
|
||||
lambda *idx: max_abs(*idx).astype(self.model_dtype) / max_int,
|
||||
name="scale",
|
||||
)
|
||||
# compute scaled weight
|
||||
scaled_weight = te.compute(
|
||||
shape=weight.shape,
|
||||
fcompute=lambda *idx: tirx.min(
|
||||
tirx.max(
|
||||
tirx.round(
|
||||
weight(*idx)
|
||||
/ scale(*idx[:axis], idx[axis] // self.group_size, *idx[axis + 1 :])
|
||||
+ max_int
|
||||
),
|
||||
tirx.const(0, self.model_dtype),
|
||||
),
|
||||
max_int * 2,
|
||||
).astype(self.storage_dtype),
|
||||
)
|
||||
# compute quantized weight per storage
|
||||
num_storage = self.num_storage_per_group * num_group
|
||||
quantized_weight_shape = (*shape[:axis], num_storage, *shape[axis + 1 :])
|
||||
quantized_weight = pack_weight(
|
||||
scaled_weight,
|
||||
axis=axis,
|
||||
num_elem_per_storage=self.num_elem_per_storage,
|
||||
weight_dtype=self.quantize_dtype,
|
||||
storage_dtype=self.storage_dtype,
|
||||
out_shape=quantized_weight_shape,
|
||||
)
|
||||
if output_transpose:
|
||||
if len(quantized_weight.shape) != 2 or len(scale.shape) != 2:
|
||||
raise ValueError(
|
||||
"Does not support transpose output quantized weight with ndim != 2"
|
||||
)
|
||||
quantized_weight = topi.transpose(quantized_weight)
|
||||
scale = topi.transpose(scale)
|
||||
return quantized_weight, scale
|
||||
|
||||
|
||||
class GroupQuantizeLinear(nn.Module):
|
||||
"""An nn.Linear module with group quantization"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: Union[int, tirx.Var],
|
||||
config: GroupQuantize,
|
||||
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
|
||||
num_group = tirx.ceildiv(in_features, config.group_size)
|
||||
num_shards = config.tensor_parallel_shards
|
||||
if num_shards > 1 and (in_features * num_shards // config.group_size) % num_shards != 0:
|
||||
raise ValueError(
|
||||
f"The linear dimension {in_features * num_shards} has "
|
||||
f"{in_features * num_shards // config.group_size} groups under group size "
|
||||
f"{config.group_size}. The groups cannot be evenly distributed on "
|
||||
f"{num_shards} GPUs.\n"
|
||||
"Possible solutions: reduce number of GPUs, or use quantization with smaller "
|
||||
"group size."
|
||||
)
|
||||
if config.linear_weight_layout == "KN":
|
||||
self.q_weight = nn.Parameter(
|
||||
(config.num_storage_per_group * num_group, out_features),
|
||||
config.storage_dtype,
|
||||
)
|
||||
self.q_scale = nn.Parameter((num_group, out_features), config.model_dtype)
|
||||
else:
|
||||
self.q_weight = nn.Parameter(
|
||||
(out_features, config.num_storage_per_group * num_group),
|
||||
config.storage_dtype,
|
||||
)
|
||||
self.q_scale = nn.Parameter((out_features, num_group), 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: GroupQuantize) -> "GroupQuantizeLinear":
|
||||
"""
|
||||
Converts a non-quantized nn.Linear to a group quantized GroupQuantizeLinear
|
||||
|
||||
Parameters
|
||||
----------
|
||||
src : nn.Linear
|
||||
The non-quantized nn.Linear.
|
||||
|
||||
config : GroupQuantize
|
||||
The group quantization config.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : GroupQuantizeLinear
|
||||
The group quantized GroupQuantizeLinear layer.
|
||||
"""
|
||||
# For dynamic shape, src.out_features is `"name"`; src.weight.shape[0] is `tirx.Var("name")`
|
||||
out_features, in_features = src.weight.shape
|
||||
quantized_linear = GroupQuantizeLinear(
|
||||
in_features=in_features,
|
||||
out_features=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
|
||||
if "shard_strategy" in src.weight.attrs:
|
||||
shard = src.weight.attrs["shard_strategy"]
|
||||
apply_sharding(shard, f"{shard.name}_q_weight", quantized_linear.q_weight)
|
||||
apply_sharding(shard, f"{shard.name}_q_scale", quantized_linear.q_scale)
|
||||
return quantized_linear
|
||||
|
||||
def forward(self, x: nn.Tensor) -> nn.Tensor:
|
||||
"""
|
||||
Forward method for group quantized linear layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : nn.Tensor
|
||||
The input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : nn.Tensor
|
||||
The output tensor for the group quantized linear layer.
|
||||
"""
|
||||
w = nn.op.tensor_expr_op(
|
||||
lambda weight, scale: self.config._dequantize(
|
||||
weight,
|
||||
scale,
|
||||
axis=self.config.linear_quant_axis,
|
||||
out_shape=(
|
||||
[
|
||||
(
|
||||
tirx.IntImm("int64", self.out_features)
|
||||
if isinstance(self.out_features, int)
|
||||
else weight.shape[0]
|
||||
), # Reuse same tirx.Var for symbolic shape (after Exporter)
|
||||
tirx.IntImm("int64", self.in_features),
|
||||
]
|
||||
if self.config.linear_weight_layout == "NK"
|
||||
else [
|
||||
tirx.IntImm("int64", self.in_features),
|
||||
(
|
||||
tirx.IntImm("int64", self.out_features)
|
||||
if isinstance(self.out_features, int)
|
||||
else weight.shape[1]
|
||||
), # Reuse same tirx.Var for symbolic shape (after Exporter)
|
||||
]
|
||||
),
|
||||
),
|
||||
name_hint="dequantize",
|
||||
args=[self.q_weight, self.q_scale],
|
||||
)
|
||||
if self.config.linear_weight_layout == "NK":
|
||||
w = nn.op.permute_dims(w)
|
||||
x = nn.op.matmul(x, w, out_dtype=self.out_dtype)
|
||||
if self.bias is not None:
|
||||
x = x + self.bias
|
||||
return x
|
||||
|
||||
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
|
||||
|
||||
|
||||
class GroupQuantizeEmbedding(nn.Module):
|
||||
"""An nn.Embedding module with group quantization"""
|
||||
|
||||
def __init__(self, num: Union[int, tirx.Var], dim: int, config: GroupQuantize):
|
||||
self.num = num
|
||||
self.dim = dim
|
||||
self.config = config
|
||||
num_group = tirx.ceildiv(dim, config.group_size)
|
||||
self.q_weight = nn.Parameter(
|
||||
(num, config.num_storage_per_group * num_group), config.storage_dtype
|
||||
)
|
||||
self.q_scale = nn.Parameter((num, num_group), config.model_dtype)
|
||||
|
||||
@staticmethod
|
||||
def from_embedding(embedding: nn.Embedding, config: GroupQuantize) -> "GroupQuantizeEmbedding":
|
||||
"""
|
||||
Converts a non-quantized nn.Embedding to a group quantized GroupQuantizeEmbedding
|
||||
|
||||
Parameters
|
||||
----------
|
||||
linear : nn.Embedding
|
||||
The non-quantized nn.Embedding.
|
||||
|
||||
config : GroupQuantize
|
||||
The group quantization config.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : GroupQuantizeEmbedding
|
||||
The group quantized GroupQuantizeEmbedding layer.
|
||||
"""
|
||||
num, dim = embedding.weight.shape
|
||||
return GroupQuantizeEmbedding(num, dim, config)
|
||||
|
||||
def forward(self, x: nn.Tensor):
|
||||
"""
|
||||
Forward method for group quantized embedding layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : nn.Tensor
|
||||
The input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : nn.Tensor
|
||||
The output tensor for the embedding layer.
|
||||
"""
|
||||
w = nn.op.tensor_expr_op(
|
||||
lambda weight, scale: self.config._dequantize(
|
||||
weight,
|
||||
scale,
|
||||
axis=-1,
|
||||
out_shape=[
|
||||
(
|
||||
tirx.IntImm("int64", self.num)
|
||||
if isinstance(self.num, int)
|
||||
else weight.shape[0]
|
||||
), # Reuse same tirx.Var for symbolic shape (after Exporter)
|
||||
tirx.IntImm("int64", self.dim),
|
||||
],
|
||||
),
|
||||
name_hint="dequantize",
|
||||
args=[self.q_weight, self.q_scale],
|
||||
)
|
||||
if x.ndim == 1:
|
||||
return nn.op.take(w, x, axis=0)
|
||||
return nn.op.reshape(
|
||||
nn.op.take(w, nn.op.reshape(x, shape=[-1]), axis=0),
|
||||
shape=[*x.shape, self.dim],
|
||||
)
|
||||
|
||||
def lm_head_forward(self, x: nn.Tensor):
|
||||
"""The lm_head forwarding, which dequantizes the weight
|
||||
and multiplies it with the input tensor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : nn.Tensor
|
||||
The input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : nn.Tensor
|
||||
The output tensor for the lm_head layer.
|
||||
"""
|
||||
w = nn.op.tensor_expr_op(
|
||||
lambda weight, scale: self.config._dequantize(
|
||||
weight,
|
||||
scale,
|
||||
axis=-1,
|
||||
out_shape=[
|
||||
(
|
||||
tirx.IntImm("int64", self.num)
|
||||
if isinstance(self.num, int)
|
||||
else weight.shape[0]
|
||||
),
|
||||
tirx.IntImm("int64", self.dim),
|
||||
],
|
||||
),
|
||||
name_hint="dequantize",
|
||||
args=[self.q_weight, self.q_scale],
|
||||
)
|
||||
w = nn.op.permute_dims(w)
|
||||
return nn.op.matmul(x, w, out_dtype="float32")
|
||||
|
||||
|
||||
class GroupQuantizeMixtralExperts(nn.Module):
|
||||
"""An MixtralExperts module with group quantization"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_local_experts,
|
||||
in_features,
|
||||
out_features,
|
||||
config: GroupQuantize,
|
||||
):
|
||||
self.num_local_experts = num_local_experts
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.config = config
|
||||
num_group = tirx.ceildiv(in_features, config.group_size)
|
||||
self.q_weight = nn.Parameter(
|
||||
(num_local_experts, out_features, config.num_storage_per_group * num_group),
|
||||
config.storage_dtype,
|
||||
)
|
||||
self.q_scale = nn.Parameter(
|
||||
(num_local_experts, out_features, num_group), config.model_dtype
|
||||
)
|
||||
self.quantize_dtype = config.quantize_dtype
|
||||
self.group_size = config.group_size
|
||||
self.dtype = config.model_dtype
|
||||
if config.linear_weight_layout == "KN":
|
||||
raise NotImplementedError("GroupQuantizeMixtralExperts does not support KN layout now.")
|
||||
|
||||
@staticmethod
|
||||
def from_mixtral_experts(
|
||||
src: "MixtralExperts", config: GroupQuantize
|
||||
) -> "GroupQuantizeMixtralExperts":
|
||||
"""
|
||||
Converts a non-quantized MixtralExperts to a group quantized GroupQuantizeMixtralExperts
|
||||
|
||||
Parameters
|
||||
----------
|
||||
src : MixtralExperts
|
||||
The non-quantized MixtralExperts
|
||||
|
||||
config : GroupQuantize
|
||||
The group quantization config.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : GroupQuantizeMixtralExperts
|
||||
The group quantized GroupQuantizeMixtralExperts layer.
|
||||
"""
|
||||
quantized_mistral_experts = GroupQuantizeMixtralExperts(
|
||||
num_local_experts=src.num_local_experts,
|
||||
in_features=src.in_features,
|
||||
out_features=src.out_features,
|
||||
config=config,
|
||||
)
|
||||
if "shard_strategy" in src.weight.attrs:
|
||||
shard = src.weight.attrs["shard_strategy"]
|
||||
apply_sharding(shard, f"{shard.name}_q_weight", quantized_mistral_experts.q_weight)
|
||||
apply_sharding(shard, f"{shard.name}_q_scale", quantized_mistral_experts.q_scale)
|
||||
return quantized_mistral_experts
|
||||
|
||||
def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor:
|
||||
"""Forward method for group quantized mistral experts.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : nn.Tensor
|
||||
The input tensor.
|
||||
|
||||
indptr: nn.Tensor
|
||||
The indptr tensor
|
||||
|
||||
single_batch_decode: bool
|
||||
Whether to use single-batch decode
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : nn.Tensor
|
||||
The output tensor for the group quantized mistral experts layer.
|
||||
"""
|
||||
from mlc_llm.op import moe_matmul
|
||||
|
||||
assert x.ndim == 2
|
||||
if indptr.ndim == 2: # single-batch
|
||||
assert indptr.shape[0] == 1
|
||||
return moe_matmul.dequantize_gemv(
|
||||
x,
|
||||
self.q_weight,
|
||||
self.q_scale,
|
||||
indptr,
|
||||
quantize_dtype=self.quantize_dtype,
|
||||
group_size=self.group_size,
|
||||
)
|
||||
assert indptr.ndim == 1
|
||||
return moe_matmul.dequantize_group_gemm(
|
||||
x,
|
||||
self.q_weight,
|
||||
self.q_scale,
|
||||
indptr,
|
||||
quantize_dtype=self.quantize_dtype,
|
||||
indptr_dtype=indptr.dtype,
|
||||
group_size=self.group_size,
|
||||
)
|
||||
Reference in New Issue
Block a user