wehub-resource-sync
189 lines
6.3 KiB
Python
189 lines
6.3 KiB
Python
"""Common utilities for quantization"""
|
|
|
|
from collections.abc import Sequence
|
|
from typing import Callable, List, Optional # noqa: UP035
|
|
|
|
from tvm import IRModule, relax, te, tirx
|
|
from tvm.relax.frontend import nn
|
|
from tvm.runtime import DataType, DataTypeCode
|
|
from tvm.s_tir import dlight as dl
|
|
from tvm.target import Target
|
|
|
|
from mlc_llm.support import tensor_parallel as tp
|
|
|
|
|
|
def convert_uint_to_float(
|
|
weight: te.Tensor,
|
|
bits: int,
|
|
num_elem_per_storage: int,
|
|
storage_dtype: str,
|
|
model_dtype: str,
|
|
axis: int = -1,
|
|
out_shape: Optional[List[tirx.Expr]] = None, # noqa: UP006
|
|
ft_reorder: Optional[bool] = False,
|
|
) -> te.Tensor:
|
|
"""Convert a quantized uint weight to an unquantized float weight."""
|
|
tir_bin_mask = tirx.const((1 << bits) - 1, storage_dtype)
|
|
if out_shape is None:
|
|
out_shape = weight.shape
|
|
out_shape[axis] *= num_elem_per_storage
|
|
axis = axis if axis >= 0 else len(out_shape) + axis
|
|
return te.compute(
|
|
shape=out_shape,
|
|
fcompute=lambda *idx: tirx.bitwise_and(
|
|
tirx.shift_right(
|
|
weight(*idx[:axis], idx[axis] // num_elem_per_storage, *idx[axis + 1 :]),
|
|
(
|
|
(
|
|
(idx[axis] % num_elem_per_storage) % 2 * 4
|
|
+ (idx[axis] % num_elem_per_storage) // 2
|
|
)
|
|
* bits
|
|
if ft_reorder
|
|
else (idx[axis] % num_elem_per_storage) * bits
|
|
).astype(storage_dtype),
|
|
),
|
|
tir_bin_mask,
|
|
).astype(model_dtype),
|
|
)
|
|
|
|
|
|
def is_final_fc(name: str) -> bool:
|
|
"""Determines whether the parameter is the last layer based on its name."""
|
|
# TODO: use more specious condition to determine final fc
|
|
return name in ["head", "lm_head", "lm_head.linear", "embed_out"]
|
|
|
|
|
|
def is_moe_gate(name: str, node: nn.Linear) -> bool:
|
|
"""Check whether the parameter is the MoE gate layer."""
|
|
return name.endswith("gate") and isinstance(node.out_features, int) and node.out_features <= 256
|
|
|
|
|
|
def compile_quantize_func(mod: IRModule, device) -> Callable:
|
|
"""Compile a quantization function for a given device."""
|
|
device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
|
|
if device_type in ["cuda", "rocm", "metal", "vulkan", "opencl"]:
|
|
target = Target.current()
|
|
if target is None:
|
|
target = Target.from_device(device)
|
|
with target:
|
|
mod = dl.ApplyDefaultSchedule(
|
|
dl.gpu.Reduction(),
|
|
dl.gpu.GeneralReduction(),
|
|
dl.gpu.Fallback(),
|
|
)(mod)
|
|
elif device_type == "cpu":
|
|
target = "llvm"
|
|
mod = relax.transform.LegalizeOps()(mod)
|
|
else:
|
|
raise NotImplementedError(f"Device type {device_type} is not supported")
|
|
ex = relax.build(mod, target=target)
|
|
vm = relax.VirtualMachine(ex, device)
|
|
return vm["main"]
|
|
|
|
|
|
def apply_sharding(shard_strategy, name: str, weight: nn.Parameter):
|
|
"""Apply sharding strategy to a weight."""
|
|
if isinstance(shard_strategy, tp.ShardSingleDim):
|
|
weight.attrs["shard_strategy"] = tp.ShardSingleDim(
|
|
name=name,
|
|
dim=shard_strategy.dim,
|
|
segs=shard_strategy.segs,
|
|
)
|
|
else:
|
|
raise NotImplementedError(f"Unknowing sharding strategy: {shard_strategy}")
|
|
|
|
|
|
def convert_uint_packed_fp8_to_float(
|
|
weight: te.Tensor,
|
|
num_elem_per_storage: int,
|
|
storage_dtype: str,
|
|
model_dtype: str,
|
|
quant_dtype: str,
|
|
axis: int = -1,
|
|
out_shape: Optional[Sequence[tirx.Expr]] = None,
|
|
) -> te.Tensor:
|
|
"""Unpack a fp8 value from the storage dtype and convert to float."""
|
|
assert quant_dtype in ["float8_e4m3fn", "float8_e5m2"]
|
|
assert DataType(storage_dtype).type_code == DataTypeCode.UINT
|
|
bits = DataType(quant_dtype).bits
|
|
elem_storage_dtype = DataType(f"uint{bits}")
|
|
tir_bin_mask = tirx.const((1 << bits) - 1, "uint8")
|
|
if axis < 0:
|
|
axis += len(weight.shape)
|
|
if out_shape is None:
|
|
out_shape = (
|
|
*weight.shape[:axis],
|
|
weight.shape[axis] * num_elem_per_storage,
|
|
*weight.shape[axis + 1 :],
|
|
)
|
|
axis = axis if axis >= 0 else len(out_shape) + axis
|
|
return te.compute(
|
|
shape=out_shape,
|
|
fcompute=lambda *idx: tirx.reinterpret(
|
|
quant_dtype,
|
|
tirx.bitwise_and(
|
|
tirx.shift_right(
|
|
weight(*idx[:axis], idx[axis] // num_elem_per_storage, *idx[axis + 1 :]),
|
|
((idx[axis] % num_elem_per_storage) * bits).astype(storage_dtype),
|
|
).astype(elem_storage_dtype),
|
|
tir_bin_mask,
|
|
),
|
|
).astype(model_dtype),
|
|
)
|
|
|
|
|
|
def pack_weight(
|
|
weight: te.Tensor,
|
|
axis: int,
|
|
num_elem_per_storage: int,
|
|
weight_dtype: str,
|
|
storage_dtype: str,
|
|
out_shape: Optional[Sequence[tirx.Expr]] = None,
|
|
):
|
|
"""Convert a tensor to a packed format by packing consecutive bits.
|
|
This can be useful for sub-byte quantization.
|
|
|
|
Parameters
|
|
----------
|
|
weight : te.Tensor
|
|
The weight
|
|
axis : int
|
|
The axis to pack.
|
|
num_elem_per_storage : int
|
|
The number of elements per storage.
|
|
weight_dtype : str
|
|
The dtype of the input tensor.
|
|
storage_dtype : str
|
|
The dtype of the packed tensor.
|
|
out_shape : Optional[Sequence[tirx.Expr]]
|
|
The output shape of the packed tensor. Zero-padding is added if needed.
|
|
"""
|
|
assert weight.dtype == storage_dtype
|
|
shape = weight.shape
|
|
if axis < 0:
|
|
axis += len(shape)
|
|
k = shape[axis]
|
|
axis = axis if axis >= 0 else len(shape) + axis
|
|
if out_shape is None:
|
|
out_shape = (
|
|
*shape[:axis],
|
|
tirx.ceildiv(k, num_elem_per_storage),
|
|
*shape[axis + 1 :],
|
|
)
|
|
r = te.reduce_axis((0, num_elem_per_storage), name="r")
|
|
packed_weight = te.compute(
|
|
shape=out_shape,
|
|
fcompute=lambda *idx: tirx.sum(
|
|
tirx.if_then_else(
|
|
idx[axis] * num_elem_per_storage + r < k,
|
|
weight(*idx[:axis], idx[axis] * num_elem_per_storage + r, *idx[axis + 1 :])
|
|
<< (r * DataType(weight_dtype).bits),
|
|
tirx.const(0, storage_dtype),
|
|
),
|
|
axis=r,
|
|
),
|
|
name="packed_weight",
|
|
).astype(storage_dtype)
|
|
return packed_weight
|