(model_uri, lora_uri, cache_dir, bits = 32, groupsize=-1, device_map="auto", seqlen=2048, max_memory=None)
| 171 | |
| 172 | |
| 173 | def load_llama_model_lora(model_uri, lora_uri, cache_dir, bits = 32, groupsize=-1, device_map="auto", seqlen=2048, max_memory=None): |
| 174 | import accelerate |
| 175 | from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer |
| 176 | |
| 177 | if max_memory is None: |
| 178 | max_memory = {0: '24Gib', 'cpu': '48Gib'} |
| 179 | |
| 180 | print(Style.BRIGHT + Fore.CYAN + "Loading Model ...") |
| 181 | t0 = time.time() |
| 182 | |
| 183 | with accelerate.init_empty_weights(): |
| 184 | config = LlamaConfig.from_pretrained(model_uri, cache_dir=cache_dir) |
| 185 | model = LlamaForCausalLM(config) |
| 186 | model = model.eval() |
| 187 | layers = find_layers(model) |
| 188 | for name in ['lm_head']: |
| 189 | if name in layers: |
| 190 | del layers[name] |
| 191 | make_quant(model, layers, groupsize=groupsize, bits = bits) |
| 192 | |
| 193 | accelerate.load_checkpoint_in_model( |
| 194 | model, |
| 195 | checkpoint=hf_hub_download(repo_id=model_uri, filename="pytorch_model.bin", cache_dir=cache_dir), |
| 196 | device_map={'': 'cpu'} |
| 197 | ) |
| 198 | |
| 199 | model.seqlen = seqlen |
| 200 | # rotary_emb fix |
| 201 | for n, m in model.named_modules(): |
| 202 | if 'rotary_emb' in n: |
| 203 | cos_cached = m.cos_cached.clone().cpu() |
| 204 | sin_cached = m.sin_cached.clone().cpu() |
| 205 | break |
| 206 | |
| 207 | from peft import PeftModel |
| 208 | from peft_tuners_lora import LinearLowbitLt |
| 209 | |
| 210 | _ = hf_hub_download(repo_id=lora_uri, filename="adapter_config.json", cache_dir=cache_dir) |
| 211 | lora_model = hf_hub_download(repo_id=lora_uri, filename="adapter_model.bin", cache_dir=cache_dir) |
| 212 | lora_path = Path(lora_model).parent |
| 213 | |
| 214 | model = PeftModel.from_pretrained(model, lora_path, device_map={'': 'cpu'}, torch_dtype=torch.float32, is_trainable=True) |
| 215 | print(Style.BRIGHT + Fore.GREEN + '{} Lora Applied.'.format(lora_path)) |
| 216 | |
| 217 | model.seqlen = seqlen |
| 218 | |
| 219 | print('Apply half ...') |
| 220 | for n, m in model.named_modules(): |
| 221 | if isinstance(m, QuantLinear) or ((lora_path is not None) and isinstance(m, LinearLowbitLt)): |
| 222 | m.qscales_scales=m.qscales_scales.half() |
| 223 | m.qscales_zeros=m.qscales_zeros.half() |
| 224 | if m.bias is not None: |
| 225 | m.bias = m.bias.half() |
| 226 | |
| 227 | print('Dispatching model ...') |
| 228 | device_map = accelerate.infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=["LlamaDecoderLayer"]) |
| 229 | model = accelerate.dispatch_model(model, device_map=device_map, offload_buffers=True, main_device=0) |
| 230 | torch.cuda.empty_cache() |
no test coverage detected