(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda")
| 24 | |
| 25 | |
| 26 | def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda"): |
| 27 | kwargs = {"device_map": device_map} |
| 28 | |
| 29 | if load_8bit: |
| 30 | kwargs['load_in_8bit'] = True |
| 31 | elif load_4bit: |
| 32 | kwargs['load_in_4bit'] = True |
| 33 | kwargs['quantization_config'] = BitsAndBytesConfig( |
| 34 | load_in_4bit=True, |
| 35 | bnb_4bit_compute_dtype=torch.float16, |
| 36 | bnb_4bit_use_double_quant=True, |
| 37 | bnb_4bit_quant_type='nf4' |
| 38 | ) |
| 39 | else: |
| 40 | kwargs['torch_dtype'] = torch.float16 |
| 41 | |
| 42 | if 'llava' in model_name.lower() or 'intern' in model_name.lower(): |
| 43 | # Load LLaVA model |
| 44 | if 'lora' in model_name.lower() and model_base is None: |
| 45 | warnings.warn('There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, please provide the `model_base` argument. Detailed instruction: https://github.com/haotian-liu/LLaVA#launch-a-model-worker-lora-weights-unmerged.') |
| 46 | if 'lora' in model_name.lower() and model_base is not None: |
| 47 | lora_cfg_pretrained = AutoConfig.from_pretrained(model_path) |
| 48 | tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False) |
| 49 | print('Loading LLaVA from base model...') |
| 50 | model = LlavaLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs) |
| 51 | token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features |
| 52 | if model.lm_head.weight.shape[0] != token_num: |
| 53 | model.lm_head.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype)) |
| 54 | model.model.embed_tokens.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype)) |
| 55 | |
| 56 | print('Loading additional LLaVA weights...') |
| 57 | if os.path.exists(os.path.join(model_path, 'non_lora_trainables.bin')): |
| 58 | non_lora_trainables = torch.load(os.path.join(model_path, 'non_lora_trainables.bin'), map_location='cpu') |
| 59 | else: |
| 60 | # this is probably from HF Hub |
| 61 | from huggingface_hub import hf_hub_download |
| 62 | def load_from_hf(repo_id, filename, subfolder=None): |
| 63 | cache_file = hf_hub_download( |
| 64 | repo_id=repo_id, |
| 65 | filename=filename, |
| 66 | subfolder=subfolder) |
| 67 | return torch.load(cache_file, map_location='cpu') |
| 68 | non_lora_trainables = load_from_hf(model_path, 'non_lora_trainables.bin') |
| 69 | non_lora_trainables = {(k[11:] if k.startswith('base_model.') else k): v for k, v in non_lora_trainables.items()} |
| 70 | if any(k.startswith('model.model.') for k in non_lora_trainables): |
| 71 | non_lora_trainables = {(k[6:] if k.startswith('model.') else k): v for k, v in non_lora_trainables.items()} |
| 72 | model.load_state_dict(non_lora_trainables, strict=False) |
| 73 | |
| 74 | from peft import PeftModel |
| 75 | print('Loading LoRA weights...') |
| 76 | model = PeftModel.from_pretrained(model, model_path) |
| 77 | print('Merging LoRA weights...') |
| 78 | model = model.merge_and_unload() |
| 79 | print('Model is loaded...') |
| 80 | elif model_base is not None: |
| 81 | # this may be mm projector only |
| 82 | print('Loading LLaVA from base model...') |
| 83 | if 'mpt' in model_name.lower(): |
no test coverage detected