import logging import os import pathlib import random from dataclasses import dataclass, field from typing import Optional import numpy as np import torch import transformers from transformers import set_seed from vita import conversation as conversation_lib from vita.model import * from vita.train.vita_trainer import VITATrainer from vita.util.data_utils_video_audio_neg_patch import make_supervised_data_module, DataArguments #from vita.util.data_utils_video_audio_neg_patch_fo import make_supervised_data_module, DataArguments #from vita.util.data_utils_video_audio_patch import make_supervised_data_module, DataArguments #from vita.util.data_utils_video_audio_patch_sf import make_supervised_data_module, DataArguments #from vita.util.data_utils_video_patch_audio import make_supervised_data_module, DataArguments def set_random_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) set_seed(seed) set_random_seed(42) local_rank = None def rank0_print(*args): if local_rank == 0: print(*args) @dataclass class ModelArguments: model_name_or_path: Optional[str] = field(default=None) model_type: Optional[str] = field(default=None) version: Optional[str] = field(default=None) freeze_backbone: bool = field(default=False) tune_mm_mlp_adapter: bool = field(default=False) tune_audio_mlp_adapter: bool = field(default=False) audio_prompt_finetune: bool = field(default=False) audio_prompt_num: Optional[int] = field(default=None) audio_state_predictor_tuning: bool = field(default=False) vision_tower: Optional[str] = field(default=None) audio_encoder: Optional[str] = field(default=None) freeze_audio_encoder: bool = field(default=True) freeze_audio_encoder_adapter: bool = field(default=True) unfreeze_vision_tower: bool = field(default=False) use_s2: bool = field(default=False) pretrain_audio_mlp_adapter: Optional[str] = field(default=None) pretrain_mm_mlp_adapter: Optional[str] = field(default=None) mm_projector_type: Optional[str] = field(default="mlp2x_gelu") @dataclass class TrainingArguments(transformers.TrainingArguments): cache_dir: Optional[str] = field(default=None) optim: str = field(default="adamw_torch") remove_unused_columns: bool = field(default=False) freeze_mm_mlp_adapter: bool = field(default=False) mpt_attn_impl: Optional[str] = field(default="triton") model_max_length: int = field( default=512, metadata={ "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)." }, ) double_quant: bool = field( default=True, metadata={"help": "Compress the quantization statistics through double quantization."}, ) quant_type: str = field( default="nf4", metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}, ) bits: int = field(default=16, metadata={"help": "How many bits to use."}) lora_enable: bool = False lora_r: int = 64 lora_alpha: int = 16 lora_dropout: float = 0.05 lora_weight_path: str = "" lora_bias: str = "none" mm_projector_lr: Optional[float] = None group_by_modality_length: bool = field(default=False) def maybe_zero_3(param, ignore_status=False, name=None): from deepspeed import zero from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus if hasattr(param, "ds_id"): if param.ds_status == ZeroParamStatus.NOT_AVAILABLE: if not ignore_status: logging.warning( f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}" ) with zero.GatheredParameters([param]): param = param.data.detach().cpu().clone() else: param = param.detach().cpu().clone() return param # Borrowed from peft.util.get_peft_model_state_dict def get_peft_state_maybe_zero_3(named_params, bias): if bias == "none": to_return = {k: t for k, t in named_params if "lora_" in k} elif bias == "all": to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k} elif bias == "lora_only": to_return = {} maybe_lora_bias = {} lora_bias_names = set() for k, t in named_params: if "lora_" in k: to_return[k] = t bias_name = k.split("lora_")[0] + "bias" lora_bias_names.add(bias_name) elif "bias" in k: maybe_lora_bias[k] = t for k, t in maybe_lora_bias: if bias_name in lora_bias_names: to_return[bias_name] = t else: raise NotImplementedError to_return = {k: maybe_zero_3(v, ignore_status=True) for k, v in to_return.items()} return to_return def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True): to_return = {k: t for k, t in named_params if "lora_" not in k} if require_grad_only: to_return = {k: t for k, t in to_return.items() if t.requires_grad} to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()} return to_return def get_mm_adapter_state_maybe_zero_3(named_params, keys_to_match): to_return = { k: t for k, t in named_params if any(key_match in k for key_match in keys_to_match) } to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()} return to_return def find_all_linear_names(model): cls = torch.nn.Linear lora_module_names = set() multimodal_keywords = ["mm_projector", "vision_tower", "vision_resampler"] for name, module in model.named_modules(): if any(mm_keyword in name for mm_keyword in multimodal_keywords): continue if isinstance(module, cls): names = name.split(".") lora_module_names.add(names[0] if len(names) == 1 else names[-1]) if "lm_head" in lora_module_names: # needed for 16-bit lora_module_names.remove("lm_head") return list(lora_module_names) def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str): """Collects the state dict and dump to disk.""" if getattr(trainer.args, "tune_mm_mlp_adapter", False): # Only save Adapter keys_to_match = ["mm_projector"] if getattr(trainer.args, "use_im_start_end", False): keys_to_match.extend(["embed_tokens", "embed_in"]) weight_to_save = get_mm_adapter_state_maybe_zero_3( trainer.model.named_parameters(), keys_to_match ) trainer.model.config.save_pretrained(output_dir) current_folder = output_dir.split("/")[-1] parent_folder = os.path.dirname(output_dir) if trainer.args.local_rank == 0 or trainer.args.local_rank == -1: if current_folder.startswith("checkpoint-"): mm_projector_folder = os.path.join(parent_folder, "mm_projector") os.makedirs(mm_projector_folder, exist_ok=True) torch.save( weight_to_save, os.path.join(mm_projector_folder, f"{current_folder}.bin") ) else: torch.save(weight_to_save, os.path.join(output_dir, f"mm_projector.bin")) return if trainer.deepspeed: torch.cuda.synchronize() trainer.save_model(output_dir) return state_dict = trainer.model.state_dict() if trainer.args.should_save: cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()} del state_dict trainer._save(output_dir, state_dict=cpu_state_dict) # noqa def train(): global local_rank parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) model_args, data_args, training_args = parser.parse_args_into_dataclasses() local_rank = training_args.local_rank compute_dtype = ( torch.float16 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32) ) bnb_model_from_pretrained_args = {} if training_args.bits in [4, 8]: from transformers import BitsAndBytesConfig bnb_model_from_pretrained_args.update( dict( device_map={"": training_args.device}, load_in_4bit=training_args.bits == 4, load_in_8bit=training_args.bits == 8, quantization_config=BitsAndBytesConfig( load_in_4bit=training_args.bits == 4, load_in_8bit=training_args.bits == 8, llm_int8_skip_modules=["mm_projector"], llm_int8_threshold=6.0, llm_int8_has_fp16_weight=False, bnb_4bit_compute_dtype=compute_dtype, bnb_4bit_use_double_quant=training_args.double_quant, bnb_4bit_quant_type=training_args.quant_type, # {'fp4', 'nf4'} ), ) ) assert model_args.vision_tower is not None if model_args.model_type in {"mixtral-8x7b", "nemo", "qwen2p5_instruct", "qwen2p5_fo_instruct"}: tokenizer = transformers.AutoTokenizer.from_pretrained( model_args.model_name_or_path, cache_dir=training_args.cache_dir, model_max_length=training_args.model_max_length, padding_side="right", use_fast=True, ) if tokenizer.unk_token is not None and tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.unk_token if model_args.model_type == "llama3-8b": tokenizer.pad_token = tokenizer.eos_token if model_args.model_type == "mixtral-8x7b": torch_dtype = torch.float16 if training_args.fp16 else torch.bfloat16 model = VITAMixtralForCausalLM.from_pretrained( model_args.model_name_or_path, cache_dir=training_args.cache_dir, torch_dtype=torch_dtype, attn_implementation="flash_attention_2", **bnb_model_from_pretrained_args, ) elif model_args.model_type == "nemo": torch_dtype = torch.float16 if training_args.fp16 else torch.bfloat16 model = VITAMistralForCausalLM.from_pretrained( model_args.model_name_or_path, cache_dir=training_args.cache_dir, torch_dtype=torch_dtype, attn_implementation="flash_attention_2", **bnb_model_from_pretrained_args, ) elif model_args.model_type == "qwen2p5_instruct": torch_dtype = torch.float16 if training_args.fp16 else torch.bfloat16 model = VITAQwen2ForCausalLM.from_pretrained( model_args.model_name_or_path, cache_dir=training_args.cache_dir, torch_dtype=torch_dtype, attn_implementation="flash_attention_2", **bnb_model_from_pretrained_args, ) elif model_args.model_type == "qwen2p5_fo_instruct": torch_dtype = torch.float16 if training_args.fp16 else torch.bfloat16 model = VITAFOQwen2ForCausalLM.from_pretrained( model_args.model_name_or_path, cache_dir=training_args.cache_dir, torch_dtype=torch_dtype, attn_implementation="flash_attention_2", **bnb_model_from_pretrained_args, ) else: raise ValueError(f"Unknown Model Type {model_args.model_type}") model.config.use_cache = False if model_args.freeze_backbone: model.model.requires_grad_(False) if training_args.bits in [4, 8]: from peft import prepare_model_for_kbit_training model.config.torch_dtype = ( torch.float32 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32) ) model = prepare_model_for_kbit_training( model, use_gradient_checkpointing=training_args.gradient_checkpointing ) if training_args.gradient_checkpointing: if hasattr(model, "enable_input_require_grads"): model.enable_input_require_grads() else: def make_inputs_require_grad(module, input, output): output.requires_grad_(True) model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) if training_args.lora_enable: from peft import LoraConfig, get_peft_model lora_config = LoraConfig( r=training_args.lora_r, lora_alpha=training_args.lora_alpha, target_modules=find_all_linear_names(model), lora_dropout=training_args.lora_dropout, bias=training_args.lora_bias, task_type="CAUSAL_LM", ) if training_args.bits == 16: if training_args.bf16: model.to(torch.bfloat16) if training_args.fp16: model.to(torch.float16) rank0_print("Adding LoRA adapters...") model = get_peft_model(model, lora_config) if model_args.version in conversation_lib.conv_templates: conversation_lib.default_conversation = conversation_lib.conv_templates[model_args.version] else: conversation_lib.default_conversation = conversation_lib.conv_templates["default"] model.get_model().initialize_vision_modules(model_args=model_args) model.config.freeze_audio_encoder = model_args.freeze_audio_encoder model.config.freeze_audio_encoder_adapter = model_args.freeze_audio_encoder_adapter model.config.audio_prompt_finetune = model_args.audio_prompt_finetune model.config.audio_prompt_num = model_args.audio_prompt_num model.get_model().initialize_audio_modules(model_args=model_args) vision_tower = model.get_vision_tower() vision_tower.to( dtype=torch.bfloat16 if training_args.bf16 else torch.float16, device=training_args.device ) audio_encoder = model.get_audio_encoder() audio_encoder.to( dtype=torch.bfloat16 if training_args.bf16 else torch.float16, device=training_args.device ) data_args.image_processor = vision_tower.image_processor data_args.audio_processor = audio_encoder.audio_processor model.config.image_aspect_ratio = data_args.image_aspect_ratio model.config.tokenizer_padding_side = tokenizer.padding_side model.config.tokenizer_model_max_length = tokenizer.model_max_length model.config.tune_mm_mlp_adapter = ( training_args.tune_mm_mlp_adapter ) = model_args.tune_mm_mlp_adapter if model_args.tune_mm_mlp_adapter: model.requires_grad_(False) for p in model.get_model().mm_projector.parameters(): p.requires_grad = True model.config.tune_audio_mlp_adapter = ( training_args.tune_audio_mlp_adapter ) = model_args.tune_audio_mlp_adapter if model_args.tune_audio_mlp_adapter: model.requires_grad_(False) for p in model.model.audio_encoder.adpter.parameters(): p.requires_grad = True model.config.audio_prompt_finetune = ( training_args.audio_prompt_finetune ) = model_args.audio_prompt_finetune model.config.audio_state_predictor_tuning = ( training_args.audio_state_predictor_tuning ) = model_args.audio_state_predictor_tuning if model_args.audio_prompt_finetune or model_args.audio_state_predictor_tuning: model.requires_grad_(False) if model_args.audio_prompt_finetune: for p in model.model.audio_encoder.prompt_embeddings.parameters(): p.requires_grad = True if model_args.audio_state_predictor_tuning: for p in model.predictor_head.parameters(): p.requires_grad = True model.config.freeze_mm_mlp_adapter = training_args.freeze_mm_mlp_adapter if training_args.freeze_mm_mlp_adapter: for p in model.get_model().mm_projector.parameters(): p.requires_grad = False if training_args.bits in [4, 8]: model.get_model().mm_projector.to(dtype=compute_dtype, device=training_args.device) model.config.mm_projector_lr = training_args.mm_projector_lr model.config.use_s2 = model_args.use_s2 model.config.unfreeze_vision_tower = ( training_args.unfreeze_vision_tower ) = model_args.unfreeze_vision_tower if training_args.unfreeze_vision_tower: for p in model.get_model().vision_tower.parameters(): p.requires_grad = True if training_args.bits in [4, 8]: from peft.tuners.lora import LoraLayer for name, module in model.named_modules(): if isinstance(module, LoraLayer): if training_args.bf16: module = module.to(torch.bfloat16) if "norm" in name: module = module.to(torch.float32) if "lm_head" in name or "embed_tokens" in name: if hasattr(module, "weight"): if training_args.bf16 and module.weight.dtype == torch.float32: module = module.to(torch.bfloat16) data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args) trainer = VITATrainer(model=model, tokenizer=tokenizer, args=training_args, **data_module) if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")): trainer.train(resume_from_checkpoint=True) else: trainer.train() trainer.save_state() model.config.use_cache = True if training_args.lora_enable: state_dict = get_peft_state_maybe_zero_3(model.named_parameters(), training_args.lora_bias) non_lora_state_dict = get_peft_state_non_lora_maybe_zero_3(model.named_parameters()) if training_args.local_rank == 0 or training_args.local_rank == -1: model.config.save_pretrained(training_args.output_dir) model.save_pretrained(training_args.output_dir, state_dict=state_dict) torch.save( non_lora_state_dict, os.path.join(training_args.output_dir, "non_lora_trainables.bin"), ) else: safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir) if __name__ == "__main__": train()