code: def run_training(args, train_data, val_data):
      print("Loading the model")
  
      lora_config = LoraConfig(
          r=16,
          lora_alpha=32,
          lora_dropout=0.05,
          bias="none",
          task_type="CAUSAL_LM",
      )
  
      train_data.start_iteration = 0
  
      print("Starting main loop")
  
      training_args = TrainingArguments(
          output_dir=args.output_dir,
          dataloader_drop_last=True,
          evaluation_strategy="steps",
          max_steps=args.max_steps,
          eval_steps=args.eval_freq,
          save_steps=args.save_freq,
          logging_steps=args.log_freq,
          per_device_train_batch_size=args.batch_size,
          per_device_eval_batch_size=args.batch_size,
          learning_rate=args.learning_rate,
          lr_scheduler_type=args.lr_scheduler_type,
          warmup_steps=args.num_warmup_steps,
          gradient_accumulation_steps=args.gradient_accumulation_steps,
          gradient_checkpointing=args.gradient_checkpointing,
          fp16=args.fp16,
          bf16=args.bf16,
          weight_decay=args.weight_decay,
          run_name="llama-7b-finetuned",
          report_to="wandb",
          ddp_find_unused_parameters=False,
      )
  
      model = AutoModelForCausalLM.from_pretrained(
          args.model_path,
          load_in_8bit=True,
          device_map={"": Accelerator().process_index},
      )
  
      trainer = SFTTrainer(
          model=model,
          args=training_args,
          train_dataset=train_data,
          eval_dataset=val_data,
          peft_config=lora_config,
          packing=True,
      )
  
      print_trainable_parameters(trainer.model)
  
      print("Training...")
      trainer.train()
  
      print("Saving last checkpoint of the model")
      trainer.model.save_pretrained(
          os.path.join(args.output_dir, "final_checkpoint/")
      )
  
file_path: trl/examples/research_projects/stack_llama/scripts/supervised_finetuning.py
imports: import argparse
  import os
  from accelerate import Accelerator
  from datasets import load_dataset
  from peft import LoraConfig
  from tqdm import tqdm
  from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, logging, set_seed
  from trl import SFTTrainer
  from trl.trainer import ConstantLengthDataset
label_desc: Use image node_id 6 for calling a global function with example usage: run_training(args, train_data, val_data) without return types