(
model, w_bit, q_config, quant_type="int", model_path=None
)
| 64 | |
| 65 | @torch.no_grad() |
| 66 | def pseudo_quantize_model_weight( |
| 67 | model, w_bit, q_config, quant_type="int", model_path=None |
| 68 | ): |
| 69 | scale = torch.tensor([1.0], dtype=torch.float32, device="cuda").squeeze(0) |
| 70 | zero = torch.tensor([0.0], dtype=torch.float32, device="cuda").squeeze(0) |
| 71 | if quant_type == "int": |
| 72 | layers = model.model.layers |
| 73 | for i in tqdm(range(len(layers)), desc="pseudo weight quantization..."): |
| 74 | named_linears = get_named_linears(layers[i]) |
| 75 | for n, m in named_linears.items(): |
| 76 | # m.cuda() |
| 77 | m.weight.data = pseudo_quantize_tensor(m.weight.data, n_bit=w_bit, **q_config) |
| 78 | # m.cpu() |
| 79 | elif "lsq" in quant_type: |
| 80 | if quant_type == "lsq-n2f3" or quant_type == "clsq-n2f3" or quant_type == "clsq-n2f2": |
| 81 | model_bnb = AutoModelForCausalLM.from_pretrained( |
| 82 | model_path, |
| 83 | load_in_4bit=True, |
| 84 | device_map='auto', |
| 85 | quantization_config=BitsAndBytesConfig( |
| 86 | load_in_4bit=True, |
| 87 | bnb_4bit_quant_type=quant_type, |
| 88 | ), |
| 89 | torch_dtype=torch.bfloat16 |
| 90 | ) |
| 91 | layers_ori = model.model.layers |
| 92 | layers_bnb = model_bnb.model.layers |
| 93 | for i in tqdm(range(len(layers_ori)), desc=f"pseudo {quant_type} weight quantization..."): |
| 94 | named_linears_ori = get_named_linears(layers_ori[i]) |
| 95 | named_linears_bnb = get_named_bnb_linears(layers_bnb[i]) |
| 96 | for ori_linear, bnb_linear in zip(named_linears_ori.items(), named_linears_bnb.items()): |
| 97 | module_ori = ori_linear[1] |
| 98 | quantizer = bnb_linear[1].weight_quantizer |
| 99 | module_ori.weight.data = quantizer(module_ori.weight.data) |
| 100 | del model_bnb, layers_bnb |
| 101 | elif "lsq" not in quant_type: |
| 102 | if quant_type == "n2f4": |
| 103 | quantizer = SteN2F4Quantizer(q_group_size=q_config["q_group_size"]) |
| 104 | elif quant_type == "nf3": |
| 105 | quantizer = SteNF3Quantizer(bit=3, weight=None, q_group_size=q_config["q_group_size"]) |
| 106 | if quant_type == "nf4": |
| 107 | quantizer = SteNF4Quantizer(bit=4, weight=None, q_group_size=q_config["q_group_size"]) |
| 108 | elif quant_type == "n2f3": |
| 109 | print("quant_type: n2f3") |
| 110 | quantizer = SteN2F3Quantizer(bit=3, weight=None, q_group_size=q_config["q_group_size"]) |
| 111 | elif quant_type == "n2f2": |
| 112 | print("quant_type: n2f2") |
| 113 | quantizer = SteN2F2Quantizer(bit=2, weight=None, q_group_size=32) |
| 114 | elif quant_type == "int3-sym": |
| 115 | quantizer = SteInt3SymQuantizer(q_group_size=q_config["q_group_size"]) |
| 116 | elif quant_type == "int4-asym": |
| 117 | quantizer = SteInt4AsymQuantizer(q_group_size=q_config["q_group_size"]) |
| 118 | elif quant_type == "int3-asym": |
| 119 | quantizer = SteInt3AsymQuantizer(q_group_size=q_config["q_group_size"]) |
| 120 | elif quant_type == "int2-asym": |
| 121 | quantizer = SteInt2AsymQuantizer(q_group_size=q_config["q_group_size"]) |
| 122 | |
| 123 | layers = model.model.layers |
nothing calls this directly
no test coverage detected