(model_path)
| 33 | print("Quantization config:", q_config) |
| 34 | |
| 35 | def build_model_and_enc(model_path): |
| 36 | if not os.path.exists(model_path): # look into ssd |
| 37 | raise FileNotFoundError(f"{model_path} not found!") |
| 38 | print(f"* Building model {model_path}") |
| 39 | |
| 40 | # all hf model |
| 41 | config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) |
| 42 | if "mpt" in config.__class__.__name__.lower(): |
| 43 | enc = AutoTokenizer.from_pretrained( |
| 44 | config.tokenizer_name, trust_remote_code=True |
| 45 | ) |
| 46 | else: |
| 47 | enc = AutoTokenizer.from_pretrained( |
| 48 | model_path, use_fast=False, trust_remote_code=True |
| 49 | ) |
| 50 | |
| 51 | kwargs = {"torch_dtype": torch.float16, "low_cpu_mem_usage": True} |
| 52 | model = AutoModelForCausalLM.from_pretrained( |
| 53 | model_path, config=config, trust_remote_code=True, **kwargs |
| 54 | ) |
| 55 | |
| 56 | model.eval() |
| 57 | |
| 58 | real_quantize_model_weight(model, w_bit=args.w_bit, q_config=q_config) |
| 59 | |
| 60 | dirpath = os.path.dirname(args.dump_quant) |
| 61 | os.makedirs(dirpath, exist_ok=True) |
| 62 | |
| 63 | print(f"Saving the quantized model at {args.dump_quant}...") |
| 64 | torch.save(model.cpu().state_dict(), args.dump_quant) |
| 65 | exit(0) |
| 66 | |
| 67 | return model, enc |
| 68 | |
| 69 | if __name__ == "__main__": |
| 70 | build_model_and_enc(args.model_path) |
no test coverage detected