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