code: def run_train(
      cfg: Any,
      model: torch.nn.Module,
      optimizer,
      scheduler,
      epoch_steps,
      train_dataloader,
      val_dataloader,
      val_df: pd.DataFrame,
  ):
      """Runs the training loop.
  
      Args:
          cfg: config object
          model: model
          train_dataloader: custom training Dataloader
          train_df: train DataFrame
          val_dataloader: custom validation Dataloader
          val_df: validation DataFrame
  
      Returns:
          Validation prediction output
          Validation loss
          Validation metric
          Last train batch
      """
      if (
          hasattr(cfg.augmentation, "neftune_noise_alpha")
          and cfg.augmentation.neftune_noise_alpha > 0
      ):
          activate_neftune(model, cfg.augmentation.neftune_noise_alpha)
  
      scaler: GradScaler | None = None
      if cfg.environment.mixed_precision:
          scaler = GradScaler()
  
      optimizer.zero_grad(set_to_none=True)
  
      # Prepare NLP Augmentation
      nlp_augment = None
      if hasattr(cfg.augmentation, "nlp_augmentations_class"):
          nlp_augment = cfg.augmentation.nlp_augmentations_class(
              cfg=cfg
          )
  
      start_epoch = 0
  
      _, metric_mode, _ = cfg.prediction.metric_class.get(
          cfg.prediction.metric
      )
      objective_op: Callable[[float, float], bool]
      if metric_mode == "max":
          best_val_metric = -np.inf
          objective_op = np.greater
      else:
          best_val_metric = np.inf
          objective_op = np.less
  
      if cfg.training.evaluate_before_training:
          val_loss, val_metric = run_eval(
              cfg=cfg,
              model=model,
              val_dataloader=val_dataloader,
              val_df=val_df,
          )
  
      for epoch in range(start_epoch, cfg.training.epochs):
          set_seed(
              cfg.environment._seed
              + epoch
              * cfg.environment._world_size
              * cfg.environment.number_of_workers
              + cfg.environment._local_rank
              * cfg.environment.number_of_workers
          )
          if cfg.environment._local_rank == 0:
              logger.info(
                  f"Training Epoch: {epoch + 1} / {cfg.training.epochs}"
              )
  
          if (
              cfg.environment._distributed
              and not cfg.environment.use_deepspeed
              and hasattr(train_dataloader.sampler, "set_epoch")
          ):
              train_dataloader.sampler.set_epoch(epoch)  # type: ignore
  
          tqdm_out = TqdmToLogger(logger, level=logging.INFO)
          progress_bar = tqdm(
              total=epoch_steps,
              disable=cfg.environment._local_rank != 0,
              file=tqdm_out,
              ascii=True,
              desc="train loss",
              mininterval=0,
          )
          tr_it = iter(train_dataloader)
  
          losses = []
          model.train()
  
          log_update_steps = max(epoch_steps // 20, 1)
          evaluation_step = max(
              int(epoch_steps * cfg.training.evaluation_epochs), 1
          )
          logger.info(f"Evaluation step: {evaluation_step}")
  
          for itr, data in enumerate(tr_it):
              cfg.environment._curr_step += (
                  cfg.training.batch_size * cfg.environment._world_size
              )
  
              # Batch to device
              batch = cfg.dataset.dataset_class.batch_to_device(
                  data, cfg.environment._device
              )
  
              # NLP augmentation
              if nlp_augment is not None:
                  batch = nlp_augment(batch)
  
              # Plot first batch
              if (
                  epoch == 0
                  and itr == 0
                  and cfg.environment._local_rank == 0
              ):
                  plot = cfg.logging.plots_class.plot_batch(
                      batch=batch, cfg=cfg
                  )
                  log_plot(cfg, plot, "train_data")
  
              # Forward pass
              with autocast(enabled=cfg.environment.mixed_precision):
                  output_dict = model.forward(batch)
  
              loss = output_dict["loss"]
              if ~np.isfinite(loss.item()) and (
                  epoch > start_epoch or itr > 20
              ):
                  raise LLMTrainingException(
                      "NaN caught in loss during training. "
                      "Please, reduce learning rate, change dtype, "
                      "or disable mixed precision. Alternatively, "
                      "gradient clipping may help to stabilize training."
                  )
              losses.append(loss.item())
  
              # loss is a mean loss per batch/sample
              # as grad_accumulations sums up the gradients, this loss must be scaled
              # by the number of grad_accumulations, to have similar behavior for
              # BS * grad_accumulations = const.
              if cfg.training.grad_accumulation != 1:
                  loss = loss / cfg.training.grad_accumulation
  
              # Backward pass
              if (
                  cfg.environment.mixed_precision
                  and not cfg.environment.use_deepspeed
              ):
                  scaler.scale(loss).backward()  # type: ignore
                  if itr % cfg.training.grad_accumulation == 0:
                      if cfg.training.gradient_clip > 0:
                          scaler.unscale_(optimizer)  # type: ignore
                          torch.nn.utils.clip_grad_norm_(
                              model.parameters(),
                              cfg.training.gradient_clip,
                          )
                      scaler.step(optimizer)  # type: ignore
                      scaler.update()
                      optimizer.zero_grad(set_to_none=True)
              else:
                  if cfg.environment.use_deepspeed:
                      model.backward(loss)  # type: ignore[operator]
                  else:
                      loss.backward()
                  if itr % cfg.training.grad_accumulation == 0:
                      if cfg.training.gradient_clip > 0:
                          torch.nn.utils.clip_grad_norm_(
                              model.parameters(),
                              cfg.training.gradient_clip,
                          )
                      optimizer.step()
                      optimizer.zero_grad(set_to_none=True)
  
              if cfg.environment._distributed:
                  torch.cuda.synchronize(
                      device=cfg.environment._local_rank
                  )
  
              if scheduler is not None:
                  scheduler.step()
  
              if cfg.environment._local_rank == 0:
                  cfg.logging._logger.log(
                      "train",
                      "loss",
                      losses[-1],
                      step=cfg.environment._curr_step,
                  )
                  cfg.logging._logger.log(
                      "meta",
                      "lr",
                      optimizer.param_groups[0]["lr"],
                      step=cfg.environment._curr_step,
                  )
                  if cfg.training.differential_learning_rate_layers:
                      cfg.logging._logger.log(
                          "meta",
                          "lr_diff",
                          optimizer.param_groups[2]["lr"],
                          step=cfg.environment._curr_step,
                      )
  
                  cfg.logging._logger.log(
                      "internal",
                      "current_step",
                      cfg.environment._curr_step,
                      step=cfg.environment._curr_step,
                  )
  
                  for key in [
                      "chosen_rewards",
                      "rejected_rewards",
                      "reward_margin",
                  ]:
                      if key in output_dict:
                          cfg.logging._logger.log(
                              "train",
                              key,
                              output_dict[key].item(),
                              step=cfg.environment._curr_step,
                          )
  
                  # Show logs each 5% of the epoch (only if doing per epoch evaluation)
                  if (
                      (itr + 1) % log_update_steps == 0
                      or itr == epoch_steps - 1
                  ):
                      progress_bar.set_description(
                          f"train loss: {np.mean(losses[-10:]):.2f}",
                          refresh=False,
                      )
                      if (itr + 1) % log_update_steps == 0:
                          progress_bar.update(log_update_steps)
                      else:
                          progress_bar.update(
                              epoch_steps % log_update_steps
                          )
  
                  del output_dict
  
              # Validation loop
              if (itr + 1) % evaluation_step == 0:
                  if cfg.training.evaluation_epochs == 1:
                      progress_bar.close()
  
                  # TODO: Move back after fixing slow generation of deepspeed.
                  if not cfg.training.save_best_checkpoint:
                      checkpoint_path = cfg.output_directory
                      if cfg.environment._local_rank == 0:
                          logger.info(
                              f"Saving last model checkpoint to {checkpoint_path}"
                          )
                      save_checkpoint(
                          model=model, path=checkpoint_path, cfg=cfg
                      )
  
                  val_loss, val_metric = run_eval(
                      cfg=cfg,
                      model=model,
                      val_dataloader=val_dataloader,
                      val_df=val_df,
                  )
  
                  if cfg.training.save_best_checkpoint:
                      if objective_op(val_metric, best_val_metric):
                          checkpoint_path = cfg.output_directory
                          if cfg.environment._local_rank == 0:
                              logger.info(
                                  f"Saving best model checkpoint: "
                                  f"val_{cfg.prediction.metric} {best_val_metric:.5} -> "
                                  f"{val_metric:.5} to {checkpoint_path}"
                              )
                          save_checkpoint(
                              model=model, path=checkpoint_path, cfg=cfg
                          )
                          best_val_metric = val_metric
  
                  model.train()
  
          progress_bar.close()
          del progress_bar
  
          if cfg.environment._distributed:
              torch.cuda.synchronize(device=cfg.environment._local_rank)
              torch.distributed.barrier()
  
          if cfg.environment._local_rank == 0:
              cfg.logging._logger.log(
                  "internal",
                  "epoch",
                  epoch + 1,
                  step=cfg.environment._curr_step,
              )
  
      if cfg.environment._distributed:
          torch.distributed.barrier()
  
      return val_loss, val_metric
  
imports: import os
  import argparse
  import gc
  import logging
  import sys
  import time
  from distutils import util
  from typing import Any, Callable, Dict, Tuple
  import deepspeed
  import numpy
  import pandas
  import torch
  from torch.cuda.amp import GradScaler, autocast
  from torch.utils.data import DataLoader
  from tqdm import tqdm
  from transformers.deepspeed import HfDeepSpeedConfig
  from llm_studio.src.datasets.text_utils import get_tokenizer
  from llm_studio.src.loggers import MainLogger
  from llm_studio.src.trl.trainer import PPOTrainer
  from llm_studio.src.utils.config_utils import load_config_py, load_config_yaml, save_config_yaml
  from llm_studio.src.utils.data_utils import get_data, get_inference_batch_size, get_train_dataloader, get_train_dataset, get_val_dataloader, get_val_dataset
  from llm_studio.src.utils.exceptions import LLMTrainingException
  from llm_studio.src.utils.export_utils import save_prediction_outputs
  from llm_studio.src.utils.gpu_utils import sync_across_processes
  from llm_studio.src.utils.logging_utils import TqdmToLogger, initialize_logging, log_plot, write_flag
  from llm_studio.src.utils.modeling_utils import activate_neftune, check_disk_space, get_ds_config, get_number_of_validation_epochs, get_optimizer, get_scheduler, load_checkpoint, reduce_metric, run_inference, save_checkpoint, save_predictions, unwrap_model, wrap_model_distributed
  from llm_studio.src.utils.utils import kill_ddp_processes, set_environment, set_seed
label_desc: Use image node_id 2 for calling a global function with example usage: run_train(cfg, model, optimizer, scheduler, epoch_steps, train_dataloader, val_dataloader, val_df) and returns: val_loss,  val_metric