| import torch |
| from model.internlm_xcomposer.configuration_internlm_xcomposer2 import InternLMXcomposer2Config |
| from model.internlm_xcomposer.modeling_internlm_xcomposer2 import InternLMXComposer2ForCausalLM |
| from model.internlm_xcomposer.tokenization_internlm_xcomposer2 import InternLMXComposer2Tokenizer |
| from peft import LoraConfig, get_peft_model |
|
|
| from .base import BaseModel |
|
|
| class InternLM(BaseModel): |
| def __init__(self, **kwargs): |
| super().__init__(**kwargs) |
|
|
| def load_model_tokenizer(self): |
| config = InternLMXcomposer2Config.from_pretrained(self.model_path) |
| config.use_cache = False |
| config.max_length = self.training_args.max_length |
| |
| model = InternLMXComposer2ForCausalLM.from_pretrained( |
| self.model_path, |
| config=config, |
| device_map=None, |
| use_caption=self.use_caption |
| ) |
|
|
| if self.data_args.img_size != 336: |
| model.vit.resize_pos() |
|
|
| tokenizer_path = self.model_path if self.lora_args.lora_weight_path == '' \ |
| else self.lora_args.lora_weight_path |
| tokenizer = InternLMXComposer2Tokenizer.from_pretrained( |
| tokenizer_path, |
| padding_side='right', |
| use_fast=False, |
| ) |
|
|
| self.model = model |
| self.tokenizer = tokenizer |
|
|
| def configure_training_args(self): |
| training_args = self.training_args |
| if training_args.fix_vit: |
| self.model.vit.requires_grad_(False) |
| else: |
| self.model.vit.requires_grad_(True) |
| self.model.vit.vision_tower.vision_model.post_layernorm = torch.nn.Identity() |
|
|
| if training_args.fix_sampler or self.use_caption: |
| self.model.vision_proj.requires_grad_(False) |
| else: |
| self.model.vision_proj.requires_grad_(True) |
|
|
| def configure_peft(self): |
| if not self.training_args.use_lora: |
| for name, param in self.model.model.named_parameters(): |
| if 'vision_cross' not in name: |
| param.requires_grad = False |
| return |
| |
| lora_args = self.lora_args |
| if lora_args.lora_type == 'lora': |
| for name, param in self.model.model.named_parameters(): |
| if 'vision_cross' in name: |
| continue |
| param.requires_grad = False |
| lora_config = LoraConfig( |
| r=lora_args.lora_r, |
| lora_alpha=lora_args.lora_alpha, |
| target_modules=lora_args.lora_target_modules, |
| lora_dropout=lora_args.lora_dropout, |
| bias=lora_args.lora_bias, |
| task_type='CAUSAL_LM', |
| ) |
|
|
| self.model = get_peft_model(self.model, lora_config) |
| elif lora_args.lora_type == 'plora': |
| for name, param in self.model.model.named_parameters(): |
| if 'Plora' not in name: |
| param.requires_grad = False |
|
|
| if self.use_caption: |
| if lora_args.lora_type == 'lora': |
| self.model.model.vision_proj.requires_grad_(True) |
| self.model.model.model.tok_embeddings.requires_grad_(True) |
| self.model.model.logit_scale.requires_grad_(True) |
| else: |
| self.model.vision_proj.requires_grad_(True) |
| self.model.model.tok_embeddings.requires_grad_(True) |
| self.model.logit_scale.requires_grad_(True) |
|
|
| |
| for name, param in self.model.model.named_parameters(): |
| if 'vision_cross' in name: |
| param.requires_grad = True |
|
|
| if self.training_args.gradient_checkpointing: |
| self.model.enable_input_require_grads() |
| |
| self.model.model.vit.vision_tower.gradient_checkpointing_enable({"use_reentrant": True}) |
|
|
|
|
| |
|
|
|
|
