(w, n_bit=8,
zero_point=True, q_group_size=-1,
inplace=False,
get_scale_zp=False
)
| 23 | |
| 24 | # core quantization method (simulated quantization) |
| 25 | def pseudo_quantize_tensor(w, n_bit=8, |
| 26 | zero_point=True, q_group_size=-1, |
| 27 | inplace=False, |
| 28 | get_scale_zp=False |
| 29 | ): |
| 30 | org_w_shape = w.shape |
| 31 | if q_group_size > 0: |
| 32 | assert org_w_shape[-1] % q_group_size == 0 |
| 33 | w = w.reshape(-1, q_group_size) |
| 34 | elif q_group_size == -1: |
| 35 | w = w.reshape(-1, w.shape[-1]) |
| 36 | assert w.dim() == 2 |
| 37 | if zero_point: |
| 38 | max_val = w.amax(dim=1, keepdim=True) |
| 39 | min_val = w.amin(dim=1, keepdim=True) |
| 40 | max_int = 2 ** n_bit - 1 |
| 41 | min_int = 0 |
| 42 | scales = (max_val - min_val).clamp(min=1e-5) / max_int |
| 43 | zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int) |
| 44 | else: # we actually never used this |
| 45 | assert min_val is None |
| 46 | max_val = w.abs().amax(dim=1, keepdim=True) |
| 47 | max_val = max_val.clamp(min=1e-5) |
| 48 | max_int = 2 ** (n_bit - 1) - 1 |
| 49 | min_int = - 2 ** (n_bit - 1) |
| 50 | scales = max_val / max_int |
| 51 | zeros = 0 |
| 52 | |
| 53 | assert torch.isnan(scales).sum() == 0 |
| 54 | assert torch.isnan(w).sum() == 0 |
| 55 | |
| 56 | if inplace: |
| 57 | ((w.div_(scales).round_().add_(zeros)).clamp_( |
| 58 | min_int, max_int).sub_(zeros)).mul_(scales) |
| 59 | else: |
| 60 | w = (torch.clamp(torch.round(w / scales) + |
| 61 | zeros, min_int, max_int) - zeros) * scales |
| 62 | assert torch.isnan(w).sum() == 0 |
| 63 | |
| 64 | w = w.reshape(org_w_shape) |
| 65 | |
| 66 | if get_scale_zp: |
| 67 | return w, scales.view(w.shape[0], -1), zeros.view(w.shape[0], -1) |
| 68 | else: |
| 69 | return w |
| 70 | |
| 71 | |
| 72 |
no outgoing calls
no test coverage detected