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 logging21from copy import deepcopy22import pathlib23from typing import Dict, Optional, Sequence, List24 25import torch26import transformers27import sys28sys.path.append(root_dir)29from vtimellm import conversation as conversation_lib30from vtimellm.train.vtimellm_trainer import VTimeLLMTrainer31from vtimellm.train.dataset import DataArguments, make_supervised_dpo_data_module32from vtimellm.model import VTimeLLMLlamaForCausalLM, VTimeLLMChatGLMForCausalLM33from vtimellm.model.builder import load_lora34from vtimellm.mm_utils import print_trainable_parameters35from trl.trl.trainer import DPOTrainer36 37import torch.multiprocessing38torch.multiprocessing.set_sharing_strategy('file_system')39 40local_rank = None41 42def rank0_print(*args):43    if local_rank == 0:44        print(*args)45 46@dataclass47class ModelArguments:48    model_name_or_path: Optional[str] = field(default="checkpoints/vtimellm/vicuna-7b-v1.5")49    stage2_path: Optional[str] = field(default='checkpoints/vtimellm/vtimellm-vicuna-v1-5-7b-stage2')50    stage3_path: Optional[str] = field(default='checkpoints/vtimellm/vtimellm-vicuna-v1-5-7b-stage3')51    stage4_path: Optional[str] = field(default='checkpoints/vtimellm/vtimellm-vicuna-v1-5-7b-activitynet-stage4')52    version: Optional[str] = field(default="v1")53    tune_mm_mlp_adapter: bool = field(default=False)54    pretrain_mm_mlp_adapter: Optional[str] = field(default=None)55 56 57@dataclass58class TrainingArguments(transformers.TrainingArguments):59    training_stage: int = field(default=3)60    finetuning: bool = field(default=True)61    cache_dir: Optional[str] = field(default=None)62    optim: str = field(default="adamw_torch")63    remove_unused_columns: bool = field(default=False)64    freeze_mm_mlp_adapter: bool = field(default=False)65    model_max_length: int = field(66        default=512,67        metadata={68            "help":69            "Maximum sequence length. Sequences will be right padded (and possibly truncated)."70        },71    )72    double_quant: bool = field(73        default=True,74        metadata={"help": "Compress the quantization statistics through double quantization."}75    )76    quant_type: str = field(77        default="nf4",78        metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}79    )80    bits: int = field(81        default=16,82        metadata={"help": "How many bits to use."}83    )84    lora_enable: bool = False85    lora_r: int = 6486    lora_alpha: int = 12887    lora_dropout: float = 0.0588    lora_weight_path: str = ""89    lora_bias: str = "none"90 91    beta: float = 0.192    train4dpo: bool = field(default=True)93    generate_during_eval: bool = field(default=False)94    dpo_alpha: float = field(default=1.0)95    gamma: float = 0.196 97 98 99def maybe_zero_3(param, ignore_status=False, name=None):100    from deepspeed import zero101    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus102    if hasattr(param, "ds_id"):103        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:104            if not ignore_status:105                logging.warning(f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}")106        with zero.GatheredParameters([param]):107            param = param.data.detach().cpu().clone()108    else:109        param = param.detach().cpu().clone()110    return param111 112 113# Borrowed from peft.utils.get_peft_model_state_dict114def get_peft_state_maybe_zero_3(named_params, bias):115    if bias == "none":116        to_return = {k: t for k, t in named_params if "lora_" in k}117    elif bias == "all":118        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}119    elif bias == "lora_only":120        to_return = {}121        maybe_lora_bias = {}122        lora_bias_names = set()123        for k, t in named_params:124            if "lora_" in k:125                to_return[k] = t126                bias_name = k.split("lora_")[0] + "bias"127                lora_bias_names.add(bias_name)128            elif "bias" in k:129                maybe_lora_bias[k] = t130        for k, t in maybe_lora_bias:131            if bias_name in lora_bias_names:132                to_return[bias_name] = t133    else:134        raise NotImplementedError135    to_return = {k: maybe_zero_3(v, name=k) for k, v in to_return.items()}136    return to_return137 138 139def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True):140    to_return = {k: t for k, t in named_params if "lora_" not in k}141    if require_grad_only:142        to_return = {k: t for k, t in to_return.items() if t.requires_grad}143    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}144    return to_return145 146 147def get_mm_adapter_state_maybe_zero_3(named_params, keys_to_match):148    to_return = {k: t for k, t in named_params if any(key_match in k for key_match in keys_to_match)}149    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}150    return to_return151 152 153def find_all_linear_explicit_names(model):154    cls = torch.nn.Linear155    lora_module_names = set()156    multimodal_keywords = ['mm_projector', 'vision_tower', 'vision_resampler']157    for name, module in model.named_modules():158        if any(mm_keyword in name for mm_keyword in multimodal_keywords):159            continue160        if isinstance(module, cls):161            lora_module_names.add(name)162 163    if 'lm_head' in lora_module_names: # needed for 16-bit164        lora_module_names.remove('lm_head')165    return list(lora_module_names)166 167 168def safe_save_model_for_hf_trainer(trainer: transformers.Trainer,169                                   output_dir: str):170    """Collects the state dict and dump to disk."""171 172    if getattr(trainer.args, "tune_mm_mlp_adapter", False):173        # Only save Adapter174        keys_to_match = ['mm_projector']175        if getattr(trainer.args, "use_im_start_end", False):176            keys_to_match.extend(['embed_tokens', 'embed_in'])177 178        weight_to_save = get_mm_adapter_state_maybe_zero_3(trainer.model.named_parameters(), keys_to_match)179        trainer.model.config.save_pretrained(output_dir)180 181        current_folder = output_dir.split('/')[-1]182        parent_folder = os.path.dirname(output_dir)183        if trainer.args.local_rank == 0 or trainer.args.local_rank == -1:184            if current_folder.startswith('checkpoint-'):185                mm_projector_folder = os.path.join(parent_folder, "mm_projector")186                os.makedirs(mm_projector_folder, exist_ok=True)187                torch.save(weight_to_save, os.path.join(mm_projector_folder, f'{current_folder}.bin'))188            else:189                torch.save(weight_to_save, os.path.join(output_dir, f'mm_projector.bin'))190        return191 192    if trainer.deepspeed:193        torch.cuda.synchronize()194        trainer.save_model(output_dir)195        return196 197    state_dict = trainer.model.state_dict()198    if trainer.args.should_save:199        cpu_state_dict = {200            key: value.cpu()201            for key, value in state_dict.items()202        }203        del state_dict204        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa205 206 207def smart_tokenizer_and_embedding_resize(208    special_tokens_dict: Dict,209    tokenizer: transformers.PreTrainedTokenizer,210    model: transformers.PreTrainedModel,211):212    """Resize tokenizer and embedding.213 214    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.215    """216    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)217    model.resize_token_embeddings(len(tokenizer))218 219    if num_new_tokens > 0:220        input_embeddings = model.get_input_embeddings().weight.data221        output_embeddings = model.get_output_embeddings().weight.data222 223        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(224            dim=0, keepdim=True)225        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(226            dim=0, keepdim=True)227 228        input_embeddings[-num_new_tokens:] = input_embeddings_avg229        output_embeddings[-num_new_tokens:] = output_embeddings_avg230 231def find_all_linear_names(model):232    cls = torch.nn.Linear233    lora_module_names = set()234    multimodal_keywords = ['mm_projector', 'vision_tower', 'vision_resampler']235    for name, module in model.named_modules():236        if any(mm_keyword in name for mm_keyword in multimodal_keywords):237            continue238        if isinstance(module, cls):239            names = name.split('.')240            lora_module_names.add(names[0] if len(names) == 1 else names[-1])241 242    if 'lm_head' in lora_module_names: # needed for 16-bit243        lora_module_names.remove('lm_head')244    return list(lora_module_names)245def train():246    global local_rank247 248    parser = transformers.HfArgumentParser(249        (ModelArguments, DataArguments, TrainingArguments))250    model_args, data_args, training_args = parser.parse_args_into_dataclasses()251    local_rank = training_args.local_rank252    compute_dtype = (torch.float16 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))253 254 255    bnb_model_from_pretrained_args = {}256    if training_args.bits in [4, 8]:257        from transformers import BitsAndBytesConfig258        bnb_model_from_pretrained_args.update(dict(259            device_map={"": training_args.device},260            load_in_4bit=training_args.bits == 4,261            load_in_8bit=training_args.bits == 8,262            quantization_config=BitsAndBytesConfig(263                load_in_4bit=training_args.bits == 4,264                load_in_8bit=training_args.bits == 8,265                llm_int8_threshold=6.0,266                llm_int8_has_fp16_weight=False,267                bnb_4bit_compute_dtype=compute_dtype,268                bnb_4bit_use_double_quant=training_args.double_quant,269                bnb_4bit_quant_type=training_args.quant_type # {'fp4', 'nf4'}270            )271        ))272 273    if 'chatglm' in model_args.model_name_or_path:274        model = VTimeLLMChatGLMForCausalLM.from_pretrained(275            model_args.model_name_or_path, empty_init=False, device='cuda'276        )277    else:278        model = VTimeLLMLlamaForCausalLM.from_pretrained(279            model_args.model_name_or_path,280            cache_dir=training_args.cache_dir,281            **bnb_model_from_pretrained_args282        )283    model.config.use_cache = False284 285    if training_args.bits in [4, 8]:286        from peft import prepare_model_for_kbit_training287        model.config.torch_dtype=(torch.float32 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))288        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=training_args.gradient_checkpointing)289 290    if training_args.gradient_checkpointing:291        if hasattr(model, "enable_input_require_grads"):292            model.enable_input_require_grads()293        else:294            def make_inputs_require_grad(module, input, output):295                output.requires_grad_(True)296            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)297 298 299    if 'chatglm' in model_args.model_name_or_path:300        tokenizer = transformers.AutoTokenizer.from_pretrained(301            model_args.model_name_or_path,302            trust_remote_code=True303        )304    else:305        tokenizer = transformers.AutoTokenizer.from_pretrained(306            model_args.model_name_or_path,307            cache_dir=training_args.cache_dir,308            model_max_length=training_args.model_max_length,309            padding_side="right",310            use_fast=False,311        )312        tokenizer.pad_token = tokenizer.unk_token313 314    if training_args.lora_enable:315        from peft import LoraConfig, get_peft_model316        lora_config = LoraConfig(317            r=training_args.lora_r,318            lora_alpha=training_args.lora_alpha,319            target_modules=find_all_linear_names(model),320            lora_dropout=training_args.lora_dropout,321            bias=training_args.lora_bias,322            task_type="CAUSAL_LM",323        )324 325        # Lora config for the SFT Model326        explicit_linear = find_all_linear_explicit_names(model)327        lora_config_dpo = LoraConfig(328            r=training_args.lora_r,329            lora_alpha=training_args.lora_alpha,330            target_modules=explicit_linear,331            lora_dropout=training_args.lora_dropout,332            bias=training_args.lora_bias,333            task_type="CAUSAL_LM",334        )335 336        if training_args.bits == 16:337            if training_args.bf16:338                model.to(torch.bfloat16)339            if training_args.fp16:340                model.to(torch.float16)341 342 343        model = model.cuda()344 345        # print_trainable_parameters(model)346        if training_args.training_stage > 2:347        # if False:348            model.get_model().initialize_vision_modules(model_args)349 350            model = load_lora(model, model_args.stage2_path)351            rank0_print('Merging stage 2 LoRA weights...')352            model = model.merge_and_unload()353 354            355        if training_args.finetuning:356            rank0_print("*" * 90)357            model = load_lora(model, model_args.stage3_path)358            rank0_print('Merging stage 3 LoRA weights...')359            model = model.merge_and_unload()360            rank0_print("*" * 90)361 362            if not training_args.train4dpo:363                rank0_print("Adding LoRA adapters...")364                model = get_peft_model(model, lora_config)365 366        if training_args.train4dpo:367            rank0_print("*" * 90)368            rank0_print("Preparing for stage 5 (dpo training)")369 370            model = load_lora(model, model_args.stage4_path)371            rank0_print('Merging stage 4 LoRA weights...')372            model = model.merge_and_unload()373 374            rank0_print("*" * 90)375            # build new lora config for dpo376            rank0_print("Adding LoRA adapters...")377            model = get_peft_model(model, lora_config_dpo)378 379        else:380 381            rank0_print("Adding LoRA adapters...")382            model = get_peft_model(model, lora_config)383 384 385        print_trainable_parameters(model)386 387    388    if model_args.version in conversation_lib.conv_templates:389        conversation_lib.default_conversation = conversation_lib.conv_templates[model_args.version]390    else:391        conversation_lib.default_conversation = conversation_lib.conv_templates["vicuna_v1"]392 393 394    model.config.tune_mm_mlp_adapter = training_args.tune_mm_mlp_adapter = model_args.tune_mm_mlp_adapter395    model.config.freeze_mm_mlp_adapter = training_args.freeze_mm_mlp_adapter396 397    if training_args.training_stage != 3:398        model.get_model().initialize_vision_modules(model_args=model_args)399 400 401        if model_args.tune_mm_mlp_adapter:402            model.requires_grad_(False)403            for p in model.get_model().mm_projector.parameters():404                p.requires_grad = True405 406        407        if training_args.freeze_mm_mlp_adapter:408            for p in model.get_model().mm_projector.parameters():409                p.requires_grad = False410 411 412 413 414    if training_args.bits in [4, 8]:415        model.get_model().mm_projector.to(dtype=compute_dtype, device=training_args.device)416 417    if training_args.bits in [4, 8]:418        from peft.tuners.lora import LoraLayer419        for name, module in model.named_modules():420            if isinstance(module, LoraLayer):421                if training_args.bf16:422                    module = module.to(torch.bfloat16)423            if 'norm' in name:424                module = module.to(torch.float32)425            if 'lm_head' in name or 'embed_tokens' in name:426                if hasattr(module, 'weight'):427                    if training_args.bf16 and module.weight.dtype == torch.float32:428                        module = module.to(torch.bfloat16)429 430    data_module = make_supervised_dpo_data_module(tokenizer=tokenizer,431                                              data_args=data_args)432    trainer = DPOTrainer(model=model,433                         ref_model=None,434                        tokenizer=tokenizer,435                        args=training_args,436                         beta=training_args.beta,437                         generate_during_eval=False,438                         dpo_alpha=training_args.dpo_alpha,439                         peft_config=None,440                         gamma=training_args.gamma,441                         **data_module)442 443    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):444        trainer.train(resume_from_checkpoint=True)445    else:446        trainer.train()447    trainer.save_state()448 449    model.config.use_cache = True450 451    if training_args.lora_enable:452        state_dict = get_peft_state_maybe_zero_3(453            model.named_parameters(), training_args.lora_bias454        )455        non_lora_state_dict = get_peft_state_non_lora_maybe_zero_3(456            model.named_parameters()457        )458        if training_args.local_rank == 0 or training_args.local_rank == -1:459            model.config.save_pretrained(training_args.output_dir)460            model.save_pretrained(training_args.output_dir, state_dict=state_dict)461            torch.save(non_lora_state_dict, os.path.join(training_args.output_dir, 'non_lora_trainables.bin'))462    else:463        safe_save_model_for_hf_trainer(trainer=trainer,464                                       output_dir=training_args.output_dir)465 466 467if __name__ == "__main__":468    train()469