(model, checkpoint, w_bit, group_size, device)
| 127 | |
| 128 | |
| 129 | def load_awq_llama_fast(model, checkpoint, w_bit, group_size, device): |
| 130 | layers = find_layers(model) |
| 131 | for name in ["lm_head"]: |
| 132 | if name in layers: |
| 133 | del layers[name] |
| 134 | make_quant_linear(model, layers, w_bit, group_size, device) |
| 135 | del layers |
| 136 | |
| 137 | if tinychat.utils.constants.mem_efficient_load: |
| 138 | # TODO: mem-efficient load for llama |
| 139 | assert os.path.isdir( |
| 140 | checkpoint |
| 141 | ), "You are in mem_efficient_load mode. \n Please set --load_quant the path to the folder containing all checkpoint files." |
| 142 | model = mem_efficient_load_checkpoint( |
| 143 | model, |
| 144 | checkpoint, |
| 145 | ) |
| 146 | else: |
| 147 | pbar = tqdm(range(1)) |
| 148 | pbar.set_description("Loading checkpoint") |
| 149 | for i in pbar: |
| 150 | if checkpoint.endswith(".safetensors"): |
| 151 | from safetensors.torch import load_file as safe_load |
| 152 | |
| 153 | model.load_state_dict(safe_load(checkpoint)) |
| 154 | else: |
| 155 | model.load_state_dict(torch.load(checkpoint)) |
| 156 | |
| 157 | # autotune_warmup(model) |
| 158 | |
| 159 | return model.to(device) |
| 160 | |
| 161 | |
| 162 |
no test coverage detected