5af84c8af5
Co-authored-by: Dehua Tang <dehtang@amd.com> Co-authored-by: HAI <hixiao@gmail.com> Co-authored-by: YC Tseng <yctseng@amd.com>
74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
"""
|
|
Common utilities for quark.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Tuple
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def quantize_fp8_scale_tensorwise(w: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
FP8_MAX = 448.0
|
|
scale = w.abs().amax().float() / FP8_MAX
|
|
scaled = (w / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
|
|
return scaled, scale
|
|
|
|
|
|
def quantize_int4_scale_columnwise(
|
|
w: torch.Tensor,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
S4_MAX = 7
|
|
w_flat = w.reshape(-1, w.shape[-1]).float()
|
|
scale = w_flat.abs().amax(axis=-1) / S4_MAX
|
|
scaled = torch.round(w_flat / scale[:, None]).to(torch.int8).clamp(-S4_MAX, S4_MAX)
|
|
return scaled.reshape(w.shape), scale.reshape(w.shape[:-1])
|
|
|
|
|
|
def pack_int4_to_int32(to_pack: torch.Tensor, reorder: bool = True) -> torch.Tensor:
|
|
if to_pack.ndim > 2:
|
|
raise ValueError(
|
|
"Pack: Only supports tensors with dimensions not greater than 2."
|
|
)
|
|
|
|
if reorder:
|
|
order_map = [0, 2, 4, 6, 1, 3, 5, 7]
|
|
else:
|
|
order_map = [0, 1, 2, 3, 4, 5, 6, 7]
|
|
pack_num = 8
|
|
if to_pack.ndim == 2:
|
|
packed = torch.zeros(
|
|
to_pack.shape[0],
|
|
to_pack.shape[1] // pack_num,
|
|
dtype=torch.int32,
|
|
device=to_pack.device,
|
|
)
|
|
new_c = to_pack.shape[1] // pack_num
|
|
for c in range(new_c):
|
|
for i in range(pack_num):
|
|
# Use -3 as an example, high_position is 11111111,cause bit_or generate errors, so we can't use int4 directly
|
|
packed_col = to_pack[:, c * pack_num + order_map[i]].to(torch.int32)
|
|
packed_col = packed_col & 0x0F
|
|
packed[:, c] = torch.bitwise_or(
|
|
packed[:, c], torch.bitwise_left_shift(packed_col, i * 4)
|
|
)
|
|
elif to_pack.ndim == 0:
|
|
packed = to_pack.to(torch.int32)
|
|
else:
|
|
packed = torch.zeros(
|
|
to_pack.shape[0] // pack_num, dtype=torch.int32, device=to_pack.device
|
|
)
|
|
new_c = to_pack.shape[0] // pack_num
|
|
for c in range(new_c):
|
|
for i in range(pack_num):
|
|
# Use -3 as an example, high_position is 11111111,cause bit_or generate errors, so we can't use int4 directly
|
|
packed_col = to_pack[c * pack_num + order_map[i]]
|
|
packed_col = packed_col & 0x0F
|
|
packed[c] = torch.bitwise_or(
|
|
packed[c], torch.bitwise_left_shift(packed_col, i * 4)
|
|
)
|
|
|
|
return packed.view(torch.uint32)
|