(model_path, wbits, group_size)
| 172 | |
| 173 | |
| 174 | def load_quantized_model(model_path, wbits, group_size): |
| 175 | print(f"Loading quantized model from {model_path}") |
| 176 | |
| 177 | # import pdb;pdb.set_trace() |
| 178 | tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False) |
| 179 | config = AutoConfig.from_pretrained(model_path) |
| 180 | with init_empty_weights(): |
| 181 | model = AutoModelForCausalLM.from_config(config=config,torch_dtype=torch.float16, trust_remote_code=True) |
| 182 | layers = model.model.layers |
| 183 | for i in tqdm(range(len(layers))): |
| 184 | layer = layers[i] |
| 185 | named_linears = get_named_linears(layer, torch.nn.Linear) |
| 186 | for name, module in named_linears.items(): |
| 187 | q_linear = QuantLinear(wbits, group_size, module.in_features,module.out_features,not module.bias is None) |
| 188 | q_linear.to(next(layer.parameters()).device) |
| 189 | set_op_by_name(layer, name, q_linear) |
| 190 | torch.cuda.empty_cache() |
| 191 | gc.collect() |
| 192 | model.tie_weights() |
| 193 | device_map = infer_auto_device_map(model) |
| 194 | print("Loading pre-computed quantized weights...") |
| 195 | load_checkpoint_in_model(model,checkpoint=model_path,device_map=device_map,offload_state_dict=True) |
| 196 | print("Loading pre-computed quantized weights Successfully") |
| 197 | |
| 198 | return model, tokenizer |
| 199 | |
| 200 | __all__ = ["QuantLinear","load_omniq_quantized"] |
no test coverage detected