code: 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 = dict(
          torch_dtype=(
              torch.float32
              if training_args.fp16
              else (
                  torch.bfloat16
                  if training_args.bf16
                  else torch.float32
              )
          ),
      )
      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_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'}
                  ),
              )
          )
  
      config = transformers.AutoConfig.from_pretrained(
          model_args.model_name_or_path, trust_remote_code=True
      )
      orig_rope_scaling = getattr(config, "rope_scaling", None)
      if orig_rope_scaling is None:
          orig_rope_scaling = {"factor": 1}
  
      orig_rope_scaling_factor = (
          orig_rope_scaling["factor"]
          if "factor" in orig_rope_scaling.keys()
          else 1
      )
      orig_ctx_len = getattr(config, "max_position_embeddings", None)
      if orig_ctx_len:
          orig_ctx_len *= orig_rope_scaling_factor
          if training_args.model_max_length > orig_ctx_len:
              scaling_factor = float(
                  math.ceil(
                      training_args.model_max_length / orig_ctx_len
                  )
              )
              config.rope_scaling = {
                  "type": "linear",
                  "factor": scaling_factor,
              }
  
      if model_args.vision_tower is not None:
          model = LlavaLlamaAttForCausalLM.from_pretrained(
              model_args.model_name_or_path,
              config=config,
              cache_dir=training_args.cache_dir,
              **bnb_model_from_pretrained_args,
          )
      else:
          model = transformers.LlamaForCausalLM.from_pretrained(
              model_args.model_name_or_path,
              config=config,
              cache_dir=training_args.cache_dir,
              **bnb_model_from_pretrained_args,
          )
      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 "mpt" in model_args.model_name_or_path:
          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",
          )
      else:
          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=False,
          )
  
      if model_args.version == "v0":
          if tokenizer.pad_token is None:
              smart_tokenizer_and_embedding_resize(
                  special_tokens_dict=dict(pad_token="[PAD]"),
                  tokenizer=tokenizer,
                  model=model,
              )
      elif model_args.version == "v0.5":
          tokenizer.pad_token = tokenizer.unk_token
      else:
          tokenizer.pad_token = tokenizer.unk_token
          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["vicuna_v1"]
              )
  
      if model_args.vision_tower is not None:
          model.get_model().initialize_vision_modules(
              model_args=model_args,
              fsdp=training_args.fsdp,
              max_token=training_args.model_max_length,
          )
  
          vision_tower = model.get_vision_tower()
          vision_tower.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.is_multimodal = True
  
          model.config.image_aspect_ratio = data_args.image_aspect_ratio
          model.config.image_grid_pinpoints = (
              data_args.image_grid_pinpoints
          )
  
          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.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_use_im_start_end = (
              data_args.mm_use_im_start_end
          ) = model_args.mm_use_im_start_end
          training_args.use_im_start_end = (
              model_args.mm_use_im_start_end
          )
          model.config.mm_use_im_patch_token = (
              model_args.mm_use_im_patch_token
          )
          model.initialize_vision_tokenizer(
              model_args, tokenizer=tokenizer
          )
  
      # all the attention modules require grad
      model.get_model().initialize_attention_modules(model_args)
  
      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 = LLaVATrainer(
          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
          )
  
imports: import os
  import copy
  import random
  from dataclasses import dataclass, field
  import json
  import logging
  import pathlib
  from typing import Dict, Optional, Sequence, List
  import pickle
  import math
  import torch
  import transformers
  from llamavid.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
  from torch.utils.data import Dataset
  from llamavid.train.llava_trainer import LLaVATrainer
  from llamavid import conversation
  from llamavid.model import *
  from llava.mm_utils import tokenizer_image_token
  from PIL import Image
  from decord import VideoReader, cpu
label_desc: Use image node_id 22 for calling a global function with example usage: train() without return types