| 3 | from quantizer import * |
| 4 | |
| 5 | def convertModelToQuant(model, |
| 6 | modules_to_not_convert=["lm_head"], |
| 7 | current_key_name=None, |
| 8 | has_been_replaced=False, |
| 9 | compute_dtype=torch.bfloat16, |
| 10 | quant_type="clsq-n2f3", |
| 11 | q_group_size=128): |
| 12 | for name, module in model.named_children(): |
| 13 | if current_key_name is None: |
| 14 | current_key_name = [] |
| 15 | current_key_name.append(name) |
| 16 | if (isinstance(module, nn.Linear)) and name not in modules_to_not_convert: |
| 17 | in_features = module.in_features |
| 18 | out_features = module.out_features |
| 19 | weight = module.weight |
| 20 | bias = module.bias |
| 21 | |
| 22 | model._modules[name] = QLinear( |
| 23 | in_features, |
| 24 | out_features, |
| 25 | module.bias is not None, |
| 26 | compute_dtype=compute_dtype, |
| 27 | quant_type=quant_type, |
| 28 | q_group_size=q_group_size |
| 29 | ) |
| 30 | |
| 31 | model._modules[name].weight = weight |
| 32 | model._modules[name].bias = bias |
| 33 | has_been_replaced = True |
| 34 | # Store the module class in case we need to transpose the weight later |
| 35 | model._modules[name].source_cls = type(module) |
| 36 | if len(list(module.children())) > 0: |
| 37 | _, has_been_replaced = convertModelToQuant( |
| 38 | module, |
| 39 | modules_to_not_convert, |
| 40 | current_key_name, |
| 41 | has_been_replaced, |
| 42 | compute_dtype, |
| 43 | quant_type, |
| 44 | q_group_size |
| 45 | ) |
| 46 | # Remove the last key for recursion |
| 47 | current_key_name.pop(-1) |
| 48 | return model, has_been_replaced |
| 49 | |
| 50 | class QLinear(nn.Linear): |
| 51 | def __init__(self, input_features, output_features, bias=True, compute_dtype=torch.bfloat16, quant_type="ste-n2f3", q_group_size=128, device=None): |