code: def train():
      parser = transformers.HfArgumentParser(
          (ModelArguments, DataArguments, TrainingArguments)
      )
      (
          model_args,
          data_args,
          training_args,
      ) = parser.parse_args_into_dataclasses()
  
      # NOTE: May expand supported model types in the future
      if model_args.model_type == "gpt-neox":
          replace_gpt_neox_attn(
              training_args.use_flash_attn, training_args.use_full_attn
          )
      else:
          replace_llama_attn(
              training_args.use_flash_attn, training_args.use_full_attn
          )
  
      # Set RoPE scaling factor
      config = transformers.AutoConfig.from_pretrained(
          model_args.model_name_or_path,
          cache_dir=training_args.cache_dir,
      )
  
      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,
              }
  
      # Load model and tokenizer
      model = transformers.AutoModelForCausalLM.from_pretrained(
          model_args.model_name_or_path,
          config=config,
          cache_dir=training_args.cache_dir,
          torch_dtype=torch.bfloat16,
          quantization_config=BitsAndBytesConfig(
              load_in_4bit=True,
              llm_int8_threshold=6.0,
              llm_int8_has_fp16_weight=False,
              bnb_4bit_compute_dtype=torch.bfloat16,
              bnb_4bit_use_double_quant=True,
              bnb_4bit_quant_type="nf4",
          ),
      )
  
      for param in model.parameters():
          param.requires_grad = (
              False  # freeze the model - train adapters later
          )
          if param.ndim == 1:
              # cast the small parameters (e.g. layernorm) to fp32 for stability
              param.data = param.data.to(torch.float32)
  
      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,
      )
  
      special_tokens_dict = dict()
      if tokenizer.pad_token is None:
          special_tokens_dict["pad_token"] = DEFAULT_PAD_TOKEN
      if tokenizer.eos_token is None:
          special_tokens_dict["eos_token"] = DEFAULT_EOS_TOKEN
      if tokenizer.bos_token is None:
          special_tokens_dict["bos_token"] = DEFAULT_BOS_TOKEN
      if tokenizer.unk_token is None:
          special_tokens_dict["unk_token"] = DEFAULT_UNK_TOKEN
  
      smart_tokenizer_and_embedding_resize(
          special_tokens_dict=special_tokens_dict,
          tokenizer=tokenizer,
          model=model,
      )
  
      data_module = make_supervised_data_module(
          tokenizer=tokenizer, data_args=data_args
      )
  
      if training_args.low_rank_training:
          if model_args.model_type == "gpt-neox":
              # added `dense` to match with llama as the basic LoRA would only target 'query_key_value'
              targets = ["query_key_value", "dense"]
          else:
              targets = ["q_proj", "k_proj", "v_proj", "o_proj"]
  
          config = LoraConfig(
              r=8,
              lora_alpha=16,
              target_modules=targets,
              lora_dropout=0,
              bias="none",
              task_type="CAUSAL_LM",
          )
          model = get_peft_model(model, config)
          # enable trainable params
          [
              p.requires_grad_()
              for n, p in model.named_parameters()
              if any(
                  [
                      k in n
                      for k in training_args.trainable_params.split(",")
                  ]
              )
          ]
  
      class CastOutputToFloat(nn.Sequential):
          def forward(self, x):
              return super().forward(x).to(torch.float32)
  
      model.lm_head = CastOutputToFloat(model.lm_head)
  
      # Verifying the datatypes.
      dtypes = {}
      for _, p in model.named_parameters():
          dtype = p.dtype
          if dtype not in dtypes:
              dtypes[dtype] = 0
          dtypes[dtype] += p.numel()
      total = 0
      for k, v in dtypes.items():
          total += v
      for k, v in dtypes.items():
          print(k, v, v / total)
  
      model.config.use_cache = (
          False  # required for gradient checkpointing
      )
      model.enable_input_require_grads()  # required for gradient checkpointing
      model.gradient_checkpointing_enable()  # enable gradient checkpointing
  
      trainer = Trainer(
          model=model,
          tokenizer=tokenizer,
          args=training_args,
          **data_module,
      )
      trainer.train()
      trainer.save_state()
      trainer.save_model(output_dir=training_args.output_dir)
  
imports: import io
  import os
  import copy
  import json
  import math
  import logging
  from dataclasses import dataclass, field
  from typing import Dict, Optional, Sequence
  import torch
  import torch.nn
  import transformers
  from torch.utils.data import Dataset
  from transformers import Trainer, DataCollatorForLanguageModeling, BitsAndBytesConfig
  from llama_attn_replace_sft import replace_llama_attn
  from gptneox_attn_replace import replace_gpt_neox_attn
  from peft import LoraConfig, get_peft_model
  from torch.distributed import barrier
label_desc: Use image node_id 7 for calling a global function with example usage: train() and returns: super