| | 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}) |
| |
|
| |
|
| | |
| |
|
| |
|
