simplecloud/VidChain-exercise

1# Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:2# Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:3#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li4#5#    Licensed under the Apache License, Version 2.0 (the "License");6#    you may not use this file except in compliance with the License.7#    You may obtain a copy of the License at8#9#        http://www.apache.org/licenses/LICENSE-2.010#11#    Unless required by applicable law or agreed to in writing, software12#    distributed under the License is distributed on an "AS IS" BASIS,13#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14#    See the License for the specific language governing permissions and15#    limitations under the License.16 17import os18root_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..")19from dataclasses import dataclass, field20import logging21import pathlib22from typing import Dict, Optional, Sequence, List23 24import torch25import transformers26import sys27sys.path.append(root_dir)28from vtimellm import conversation as conversation_lib29from vtimellm.train.vtimellm_trainer import VTimeLLMTrainer30from vtimellm.train.dataset import make_supervised_data_module, DataArguments31from vtimellm.model import VTimeLLMLlamaForCausalLM, VTimeLLMChatGLMForCausalLM32from vtimellm.model.builder import load_lora33from vtimellm.mm_utils import print_trainable_parameters34 35local_rank = None36 37def rank0_print(*args):38    if local_rank == 0:39        print(*args)40 41@dataclass42class ModelArguments:43    model_name_or_path: Optional[str] = field(default="checkpoints/vtimellm/vicuna-7b-v1.5")44    stage2_path: Optional[str] = field(default='checkpoints/vtimellm/vtimellm-vicuna-v1-5-7b-stage2')45    stage3_path: Optional[str] = field(default='checkpoints/vtimellm/vtimellm-vicuna-v1-5-7b-stage3')46    version: Optional[str] = field(default="v0")47    tune_mm_mlp_adapter: bool = field(default=False)48    pretrain_mm_mlp_adapter: Optional[str] = field(default=None)49 50    ##################################################################################51    # Connector Arguments52    mm_projector_type: Optional[str] = field(default='stc_connector')53    tune_mm_mlp_adapter: bool = field(default=False)54 55    # Vision tower Arguments56    vision_tower: Optional[str] = field(default=None)57    mm_vision_select_layer: Optional[int] = field(default=-2)58    mm_vision_select_feature: Optional[str] = field(default="patch")59    # Other Arguments60    mm_use_im_start_end: bool = field(default=False)61    mm_use_im_patch_token: bool = field(default=False)62    pretrain_model_name_or_path: Optional[str] = field(default=None, metadata={"help": "To train from previously trained checkpoints. E.g, further fine-tuning based on the finetuned version of the whole model."})63    ###############################################################################64 65 66@dataclass67class TrainingArguments(transformers.TrainingArguments):68    training_stage: int = field(default=2)69    finetuning: bool = field(default=False)70    cache_dir: Optional[str] = field(default=None)71    optim: str = field(default="adamw_torch")72    remove_unused_columns: bool = field(default=False)73    freeze_mm_mlp_adapter: bool = field(default=False)74    model_max_length: int = field(75        default=512,76        metadata={77            "help":78            "Maximum sequence length. Sequences will be right padded (and possibly truncated)."79        },80    )81    double_quant: bool = field(82        default=True,83        metadata={"help": "Compress the quantization statistics through double quantization."}84    )85    quant_type: str = field(86        default="nf4",87        metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}88    )89    bits: int = field(90        default=16,91        metadata={"help": "How many bits to use."}92    )93    lora_enable: bool = False94    lora_r: int = 6495    lora_alpha: int = 1696    lora_dropout: float = 0.0597    lora_weight_path: str = ""98    lora_bias: str = "none"99 100 101 102 103def maybe_zero_3(param, ignore_status=False, name=None):104    from deepspeed import zero105    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus106    if hasattr(param, "ds_id"):107        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:108            if not ignore_status:109                logging.warning(f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}")110        with zero.GatheredParameters([param]):111            param = param.data.detach().cpu().clone()112    else:113        param = param.detach().cpu().clone()114    return param115 116 117# Borrowed from peft.utils.get_peft_model_state_dict118def get_peft_state_maybe_zero_3(named_params, bias):119    if bias == "none":120        to_return = {k: t for k, t in named_params if "lora_" in k}121    elif bias == "all":122        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}123    elif bias == "lora_only":124        to_return = {}125        maybe_lora_bias = {}126        lora_bias_names = set()127        for k, t in named_params:128            if "lora_" in k:129                to_return[k] = t130                bias_name = k.split("lora_")[0] + "bias"131                lora_bias_names.add(bias_name)132            elif "bias" in k:133                maybe_lora_bias[k] = t134        for k, t in maybe_lora_bias:135            if bias_name in lora_bias_names:136                to_return[bias_name] = t137    else:138        raise NotImplementedError139    to_return = {k: maybe_zero_3(v, name=k) for k, v in to_return.items()}140    return to_return141 142 143def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True):144    to_return = {k: t for k, t in named_params if "lora_" not in k}145    if require_grad_only:146        to_return = {k: t for k, t in to_return.items() if t.requires_grad}147    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}148    return to_return149 150 151def get_mm_adapter_state_maybe_zero_3(named_params, keys_to_match):152    to_return = {k: t for k, t in named_params if any(key_match in k for key_match in keys_to_match)}153    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}154    return to_return155 156 157def find_all_linear_names(model):158    cls = torch.nn.Linear159    lora_module_names = set()160    for name, module in model.named_modules():161        if isinstance(module, cls):162            names = name.split('.')163            lora_module_names.add(names[0] if len(names) == 1 else names[-1])164 165 166    if 'lm_head' in lora_module_names: # needed for 16-bit167        lora_module_names.remove('lm_head')168    return list(lora_module_names)169 170 171def safe_save_model_for_hf_trainer(trainer: transformers.Trainer,172                                   output_dir: str):173    """Collects the state dict and dump to disk."""174 175    if getattr(trainer.args, "tune_mm_mlp_adapter", False):176        # Only save Adapter177        keys_to_match = ['mm_projector']178        if getattr(trainer.args, "use_im_start_end", False):179            keys_to_match.extend(['embed_tokens', 'embed_in'])180 181        weight_to_save = get_mm_adapter_state_maybe_zero_3(trainer.model.named_parameters(), keys_to_match)182        trainer.model.config.save_pretrained(output_dir)183 184        current_folder = output_dir.split('/')[-1]185        parent_folder = os.path.dirname(output_dir)186        if trainer.args.local_rank == 0 or trainer.args.local_rank == -1:187            if current_folder.startswith('checkpoint-'):188                mm_projector_folder = os.path.join(parent_folder, "mm_projector")189                os.makedirs(mm_projector_folder, exist_ok=True)190                torch.save(weight_to_save, os.path.join(mm_projector_folder, f'{current_folder}.bin'))191            else:192                torch.save(weight_to_save, os.path.join(output_dir, f'mm_projector.bin'))193        return194 195    if trainer.deepspeed:196        torch.cuda.synchronize()197        trainer.save_model(output_dir)198        return199 200    state_dict = trainer.model.state_dict()201    if trainer.args.should_save:202        cpu_state_dict = {203            key: value.cpu()204            for key, value in state_dict.items()205        }206        del state_dict207        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa208 209 210def smart_tokenizer_and_embedding_resize(211    special_tokens_dict: Dict,212    tokenizer: transformers.PreTrainedTokenizer,213    model: transformers.PreTrainedModel,214):215    """Resize tokenizer and embedding.216 217    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.218    """219    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)220    model.resize_token_embeddings(len(tokenizer))221 222    if num_new_tokens > 0:223        input_embeddings = model.get_input_embeddings().weight.data224        output_embeddings = model.get_output_embeddings().weight.data225 226        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(227            dim=0, keepdim=True)228        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(229            dim=0, keepdim=True)230 231        input_embeddings[-num_new_tokens:] = input_embeddings_avg232        output_embeddings[-num_new_tokens:] = output_embeddings_avg233 234 235def train():236    global local_rank237 238    parser = transformers.HfArgumentParser(239        (ModelArguments, DataArguments, TrainingArguments))240    model_args, data_args, training_args = parser.parse_args_into_dataclasses()241    local_rank = training_args.local_rank242    compute_dtype = (torch.float16 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))243 244 245    bnb_model_from_pretrained_args = {}246    if training_args.bits in [4, 8]:247        from transformers import BitsAndBytesConfig248        bnb_model_from_pretrained_args.update(dict(249            device_map={"": training_args.device},250            load_in_4bit=training_args.bits == 4,251            load_in_8bit=training_args.bits == 8,252            quantization_config=BitsAndBytesConfig(253                load_in_4bit=training_args.bits == 4,254                load_in_8bit=training_args.bits == 8,255                llm_int8_threshold=6.0,256                llm_int8_has_fp16_weight=False,257                bnb_4bit_compute_dtype=compute_dtype,258                bnb_4bit_use_double_quant=training_args.double_quant,259                bnb_4bit_quant_type=training_args.quant_type # {'fp4', 'nf4'}260            )261        ))262 263    if 'chatglm' in model_args.model_name_or_path:264        model = VTimeLLMChatGLMForCausalLM.from_pretrained(265            model_args.model_name_or_path, empty_init=False, device='cuda'266        )267 268    elif 'VideoLLaMA2' in model_args.model_name_or_path:269        config = transformers.AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=True)270        config._attn_implementation = 'flash_attention_2'271        model = Videollama2MistralForCausalLM.from_pretrained(272            model_args.model_name_or_path,273            config=config,274            cache_dir=training_args.cache_dir,275            torch_dtype=(torch.bfloat16 if training_args.bf16 else None),276            do_sample=True,277            **bnb_model_from_pretrained_args278        )279 280    else:281        model = VTimeLLMLlamaForCausalLM.from_pretrained(282            model_args.model_name_or_path,283            cache_dir=training_args.cache_dir,284            **bnb_model_from_pretrained_args285        )286    model.config.use_cache = False287 288    if training_args.bits in [4, 8]:289        from peft import prepare_model_for_kbit_training290        model.config.torch_dtype=(torch.float32 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))291        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=training_args.gradient_checkpointing)292 293    if training_args.gradient_checkpointing:294        if hasattr(model, "enable_input_require_grads"):295            model.enable_input_require_grads()296        else:297            def make_inputs_require_grad(module, input, output):298                output.requires_grad_(True)299            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)300 301 302    if 'chatglm' in model_args.model_name_or_path:303        tokenizer = transformers.AutoTokenizer.from_pretrained(304            model_args.model_name_or_path,305            trust_remote_code=True306        )307    else:308        tokenizer = transformers.AutoTokenizer.from_pretrained(309            model_args.model_name_or_path,310            cache_dir=training_args.cache_dir,311            model_max_length=training_args.model_max_length,312            padding_side="right",313            use_fast=False,314        )315        tokenizer.pad_token = tokenizer.unk_token316 317    if training_args.lora_enable:318        from peft import LoraConfig, get_peft_model319        lora_config = LoraConfig(320            r=training_args.lora_r,321            lora_alpha=training_args.lora_alpha,322            target_modules=find_all_linear_names(model),323            lora_dropout=training_args.lora_dropout,324            bias=training_args.lora_bias,325            task_type="CAUSAL_LM",326        )327        if training_args.bits == 16:328            if training_args.bf16:329                model.to(torch.bfloat16)330            if training_args.fp16:331                model.to(torch.float16)332 333 334        model = model.cuda()335 336        # print_trainable_parameters(model)337        if training_args.training_stage == 3:338            model.get_model().initialize_vision_modules(model_args)339 340            model = load_lora(model, model_args.stage2_path)341            rank0_print('Merging stage 2 LoRA weights...')342            model = model.merge_and_unload()343            344            if training_args.finetuning:345                # ======================================================= #346                # including stage 4 training (Finetuning for the benchmark)347                rank0_print("*" * 90)348                rank0_print("Preparing for stage 4 (finetuning)")349 350                model = load_lora(model, model_args.stage3_path)351                rank0_print('Merging stage 3 LoRA weights...')352                model = model.merge_and_unload()353                rank0_print("*" * 90)354 355            rank0_print("Adding LoRA adapters...")356            model = get_peft_model(model, lora_config)357 358        else:359            rank0_print("Adding LoRA adapters...")360            model = get_peft_model(model, lora_config)361 362 363        print_trainable_parameters(model)364 365    366    if model_args.version in conversation_lib.conv_templates:367        conversation_lib.default_conversation = conversation_lib.conv_templates[model_args.version]368    else:369        conversation_lib.default_conversation = conversation_lib.conv_templates["vicuna_v1"]370 371 372    model.config.tune_mm_mlp_adapter = training_args.tune_mm_mlp_adapter = model_args.tune_mm_mlp_adapter373    model.config.freeze_mm_mlp_adapter = training_args.freeze_mm_mlp_adapter374 375    if training_args.training_stage != 3:376        model.get_model().initialize_vision_modules(model_args=model_args)377 378 379        if model_args.tune_mm_mlp_adapter:380            model.requires_grad_(False)381            for p in model.get_model().mm_projector.parameters():382                p.requires_grad = True383 384        385        if training_args.freeze_mm_mlp_adapter:386            for p in model.get_model().mm_projector.parameters():387                p.requires_grad = False388 389    if training_args.bits in [4, 8]:390        model.get_model().mm_projector.to(dtype=compute_dtype, device=training_args.device)391 392    if training_args.bits in [4, 8]:393        from peft.tuners.lora import LoraLayer394        for name, module in model.named_modules():395            if isinstance(module, LoraLayer):396                if training_args.bf16:397                    module = module.to(torch.bfloat16)398            if 'norm' in name:399                module = module.to(torch.float32)400            if 'lm_head' in name or 'embed_tokens' in name:401                if hasattr(module, 'weight'):402                    if training_args.bf16 and module.weight.dtype == torch.float32:403                        module = module.to(torch.bfloat16)404 405    data_module = make_supervised_data_module(tokenizer=tokenizer,406                                              data_args=data_args)407    trainer = VTimeLLMTrainer(model=model,408                    tokenizer=tokenizer,409                    args=training_args,410                    **data_module)411 412    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):413        trainer.train(resume_from_checkpoint=True)414    else:415        trainer.train()416    trainer.save_state()417 418    model.config.use_cache = True419 420    if training_args.lora_enable:421        state_dict = get_peft_state_maybe_zero_3(422            model.named_parameters(), training_args.lora_bias423        )424        non_lora_state_dict = get_peft_state_non_lora_maybe_zero_3(425            model.named_parameters()426        )427        if training_args.local_rank == 0 or training_args.local_rank == -1:428            model.config.save_pretrained(training_args.output_dir)429            model.save_pretrained(training_args.output_dir, state_dict=state_dict)430            torch.save(non_lora_state_dict, os.path.join(training_args.output_dir, 'non_lora_trainables.bin'))431    else:432        safe_save_model_for_hf_trainer(trainer=trainer,433                                       output_dir=training_args.output_dir)434 435 436if __name__ == "__main__":437    train()438