chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,700 @@
|
||||
"""The per-tensor quantization config"""
|
||||
|
||||
import functools
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Type, Union # noqa: UP035
|
||||
|
||||
import numpy as np
|
||||
from tvm import DataType, DataTypeCode, IRModule, relax, runtime, 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.op import cutlass, extern
|
||||
from mlc_llm.support import logging
|
||||
|
||||
from .utils import (
|
||||
apply_sharding,
|
||||
compile_quantize_func,
|
||||
convert_uint_packed_fp8_to_float,
|
||||
is_final_fc,
|
||||
is_moe_gate,
|
||||
pack_weight,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PerTensorQuantize:
|
||||
"""Configuration for per-tensor quantization"""
|
||||
|
||||
name: str
|
||||
kind: str
|
||||
activation_dtype: Literal["float8_e4m3fn", "float8_e5m2"]
|
||||
weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"]
|
||||
storage_dtype: Literal["uint32", "float8_e4m3fn", "float8_e5m2"]
|
||||
model_dtype: Literal["float16"]
|
||||
quantize_embedding: bool = True
|
||||
quantize_final_fc: bool = True
|
||||
quantize_linear: bool = True
|
||||
|
||||
num_elem_per_storage: int = 0
|
||||
max_int_value: int = 0
|
||||
use_scale: bool = True
|
||||
# The calibration mode for quantization. If set to "inference", the model is built for
|
||||
# inference. This should be used after calibration is done.
|
||||
# If set to "max", the model is built for calibration that computes the scale using max value of
|
||||
# the activations.
|
||||
calibration_mode: Literal["inference", "max"] = "inference"
|
||||
tensor_parallel_shards: int = 1
|
||||
|
||||
def __post_init__(self):
|
||||
assert self.kind == "per-tensor-quant"
|
||||
self.num_elem_per_storage = (
|
||||
DataType(self.storage_dtype).bits // DataType(self.weight_dtype).bits
|
||||
)
|
||||
self.max_int_value = int(tirx.max_value(self.weight_dtype).value)
|
||||
self._quantize_func_cache = {}
|
||||
|
||||
def quantize_model(
|
||||
self,
|
||||
model: nn.Module,
|
||||
quant_map: QuantizeMapping,
|
||||
name_prefix: str,
|
||||
tensor_parallel_shards: int,
|
||||
) -> nn.Module:
|
||||
"""
|
||||
Quantize model with per-tensor 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.
|
||||
|
||||
tensor_parallel_shards : int
|
||||
The number of tensor parallel shards.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : nn.Module
|
||||
The quantized nn.Module.
|
||||
"""
|
||||
|
||||
self.tensor_parallel_shards = tensor_parallel_shards
|
||||
|
||||
class _Mutator(nn.Mutator):
|
||||
def __init__(self, config: PerTensorQuantize, 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 per-tensor 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.
|
||||
"""
|
||||
weight_name = f"{name}.weight"
|
||||
param_names = (
|
||||
[f"{name}.q_weight", f"{name}.q_scale"]
|
||||
if self.config.use_scale
|
||||
else [
|
||||
f"{name}.q_weight",
|
||||
]
|
||||
)
|
||||
if (
|
||||
isinstance(node, nn.Linear)
|
||||
and self.config.quantize_linear
|
||||
and (not is_final_fc(name) or self.config.quantize_final_fc)
|
||||
and not is_moe_gate(name, node)
|
||||
):
|
||||
self.quant_map.param_map[weight_name] = param_names
|
||||
self.quant_map.map_func[weight_name] = self.config.quantize_weight
|
||||
op = PerTensorQuantizeLinear.from_linear(node, self.config, name)
|
||||
elif isinstance(node, nn.Embedding) and self.config.quantize_embedding:
|
||||
self.quant_map.param_map[weight_name] = param_names
|
||||
self.quant_map.map_func[weight_name] = self.config.quantize_weight
|
||||
op = PerTensorQuantizeEmbedding.from_embedding(node, self.config)
|
||||
elif isinstance(node, MixtralExperts):
|
||||
self.quant_map.param_map[weight_name] = param_names
|
||||
self.quant_map.map_func[weight_name] = self.config.quantize_weight
|
||||
op = PerTensorQuantizeMixtralExperts.from_mixtral_experts(
|
||||
node, self.config, name
|
||||
)
|
||||
else:
|
||||
return self.visit(name, node)
|
||||
|
||||
if hasattr(op, "q_calibration_scale") and op.q_calibration_scale:
|
||||
# update quant_map for calibration scale
|
||||
param_name = f"{name}.q_calibration_scale"
|
||||
old_map_func = self.quant_map.map_func[weight_name]
|
||||
|
||||
def map_func(*args, **kwargs):
|
||||
# placeholder for calibration scale, the actual value will be set after
|
||||
# calibration.
|
||||
scale = runtime.empty(
|
||||
shape=op.q_calibration_scale.shape,
|
||||
dtype=op.q_calibration_scale.dtype,
|
||||
)
|
||||
return [*old_map_func(*args, **kwargs), scale]
|
||||
|
||||
self.quant_map.param_map[weight_name].append(param_name)
|
||||
self.quant_map.map_func[weight_name] = map_func
|
||||
return op
|
||||
|
||||
model.to(dtype=self.model_dtype)
|
||||
mutator = _Mutator(self, quant_map)
|
||||
model = mutator.visit(name_prefix, model)
|
||||
return model
|
||||
|
||||
def quantize_weight(self, weight) -> List[Tensor]: # noqa: UP006
|
||||
"""
|
||||
Quantize weight with per-tensor quantization.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
weight : Tensor
|
||||
The weight to quantize.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : List[Tensor]
|
||||
The quantized weight and the scale if use_scale is True.
|
||||
"""
|
||||
device = weight.device
|
||||
device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
|
||||
|
||||
def _create_quantize_func() -> IRModule:
|
||||
if DataType(self.weight_dtype).type_code in [
|
||||
DataTypeCode.Float8E4M3FN,
|
||||
DataTypeCode.Float8E5M2,
|
||||
]:
|
||||
quantize_func = functools.partial(
|
||||
self.quantize_float8,
|
||||
quantize_dtype=self.weight_dtype,
|
||||
storage_dtype=self.storage_dtype,
|
||||
)
|
||||
else:
|
||||
assert NotImplementedError()
|
||||
|
||||
class Quantizer(nn.Module):
|
||||
"""Quantizer module for per-tensor quantization."""
|
||||
|
||||
def main(self, weight: nn.Tensor):
|
||||
return quantize_func(weight)
|
||||
|
||||
mod = Quantizer()
|
||||
mod, _ = mod.export_tvm(
|
||||
spec={"main": {"weight": nn.spec.Tensor(weight.shape, weight.dtype)}}
|
||||
)
|
||||
return mod
|
||||
|
||||
key = f"({weight.shape}, {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(), device)
|
||||
self._quantize_func_cache[key] = quantize_func
|
||||
return quantize_func(weight)
|
||||
|
||||
def quantize_float8(
|
||||
self,
|
||||
tensor: nn.Tensor,
|
||||
quantize_dtype: str,
|
||||
storage_dtype: str,
|
||||
) -> Union[Tuple[nn.Tensor], Tuple[nn.Tensor, nn.Tensor]]: # noqa: UP006
|
||||
"""Per-tensor quantization for weight tensor, defined in tensor expression."""
|
||||
|
||||
if self.use_scale:
|
||||
# min_scaling_factor taken from TRT-LLM
|
||||
def _compute_scale(x: te.Tensor) -> te.Tensor:
|
||||
max_abs = topi.max(topi.abs(x))
|
||||
min_scaling_factor = tirx.const(
|
||||
1.0 / (self.max_int_value * 512.0), self.model_dtype
|
||||
)
|
||||
scale = topi.maximum(
|
||||
max_abs.astype(self.model_dtype) / self.max_int_value,
|
||||
min_scaling_factor,
|
||||
).astype("float32")
|
||||
scale = topi.expand_dims(scale, axis=0)
|
||||
return scale
|
||||
|
||||
scale = nn.tensor_expr_op(_compute_scale, "compute_scale", args=[tensor])
|
||||
else:
|
||||
scale = None
|
||||
|
||||
def _compute_quantized_tensor(weight: te.Tensor, scale: Optional[te.Tensor]) -> te.Tensor:
|
||||
elem_storage_dtype = (
|
||||
f"uint{DataType(quantize_dtype).bits}"
|
||||
if DataType(storage_dtype).type_code == DataTypeCode.UINT
|
||||
else quantize_dtype
|
||||
)
|
||||
scaled_tensor = te.compute(
|
||||
shape=weight.shape,
|
||||
fcompute=lambda *idx: tirx.Cast(
|
||||
self.storage_dtype,
|
||||
tirx.reinterpret(
|
||||
elem_storage_dtype,
|
||||
tirx.Cast(
|
||||
quantize_dtype,
|
||||
weight(*idx) / scale(0) if scale is not None else weight(*idx),
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
if quantize_dtype == self.storage_dtype:
|
||||
return scaled_tensor
|
||||
|
||||
packed_weight = pack_weight(
|
||||
scaled_tensor,
|
||||
axis=-1,
|
||||
num_elem_per_storage=self.num_elem_per_storage,
|
||||
weight_dtype=self.weight_dtype,
|
||||
storage_dtype=self.storage_dtype,
|
||||
)
|
||||
|
||||
return packed_weight
|
||||
|
||||
quantized_tensor = nn.tensor_expr_op(
|
||||
_compute_quantized_tensor, "compute_quantized_tensor", args=[tensor, scale]
|
||||
)
|
||||
|
||||
if self.use_scale:
|
||||
return quantized_tensor, scale
|
||||
return (quantized_tensor,)
|
||||
|
||||
def _dequantize(
|
||||
self,
|
||||
q_weight: te.Tensor,
|
||||
scale: Optional[te.Tensor] = None,
|
||||
out_shape: Optional[Sequence[tirx.Expr]] = None,
|
||||
) -> te.Tensor:
|
||||
if self.use_scale:
|
||||
assert scale is not None
|
||||
if DataType(self.weight_dtype).type_code in [
|
||||
DataTypeCode.Float8E4M3FN,
|
||||
DataTypeCode.Float8E5M2,
|
||||
]:
|
||||
return self.dequantize_float8(q_weight, scale, self.weight_dtype, out_shape)
|
||||
raise NotImplementedError()
|
||||
|
||||
def dequantize_float8(
|
||||
self,
|
||||
q_tensor: te.Tensor,
|
||||
scale: Optional[te.Tensor],
|
||||
quantize_dtype: str,
|
||||
out_shape: Optional[Sequence[tirx.Expr]] = None,
|
||||
) -> te.Tensor:
|
||||
"""Dequantize a fp8 tensor (input or weight) to higher-precision float."""
|
||||
if quantize_dtype != self.storage_dtype:
|
||||
dequantized_tensor = convert_uint_packed_fp8_to_float(
|
||||
q_tensor,
|
||||
self.num_elem_per_storage,
|
||||
self.storage_dtype,
|
||||
self.model_dtype,
|
||||
quantize_dtype,
|
||||
axis=-1,
|
||||
out_shape=out_shape,
|
||||
)
|
||||
else:
|
||||
dequantized_tensor = q_tensor.astype(self.model_dtype)
|
||||
if scale is not None:
|
||||
dequantized_tensor = dequantized_tensor * scale.astype(dequantized_tensor.dtype)
|
||||
return dequantized_tensor
|
||||
|
||||
|
||||
class PerTensorQuantizeLinear(nn.Module):
|
||||
"""An nn.Linear module with per-tensor quantization."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: Union[int, tirx.Var],
|
||||
config: PerTensorQuantize,
|
||||
name: str,
|
||||
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 or config.model_dtype
|
||||
self.config = config
|
||||
self.name = name
|
||||
self.q_weight = nn.Parameter(
|
||||
(out_features, tirx.ceildiv(in_features, config.num_elem_per_storage)),
|
||||
config.storage_dtype,
|
||||
)
|
||||
self.q_calibration_scale = None
|
||||
if config.use_scale:
|
||||
self.q_scale = nn.Parameter((1,), "float32")
|
||||
if config.calibration_mode == "inference":
|
||||
self.q_calibration_scale = nn.Parameter((1,), "float32")
|
||||
else:
|
||||
self.q_scale = None
|
||||
if bias:
|
||||
self.bias = nn.Parameter(
|
||||
(out_features,), config.model_dtype if out_dtype is None else out_dtype
|
||||
)
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
@classmethod
|
||||
def from_linear(
|
||||
cls, src: nn.Linear, config: PerTensorQuantize, name: str
|
||||
) -> "PerTensorQuantizeLinear":
|
||||
"""
|
||||
Converts a non-quantized nn.Linear to a per-tensor quantized PerTensorQuantizeLinear
|
||||
|
||||
Parameters
|
||||
----------
|
||||
src : nn.Linear
|
||||
The non-quantized nn.Linear.
|
||||
|
||||
config : PerTensorQuantize
|
||||
The per-tensor quantization config.
|
||||
|
||||
name: str
|
||||
The name of the layer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : PerTensorQuantizeLinear
|
||||
The per-tensor quantized PerTensorQuantizeLinear layer.
|
||||
"""
|
||||
out_features, in_features = src.weight.shape
|
||||
quantized_linear = cls(
|
||||
in_features=in_features,
|
||||
out_features=out_features,
|
||||
config=config,
|
||||
name=name,
|
||||
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)
|
||||
# scale doesn't need to be sharded since it's the same for all shards
|
||||
return quantized_linear
|
||||
|
||||
def forward(self, x: nn.Tensor) -> nn.Tensor:
|
||||
"""
|
||||
Forward method for per-tensor quantized linear layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : nn.Tensor
|
||||
The input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : nn.Tensor
|
||||
The output tensor for the per-tensor quantized linear layer.
|
||||
"""
|
||||
# Note: Use calibration scale when calibration is enabled
|
||||
if self.config.calibration_mode == "inference":
|
||||
if self.q_calibration_scale:
|
||||
x /= self.q_calibration_scale.astype(x.dtype)
|
||||
x_q = x.astype(self.config.activation_dtype)
|
||||
x_scale = self.q_calibration_scale
|
||||
elif self.config.calibration_mode == "max":
|
||||
_, x_scale = self.config.quantize_float8(
|
||||
x,
|
||||
quantize_dtype=self.config.activation_dtype,
|
||||
storage_dtype=self.config.storage_dtype,
|
||||
)
|
||||
if self.config.tensor_parallel_shards > 1:
|
||||
x_scale = nn.ccl_allreduce(x_scale, "max")
|
||||
x_scale = nn.extern(
|
||||
"mlc_llm.calibration_observer",
|
||||
[f"{self.name}.q_calibration_scale", "max", x_scale],
|
||||
out=nn.Tensor.placeholder(x_scale.shape, x_scale.dtype),
|
||||
)
|
||||
x_q = (x / x_scale.astype(x.dtype)).astype(self.config.activation_dtype)
|
||||
x = x_q.astype(self.config.model_dtype) * x_scale.astype(self.config.model_dtype)
|
||||
else:
|
||||
raise ValueError(f"Unknown calibration mode: {self.config.calibration_mode}")
|
||||
|
||||
if (
|
||||
self.config.weight_dtype == self.config.storage_dtype
|
||||
and self.config.calibration_mode == "inference"
|
||||
):
|
||||
if (
|
||||
extern.get_store().cutlass_gemm
|
||||
and functools.reduce(lambda x, y: x * y, x_q.shape[:-1]) != 1
|
||||
):
|
||||
# Dispatch to cutlass kernel for gemm when cutlass is available.
|
||||
scale = (
|
||||
x_scale * self.q_scale
|
||||
if self.config.use_scale
|
||||
else nn.wrap_nested(
|
||||
relax.Constant(runtime.tensor(np.array([1.0]).astype("float32"))),
|
||||
"scale",
|
||||
)
|
||||
)
|
||||
return cutlass.fp8_gemm(
|
||||
x_q,
|
||||
self.q_weight,
|
||||
scale,
|
||||
self.config.weight_dtype,
|
||||
self.config.model_dtype,
|
||||
)
|
||||
x = nn.op.matmul(x_q, nn.permute_dims(self.q_weight), out_dtype="float32")
|
||||
if self.config.use_scale:
|
||||
scale = x_scale * self.q_scale
|
||||
x = x * scale
|
||||
x = x.astype(self.out_dtype)
|
||||
else:
|
||||
w = nn.op.tensor_expr_op(
|
||||
lambda weight, scale: self.config._dequantize(
|
||||
weight,
|
||||
scale,
|
||||
out_shape=[
|
||||
(
|
||||
tirx.IntImm("int64", self.out_features)
|
||||
if isinstance(self.out_features, int)
|
||||
else weight.shape[0]
|
||||
),
|
||||
tirx.IntImm("int64", self.in_features),
|
||||
],
|
||||
),
|
||||
"dequantize",
|
||||
args=[self.q_weight, self.q_scale],
|
||||
)
|
||||
x = nn.op.matmul(x, nn.permute_dims(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)
|
||||
if self.q_scale:
|
||||
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 PerTensorQuantizeEmbedding(nn.Module):
|
||||
"""An nn.Embedding module with group quantization"""
|
||||
|
||||
def __init__(self, num: Union[int, tirx.Var], dim: int, config: PerTensorQuantize):
|
||||
self.num = num
|
||||
self.dim = dim
|
||||
self.config = config
|
||||
self.q_weight = nn.Parameter(
|
||||
(num, tirx.ceildiv(dim, config.num_elem_per_storage)), config.storage_dtype
|
||||
)
|
||||
if self.config.use_scale:
|
||||
self.q_scale = nn.Parameter((1,), "float32")
|
||||
else:
|
||||
self.q_scale = None
|
||||
|
||||
@staticmethod
|
||||
def from_embedding(
|
||||
embedding: nn.Embedding, config: PerTensorQuantize
|
||||
) -> "PerTensorQuantizeEmbedding":
|
||||
"""
|
||||
Converts a non-quantized nn.Embedding to a per-tensor quantized PerTensorQuantizeEmbedding
|
||||
|
||||
Parameters
|
||||
----------
|
||||
linear : nn.Embedding
|
||||
The non-quantized nn.Embedding.
|
||||
|
||||
config : PerTensorQuantize
|
||||
The per-tensor quantization config.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : PerTensorQuantizeEmbedding
|
||||
The per-tensor quantized embedding layer.
|
||||
"""
|
||||
num, dim = embedding.weight.shape
|
||||
return PerTensorQuantizeEmbedding(num, dim, config)
|
||||
|
||||
def forward(self, x: nn.Tensor):
|
||||
"""
|
||||
Forward method for per-tensor 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,
|
||||
out_shape=[
|
||||
(
|
||||
tirx.IntImm("int64", self.num)
|
||||
if isinstance(self.num, int)
|
||||
else weight.shape[0]
|
||||
),
|
||||
tirx.IntImm("int64", self.dim),
|
||||
],
|
||||
),
|
||||
"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,
|
||||
out_shape=[
|
||||
(
|
||||
tirx.IntImm("int64", self.num)
|
||||
if isinstance(self.num, int)
|
||||
else weight.shape[0]
|
||||
),
|
||||
tirx.IntImm("int64", self.dim),
|
||||
],
|
||||
),
|
||||
"dequantize",
|
||||
args=[self.q_weight, self.q_scale],
|
||||
)
|
||||
w = nn.op.permute_dims(w)
|
||||
return nn.op.matmul(x, w, out_dtype="float32")
|
||||
|
||||
|
||||
class PerTensorQuantizeMixtralExperts(nn.Module):
|
||||
"""An MixtralExperts module with group quantization"""
|
||||
|
||||
_IMPL: ClassVar[Dict[str, Type["PerTensorQuantizeMixtralExperts"]]] = {} # noqa: UP006
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_local_experts,
|
||||
in_features,
|
||||
out_features,
|
||||
config: PerTensorQuantize,
|
||||
name: str,
|
||||
):
|
||||
self.num_local_experts = num_local_experts
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.config = config
|
||||
self.name = name
|
||||
self.q_weight = nn.Parameter(
|
||||
(
|
||||
num_local_experts,
|
||||
out_features,
|
||||
tirx.ceildiv(in_features, config.num_elem_per_storage),
|
||||
),
|
||||
config.storage_dtype,
|
||||
)
|
||||
self.q_calibration_scale = None
|
||||
if config.use_scale:
|
||||
self.q_scale = nn.Parameter((1,), "float32")
|
||||
if config.calibration_mode == "inference":
|
||||
self.q_calibration_scale = nn.Parameter((1,), "float32")
|
||||
else:
|
||||
self.q_scale = None
|
||||
|
||||
@staticmethod
|
||||
def from_mixtral_experts(
|
||||
src: "MixtralExperts",
|
||||
config: PerTensorQuantize,
|
||||
name: str,
|
||||
) -> "PerTensorQuantizeMixtralExperts":
|
||||
"""
|
||||
Converts a non-quantized MixtralExperts to a per-tensor quantized
|
||||
PerTensorQuantizeMixtralExperts
|
||||
|
||||
Parameters
|
||||
----------
|
||||
src : MixtralExperts
|
||||
The non-quantized MixtralExperts
|
||||
|
||||
config : PerTensorQuantize
|
||||
The per-tensor quantization config
|
||||
|
||||
name: str
|
||||
The name of the layer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : PerTensorQuantizeMixtralExperts
|
||||
The per-tensor quantized MixtralExperts layer
|
||||
"""
|
||||
if DataType(config.weight_dtype).type_code in [
|
||||
DataTypeCode.Float8E4M3FN,
|
||||
DataTypeCode.Float8E5M2,
|
||||
]:
|
||||
return PerTensorQuantizeMixtralExperts._IMPL["fp8"].from_mixtral_experts(
|
||||
src, config, name
|
||||
)
|
||||
raise NotImplementedError()
|
||||
|
||||
def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor:
|
||||
"""Forward method for per-tensor quantized mistral experts.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : nn.Tensor
|
||||
The input tensor.
|
||||
|
||||
indptr: nn.Tensor
|
||||
The indptr tensor
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : nn.Tensor
|
||||
The output tensor for the per-tensor quantized mistral experts layer.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
Reference in New Issue
Block a user