| 9 | |
| 10 | |
| 11 | class WeightQuantization(object): |
| 12 | |
| 13 | def __init__(self, mlp_extra_grouping=True, mp_size=1): |
| 14 | self.dense_scales = [] |
| 15 | self.qkv_scales = [] |
| 16 | self.mlp4hh_scales = [] |
| 17 | self.mlph4h_scales = [] |
| 18 | self.mlp_extra_grouping = mlp_extra_grouping |
| 19 | self.mp_size = mp_size |
| 20 | |
| 21 | def quantize_data(self, data, quantize_bits, groups, key=None): |
| 22 | data_groups = torch.split(data.float().view(-1), data.numel() // groups) |
| 23 | max_d = [max(g.max(), g.min().abs()) for g in data_groups] |
| 24 | data_scale = [float(1 << quantize_bits) / (2 * mx + 1e-5) for mx in max_d] |
| 25 | data_int = [(g * s) for g, s in zip(data_groups, data_scale)] |
| 26 | data_int = [ |
| 27 | di.round().clamp(-(1 << (quantize_bits - 1)), (((1 << (quantize_bits - 1)) - 1))) for di in data_int |
| 28 | ] |
| 29 | data_int = torch.cat(data_int).reshape(data.shape) |
| 30 | data_int = data_int.to(torch.int8) |
| 31 | data_scale = torch.cat([s.unsqueeze(0).unsqueeze(0) for s in data_scale]) |
| 32 | return data_int, data_scale |
| 33 | |
| 34 | def is_mlp(self, data, merge_count=1): |
| 35 | return ((self.mp_size *data.shape[0] * merge_count) / data.shape[1] == 4 or \ |
| 36 | (self.mp_size *data.shape[1] * merge_count) / data.shape[0] == 4) |
| 37 | |
| 38 | def is_qkv(self, data): |
| 39 | return ((self.mp_size * data.shape[0]) / data.shape[1] == 3 or \ |
| 40 | (self.mp_size * data.shape[1]) / data.shape[0] == 3) |
| 41 | |
| 42 | def Quantize(self, value_list, quantize_bits, groups, key, merge_dim=0): |
| 43 | if self.mlp_extra_grouping and self.is_mlp(value_list[0], merge_count=len(value_list)): |
| 44 | groups *= 2 |
| 45 | q_scale = [] |
| 46 | index = 0 |
| 47 | for data in value_list: |
| 48 | data_int, data_scale = self.quantize_data(data, quantize_bits, groups, key) |
| 49 | q_scale.append(data_scale) |
| 50 | value_list[index] = data_int |
| 51 | index += 1 |
| 52 | q_scale = (1 / |
| 53 | torch.cat(q_scale, dim=merge_dim).to(get_accelerator().current_device_name()).view(-1).unsqueeze(0)) |
| 54 | if "mlp.dense_4h_to_h.weight" in key: |
| 55 | self.mlp4hh_scales.append(q_scale) |
| 56 | elif "mlp.dense_h_to_4h.weight" in key: |
| 57 | self.mlph4h_scales.append(q_scale) |
| 58 | elif "attention.query_key_value.weight" in key: |
| 59 | self.qkv_scales.append(q_scale) |
| 60 | else: |
| 61 | self.dense_scales.append(q_scale) |
| 62 | return value_list |
| 63 | |
| 64 | def merge_layer_scales(self, layer_scales): |
| 65 | max_dim = max([s.shape[-1] for s in layer_scales]) |
| 66 | layer_scales = [ |
| 67 | torch.cat((s, torch.zeros((1, max_dim - s.shape[-1]), device=get_accelerator().current_device_name())), |
| 68 | dim=-1) if s.shape[-1] < max_dim else s for s in layer_scales |
no outgoing calls
no test coverage detected