wehub-resource-sync
123 lines
4.3 KiB
Python
123 lines
4.3 KiB
Python
"""Quantization techniques for FP8"""
|
|
|
|
import numpy as np
|
|
from tvm import relax, runtime
|
|
from tvm.relax.frontend import nn
|
|
|
|
from mlc_llm.nn import MixtralExperts
|
|
|
|
from ..op import cutlass, extern, moe_matmul
|
|
from . import per_tensor_quantization as ptq
|
|
from .utils import apply_sharding
|
|
|
|
|
|
class FP8PerTensorQuantizeMixtralExperts(ptq.PerTensorQuantizeMixtralExperts):
|
|
"""MixtralExperts with per-tensor quantization in FP8."""
|
|
|
|
def __init__(
|
|
self,
|
|
num_local_experts,
|
|
in_features,
|
|
out_features,
|
|
config: ptq.PerTensorQuantize,
|
|
name: str,
|
|
tensor_parallel_shards=1,
|
|
):
|
|
super().__init__(num_local_experts, in_features, out_features, config, name)
|
|
self.tensor_parallel_shards = tensor_parallel_shards
|
|
|
|
@staticmethod
|
|
def from_mixtral_experts(
|
|
src: "MixtralExperts",
|
|
config: ptq.PerTensorQuantize,
|
|
name: str,
|
|
) -> "FP8PerTensorQuantizeMixtralExperts":
|
|
"""
|
|
Converts a non-quantized MixtralExperts to a per-tensor quantized MixtralExperts.
|
|
|
|
Parameters
|
|
----------
|
|
src : MixtralExperts
|
|
The non-quantized MixtralExperts
|
|
|
|
config : PerTensorQuantize
|
|
The FP8 quantization weight_config.
|
|
|
|
name : str
|
|
The name of the layer.
|
|
|
|
Returns
|
|
-------
|
|
ret : MixtralExpertsFP8
|
|
The per-tensor quantized MixtralExperts.
|
|
"""
|
|
quantized_mistral_experts = FP8PerTensorQuantizeMixtralExperts(
|
|
num_local_experts=src.num_local_experts,
|
|
in_features=src.in_features,
|
|
out_features=src.out_features,
|
|
config=config,
|
|
name=name,
|
|
tensor_parallel_shards=src.tensor_parallel_shards,
|
|
)
|
|
|
|
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)
|
|
# scale doesn't need to be sharded since it's the same for all shards
|
|
|
|
return quantized_mistral_experts
|
|
|
|
def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor:
|
|
w = self.q_weight
|
|
|
|
if self.config.calibration_mode == "max":
|
|
_, x_scale = self.config.quantize_float8(
|
|
x,
|
|
quantize_dtype=self.config.activation_dtype,
|
|
storage_dtype=self.config.activation_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)
|
|
|
|
if indptr.ndim == 2:
|
|
assert indptr.shape[0] == 1
|
|
return moe_matmul.dequantize_float8_gemv(
|
|
x, w, self.q_scale, indptr, self.config.weight_dtype
|
|
)
|
|
|
|
if extern.get_store().cutlass_group_gemm:
|
|
if self.config.calibration_mode == "inference":
|
|
if self.q_calibration_scale is not None:
|
|
x /= self.q_calibration_scale.astype(x.dtype)
|
|
x_q = nn.op.astype(x, dtype=self.config.activation_dtype)
|
|
x_scale = self.q_calibration_scale
|
|
|
|
scale = (
|
|
x_scale * self.q_scale
|
|
if self.q_scale is not None
|
|
else nn.wrap_nested(
|
|
relax.Constant(runtime.tensor(np.array([1.0]).astype("float32"))),
|
|
"scale",
|
|
)
|
|
)
|
|
return cutlass.group_gemm(
|
|
x_q, w, indptr, scale, self.config.weight_dtype, self.config.model_dtype
|
|
)
|
|
# Note: convert_weight is target agnostic, so a fallback must be provided
|
|
w = nn.tensor_expr_op(
|
|
self.config.dequantize_float8,
|
|
"dequantize",
|
|
args=[w, self.q_scale, self.config.weight_dtype],
|
|
)
|
|
return moe_matmul.group_gemm(x, w, indptr)
|
|
|
|
|
|
ptq.PerTensorQuantizeMixtralExperts._IMPL["fp8"] = FP8PerTensorQuantizeMixtralExperts
|