(
model, w_bit, q_config,
init_only=False
)
| 72 | |
| 73 | @torch.no_grad() |
| 74 | def real_quantize_model_weight( |
| 75 | model, w_bit, q_config, |
| 76 | init_only=False |
| 77 | ): |
| 78 | from .qmodule import WQLinear |
| 79 | from .pre_quant import get_blocks, get_named_linears, set_op_by_name |
| 80 | assert q_config["zero_point"], "We only support zero_point quantization now." |
| 81 | |
| 82 | layers = get_blocks(model) |
| 83 | for i in tqdm(range(len(layers)), desc="real weight quantization..." + ("(init only)" if init_only else "")): |
| 84 | layer = layers[i] |
| 85 | named_linears = get_named_linears(layer) |
| 86 | # scale_activations(layer) |
| 87 | |
| 88 | for name, module in named_linears.items(): |
| 89 | if init_only: |
| 90 | q_linear = WQLinear.from_linear( |
| 91 | module, w_bit, q_config['q_group_size'], True) |
| 92 | q_linear.to(next(layer.parameters()).device) |
| 93 | set_op_by_name(layer, name, q_linear) |
| 94 | else: |
| 95 | module.cuda() |
| 96 | module.weight.data, scales, zeros = pseudo_quantize_tensor(module.weight.data, n_bit=w_bit, get_scale_zp=True, **q_config) |
| 97 | # scales = scales.t().contiguous() |
| 98 | # zeros = zeros.t().contiguous() |
| 99 | q_linear = WQLinear.from_linear( |
| 100 | module, w_bit, q_config['q_group_size'], False, scales, zeros) |
| 101 | module.cpu() |
| 102 | q_linear.to(next(layer.parameters()).device) |
| 103 | set_op_by_name(layer, name, q_linear) |
| 104 | torch.cuda.empty_cache() |
| 105 | gc.collect() |
| 106 | |
| 107 | torch.cuda.empty_cache() |
| 108 | gc.collect() |
| 109 | |
| 110 | |
| 111 |
no test coverage detected