(model, checkpoint, w_bit, group_size, device)
| 49 | |
| 50 | |
| 51 | def load_awq_model(model, checkpoint, w_bit, group_size, device): |
| 52 | q_config = {"zero_point": True, "q_group_size": group_size} |
| 53 | real_quantize_model_weight(model, w_bit, q_config, init_only=True) |
| 54 | |
| 55 | if hasattr(model.config, "tie_encoder_decoder"): |
| 56 | model.config.tie_encoder_decoder = False |
| 57 | if hasattr(model.config, "tie_word_embeddings"): |
| 58 | model.config.tie_word_embeddings = False |
| 59 | if tinychat.utils.constants.mem_efficient_load: |
| 60 | assert os.path.isdir( |
| 61 | checkpoint |
| 62 | ), "You are in mem_efficient_load mode. \n Please set --load_quant the path to the folder containing all checkpoint files." |
| 63 | model = mem_efficient_load_checkpoint( |
| 64 | model, |
| 65 | checkpoint, |
| 66 | ).to(device) |
| 67 | else: |
| 68 | pbar = tqdm(range(1)) |
| 69 | pbar.set_description("Loading checkpoint") |
| 70 | for i in pbar: |
| 71 | model = load_checkpoint_and_dispatch( |
| 72 | model, |
| 73 | checkpoint, |
| 74 | no_split_module_classes=[ |
| 75 | "OPTDecoderLayer", |
| 76 | "LlamaDecoderLayer", |
| 77 | "BloomBlock", |
| 78 | "MPTBlock", |
| 79 | "DecoderLayer", |
| 80 | ], |
| 81 | ).to(device) |
| 82 | return model |
| 83 | |
| 84 | |
| 85 | def make_quant_linear(module, names, w_bit, groupsize, device, name=""): |
no test coverage detected