code: def run_train_rlhf(
      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
      """
  
      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,
          )
  
      with torch.device(cfg.environment._device):
          logger.info("Using RLHF - Loading reward model")
          reward_model = cfg.architecture.reward_model_class(cfg)
          reward_model.eval()
  
      if cfg.training.offload_reward_model:
          reward_model.to("cpu")
      else:
          reward_model.to(cfg.environment._device)
  
      # initialize trainer
      tokenizer = get_tokenizer(cfg)
      ppo_trainer = PPOTrainer(
          cfg=cfg,
          model=model,
          tokenizer=tokenizer,
          optimizer=optimizer,
          lr_scheduler=scheduler,
          scaler=scaler,
      )
  
      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 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"Target evaluation step: {evaluation_step}")
  
          # Round up to the nearest multiple of cfg.training.rollout_steps
          assert cfg.training.rollout_steps <= epoch_steps, (
              f"Rollout steps ({cfg.training.rollout_steps}) must be less than the total "
              f"training steps in one epoch ({epoch_steps})"
          )
  
          # Round up to the next multiple of rollout steps
          evaluation_step = (
              (evaluation_step + cfg.training.rollout_steps - 1)
              // cfg.training.rollout_steps
          ) * cfg.training.rollout_steps
  
          logger.info(f"Actual evaluation step: {evaluation_step}")
  
          query_tensors = []
          response_tensors = []
          rewards = []
  
          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")
  
              with torch.no_grad():
                  logger.debug(
                      "Rollout: Generating response from active model"
                  )
                  output_dict = {}
                  output_dict["predicted_answer_ids"] = (
                      unwrap_model(model)
                      .generate(batch, unwrap_model(model).cfg)
                      .detach()
                  )
                  output_dict = train_dataloader.dataset.postprocess_batch_predictions(
                      output=output_dict
                  )
  
                  logger.debug("Evaluation: Score from reward model")
                  # tokenize prompt & output internally
                  if cfg.training.offload_reward_model:
                      reward_model.to(cfg.environment._device)
                  context = (
                      autocast(enabled=cfg.environment.mixed_precision)
                      if (
                          torch.cuda.is_available()
                          and len(cfg.environment.gpus) > 0
                      )
                      else torch.no_grad()
                  )
                  with context:  # type: ignore[attr-defined]
                      scores = reward_model.get_score(
                          batch["reward_model_prompt_text"],
                          output_dict["predicted_text"],
                      )
  
                  if cfg.training.offload_reward_model:
                      reward_model.to("cpu")
  
              # score by reward model
              reward = [
                  torch.tensor(score, dtype=torch.float32)
                  for score in scores
              ]
  
              # remove padding from query and response
              batch["input_ids"] = batch["input_ids"].detach().cpu()
              query_tensor = [
                  input_ids[torch.where(att_mask == 1)[0].min() :]
                  if len(torch.where(att_mask == 1)[0]) > 0
                  else input_ids
                  for input_ids, att_mask in zip(
                      batch["input_ids"].detach().cpu(),
                      batch["attention_mask"],
                  )
              ]
              pad_tok_id = (
                  unwrap_model(model).backbone.config.pad_token_id
                  or unwrap_model(model).backbone.config.eos_token_id
              )
              output_dict["predicted_answer_ids"] = (
                  output_dict["predicted_answer_ids"].detach().cpu()
              )
              response_tensor = [
                  predicted_answer_ids[
                      : torch.where(predicted_answer_ids == pad_tok_id)[
                          0
                      ].min()
                  ]
                  if len(
                      torch.where(predicted_answer_ids == pad_tok_id)[0]
                  )
                  > 0
                  else predicted_answer_ids
                  for predicted_answer_ids in output_dict[
                      "predicted_answer_ids"
                  ]
              ]
  
              del output_dict
              del batch
  
              query_tensors += query_tensor
              response_tensors += response_tensor
              rewards += reward
  
              if cfg.environment._distributed:
                  torch.cuda.synchronize(
                      device=cfg.environment._local_rank
                  )
                  torch.distributed.barrier()
  
              if (itr + 1) % cfg.training.rollout_steps == 0:
                  output_dict = ppo_trainer.step(
                      query_tensors, response_tensors, rewards
                  )
                  del query_tensors, response_tensors, rewards, scores
  
                  query_tensors = []
                  response_tensors = []
                  rewards = []
  
                  loss = output_dict["ppo/loss/total"]
                  losses.append(loss)
  
                  if cfg.environment._local_rank == 0:
                      for key in output_dict.keys():
                          if isinstance(
                              output_dict[key], (float, int)
                          ) or (
                              isinstance(output_dict[key], np.ndarray)
                              and output_dict[key].size == 1
                          ):
                              if np.isfinite(output_dict[key]):
                                  cfg.logging._logger.log(
                                      "train",
                                      key,
                                      output_dict[key],
                                      step=cfg.environment._curr_step,
                                  )
                      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,
                      )
  
                      del output_dict
  
              if cfg.environment._local_rank == 0:
                  # Show logs each 5% of the epoch (only if doing per epoch eval)
                  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
                          )
  
              if cfg.environment._distributed:
                  torch.cuda.synchronize(
                      device=cfg.environment._local_rank
                  )
                  torch.distributed.barrier()
  
              # Validation loop
              if ((itr + 1) % evaluation_step == 0) or (
                  itr == epoch_steps - 1
              ):
                  if cfg.training.evaluation_epochs == 1:
                      progress_bar.close()
  
                  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
                  else:
                      checkpoint_path = cfg.output_directory
                      if cfg.environment._local_rank == 0:
                          logger.info(
                              f"Saving last model checkpoint: "
                              f"val_loss {val_loss:.5}, val_{cfg.prediction.metric} "
                              f"{val_metric:.5} to {checkpoint_path}"
                          )
                      save_checkpoint(
                          model=model, path=checkpoint_path, cfg=cfg
                      )
  
                  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 3 for calling a global function with example usage: run_train_rlhf(cfg, model, optimizer, scheduler, epoch_steps, train_dataloader, val_dataloader, val_df) and returns: val_loss,  val_metric