| |
| from collections import OrderedDict |
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| from mmengine.config import Config, ConfigDict |
| from mmengine.model import BaseModel |
| from peft import get_peft_model, prepare_model_for_kbit_training |
| from transformers import (AutoTokenizer, BitsAndBytesConfig, GenerationConfig) |
| from xtuner.registry import BUILDER |
| from xtuner.model.utils import find_all_linear_names, get_peft_model_state_dict, guess_load_checkpoint, make_inputs_require_grad |
|
|
|
|
|
|
| class QwenVL2(BaseModel): |
|
|
| def __init__(self, |
| model_path, |
| freeze_llm=False, |
| freeze_visual_encoder=False, |
| llm_lora=None, |
| visual_encoder_lora=None, |
| quantization_vit=False, |
| quantization_llm=False, |
| pretrained_pth=None, |
| |
| special_tokens=None, |
| ): |
| super().__init__() |
| self.freeze_llm = freeze_llm |
| self.freeze_visual_encoder = freeze_visual_encoder |
| self.use_llm_lora = llm_lora is not None |
| self.use_visual_encoder_lora = visual_encoder_lora is not None |
| self.quantization_vit = quantization_vit |
| self.quantization_llm = quantization_llm |
| if quantization_vit: |
| assert visual_encoder_lora is not None |
| if quantization_llm: |
| assert quantization_llm and llm_lora is not None |
|
|
| |
|
|
| if quantization_vit is False and quantization_llm is False: |
| quantization = None |
| else: |
| llm_int8_skip_modules = ['mlp1'] |
| if quantization_llm and not quantization_vit: |
| llm_int8_skip_modules.append('vision_model') |
|
|
| if quantization_vit and not quantization_llm: |
| llm_int8_skip_modules.append('model') |
|
|
| quantization_config = dict( |
| type=BitsAndBytesConfig, |
| llm_int8_skip_modules=llm_int8_skip_modules, |
| load_in_4bit=True, |
| load_in_8bit=False, |
| llm_int8_threshold=6.0, |
| llm_int8_has_fp16_weight=False, |
| bnb_4bit_compute_dtype=torch.float16, |
| bnb_4bit_use_double_quant=True, |
| bnb_4bit_quant_type='nf4') |
| quantization_clazz = quantization_config.pop('type') |
| quantization = quantization_clazz(**quantization_config) |
|
|
| from .qwen2_vl import Qwen2VLForConditionalGeneration |
| self.model = Qwen2VLForConditionalGeneration.from_pretrained( |
| model_path, |
| torch_dtype=torch.bfloat16, |
| quantization_config=quantization, |
| |
| trust_remote_code=True) |
|
|
| tokenizer = AutoTokenizer.from_pretrained( |
| model_path, trust_remote_code=True) |
| self.tokenizer = tokenizer |
| if special_tokens is not None: |
| self._add_special_tokens(special_tokens) |
|
|
| if self.freeze_llm: |
| self.model.model.requires_grad_(False) |
| if self.freeze_visual_encoder: |
| self.model.visual.requires_grad_(False) |
|
|
| if hasattr(self.model.model, 'enable_input_require_grads'): |
| self.model.model.enable_input_require_grads() |
| else: |
| self.model.model.get_input_embeddings( |
| ).register_forward_hook(make_inputs_require_grad) |
|
|
| self.gradient_checkpointing_enable() |
|
|
| if self.use_llm_lora: |
| self._prepare_llm_for_lora(llm_lora) |
|
|
| if self.use_visual_encoder_lora: |
| self._prepare_visual_encoder_for_lora(visual_encoder_lora) |
|
|
| if pretrained_pth is not None: |
| pretrained_state_dict = guess_load_checkpoint(pretrained_pth) |
|
|
| self.load_state_dict(pretrained_state_dict, strict=False) |
| print(f'Load pretrained weight from {pretrained_pth}') |
|
|
| self._count = 0 |
|
|
|
|
| def _add_special_tokens(self, special_tokens): |
| num_new_tokens = self.tokenizer.add_tokens(special_tokens, special_tokens=True) |
| if num_new_tokens > 0: |
| |
| self.model.resize_token_embeddings(len(self.tokenizer)) |
|
|
| def _post_init(self, fast_pool_size=4, fast_pool=True): |
| if fast_pool: |
| self.fast_pool = nn.AdaptiveAvgPool2d((fast_pool_size, fast_pool_size)) |
| return |
|
|
| def _parse_lora_config(self, lora_config): |
| if isinstance(lora_config, dict) or isinstance( |
| lora_config, Config) or isinstance(lora_config, ConfigDict): |
| lora_config = BUILDER.build(lora_config) |
| return lora_config |
|
|
| def _prepare_llm_for_lora(self, |
| lora_config, |
| use_activation_checkpointing=True): |
| lora_config = self._parse_lora_config(lora_config) |
| self.model.model = prepare_model_for_kbit_training( |
| self.model.model, use_activation_checkpointing) |
| if lora_config.target_modules is None: |
| modules = find_all_linear_names(self.model.model) |
| lora_config.target_modules = modules |
| self.model.model = get_peft_model(self.model.model, lora_config) |
|
|
| def _prepare_visual_encoder_for_lora(self, lora_config): |
| lora_config = self._parse_lora_config(lora_config) |
| if lora_config.target_modules is None: |
| modules = find_all_linear_names(self.model.vision_model) |
| lora_config.target_modules = modules |
| self.model.vision_model = get_peft_model(self.model.vision_model, |
| lora_config) |
|
|
| def gradient_checkpointing_enable(self): |
| self.activation_checkpointing_enable() |
|
|
| def activation_checkpointing_enable(self): |
| self.model.model.gradient_checkpointing_enable() |
|
|
| def gradient_checkpointing_disable(self): |
| self.activation_checkpointing_disable() |
|
|
| def activation_checkpointing_disable(self): |
| self.model.model.gradient_checkpointing_disable() |
|
|
| def state_dict(self, *args, **kwargs): |
| state_dict = super().state_dict(*args, **kwargs) |
| to_return = OrderedDict() |
| |
| if self.use_visual_encoder_lora: |
| to_return.update( |
| get_peft_model_state_dict( |
| self.model.vision_model, state_dict=state_dict)) |
| elif not self.freeze_visual_encoder: |
| to_return.update({ |
| k: v |
| for k, v in state_dict.items() if 'model.visual.' in k |
| }) |
| |
| if self.use_llm_lora: |
| to_return.update( |
| get_peft_model_state_dict( |
| self.model.model, state_dict=state_dict)) |
| elif not self.freeze_llm: |
| to_return.update({ |
| k: v |
| for k, v in state_dict.items() if 'model.model.' in k |
| }) |
| |
| to_return.update( |
| {k: v |
| for k, v in state_dict.items() if 'model.lm_head.' in k}) |
| return to_return |
|
|
| def init_weights(self): |
| pass |
|
|
| def forward(self, data, data_samples=None, mode='loss'): |
| has_image = data.get('image_grid_thw', None) is not None |
| if has_image: |
| pixel_values = data['pixel_values'][0] |
| image_grid_thw = data['image_grid_thw'][0] |
| else: |
| pixel_values = None |
| image_grid_thw = None |
| input_ids = data['input_ids'] |
| |
| position_ids = None |
| attention_mask = data['attention_mask'] |
|
|
| labels = data['labels'] |
| use_cache = False |
|
|
| |
| outputs = self.model(input_ids=input_ids, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| pixel_values=pixel_values, |
| image_grid_thw=image_grid_thw, |
| labels=labels, |
| use_cache=use_cache, |
| output_hidden_states=True, |
| ) |
| return outputs |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| pixel_values: Optional[torch.FloatTensor] = None, |
| input_ids: Optional[torch.FloatTensor] = None, |
| attention_mask: Optional[torch.LongTensor] = None, |
| generation_config: Optional[GenerationConfig] = None, |
| output_hidden_states: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| position_ids = None, |
| image_grid_thw = None, |
| **generate_kwargs, |
| ) -> torch.LongTensor: |
| device = self.model.device |
| pixel_values = pixel_values.to(device) |
| image_grid_thw = image_grid_thw.to(device) |
| attention_mask = attention_mask.to(device) |
| input_ids = input_ids.to(device) |
| outputs = self.model.generate( |
| input_ids=input_ids.to(device), |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| pixel_values=pixel_values, |
| image_grid_thw=image_grid_thw, |
| generation_config=generation_config, |
| output_hidden_states=output_hidden_states, |
| return_dict=return_dict, |
| use_cache=True, |
| **generate_kwargs, |
| ) |
|
|
| return outputs |
|
|
