(num_bits, symmetric, x, group_dim, group_size, quant_type)
| 14 | |
| 15 | |
| 16 | def _quantize(num_bits, symmetric, x, group_dim, group_size, quant_type): |
| 17 | x_float = x.float() |
| 18 | x_reshaped = x_float.reshape( |
| 19 | x.shape[:group_dim] |
| 20 | + (x.shape[group_dim] // group_size, group_size) |
| 21 | + x.shape[group_dim + 1 :] |
| 22 | ) |
| 23 | if symmetric: |
| 24 | signed_max = float(2 ** (num_bits - 1)) - 1 |
| 25 | offset = signed_max if quant_type is torch.uint8 else 0.0 |
| 26 | scale_float = ( |
| 27 | x_reshaped.abs().max(dim=group_dim + 1, keepdim=True).values / signed_max |
| 28 | ) |
| 29 | quantized = ( |
| 30 | torch.round(x_reshaped / scale_float + offset) |
| 31 | .reshape(x.shape) |
| 32 | .to(quant_type) |
| 33 | ) |
| 34 | if num_bits == 4: |
| 35 | quantized = _pack_int8_to_int4(quantized) |
| 36 | return (quantized, scale_float.squeeze(group_dim + 1).to(x.dtype), None) |
| 37 | else: |
| 38 | unsigned_max = float(2 ** num_bits) - 1 |
| 39 | mn = x_reshaped.min(dim=group_dim + 1, keepdim=True).values |
| 40 | mx = x_reshaped.max(dim=group_dim + 1, keepdim=True).values |
| 41 | scale_float = (mx - mn) / unsigned_max |
| 42 | quantized = ( |
| 43 | torch.round((x_reshaped - mn) / scale_float).reshape(x.shape).to(torch.uint8) |
| 44 | ) |
| 45 | if num_bits == 4: |
| 46 | quantized = _pack_int8_to_int4(quantized) |
| 47 | return ( |
| 48 | quantized, |
| 49 | scale_float.squeeze(group_dim + 1).to(x.dtype), |
| 50 | mn.squeeze(group_dim + 1).to(x.dtype), |
| 51 | ) |
| 52 | |
| 53 | class QuantizedLinear(torch.nn.Module): |
| 54 | def __init__( |
no test coverage detected