code: def main():
      args = parse_args()
  
      if args.local_rank == -1:
          device = torch.device(get_accelerator().device_name())
      else:
          get_accelerator().set_device(args.local_rank)
          device = torch.device(
              get_accelerator().device_name(), args.local_rank
          )
          # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
          # torch.distributed.init_process_group(backend='nccl')
          deepspeed.init_distributed()
  
      args.global_rank = torch.distributed.get_rank()
  
      ds_config = get_train_ds_config(
          offload=args.offload,
          dtype=args.dtype,
          stage=args.zero_stage,
          enable_tensorboard=args.enable_tensorboard,
          tb_path=args.tensorboard_path,
          tb_name="step2_model",
      )
      ds_config[
          "train_micro_batch_size_per_gpu"
      ] = args.per_device_train_batch_size
      ds_config["train_batch_size"] = (
          args.per_device_train_batch_size
          * torch.distributed.get_world_size()
          * args.gradient_accumulation_steps
      )
  
      # If passed along, set the training seed now.
      set_random_seed(args.seed)
      torch.distributed.barrier()
  
      # load_hf_tokenizer will get the correct tokenizer and set padding tokens based on the model family
      args.end_of_conversation_token = "<|endoftext|>"
      additional_special_tokens = (
          args.end_of_conversation_token if args.add_eot_token else None
      )
      tokenizer = load_hf_tokenizer(
          args.model_name_or_path,
          fast_tokenizer=True,
          add_special_tokens=additional_special_tokens,
      )
      rm_model = create_critic_model(
          args.model_name_or_path,
          tokenizer,
          ds_config,
          args.num_padding_at_beginning,
          dropout=args.dropout,
          zero_stage=args.zero_stage,
          compute_fp32_loss=args.compute_fp32_loss,
      )
  
      # Model bigscience/bloom-560m has large variance at ln_f.weight parameter
      # This makes bf16 finetuning hard.
      # In general, since we are replacing the model head, it makes sense to reset
      # the LN that precedes it.
      force_optimize_params = []
      if "bigscience/bloom-" in args.model_name_or_path:
          zero_init_enabled = args.zero_stage == 3
          params = [
              rm_model.rwtranrsformer.ln_f.weight,
              rm_model.rwtranrsformer.ln_f.bias,
          ]
          with deepspeed.zero.GatheredParameters(
              params, modifier_rank=0, enabled=zero_init_enabled
          ):
              if (
                  deepspeed.comm.get_rank() == 0
                  or not zero_init_enabled
              ):
                  torch.nn.init.ones_(
                      rm_model.rwtransformer.ln_f.weight
                  )
                  torch.nn.init.zeros_(rm_model.rwtransformer.ln_f.bias)
          force_optimize_params.extend(
              ["rwtransformer.ln_f.weight", "rwtransformer.ln_f.bias"]
          )
  
      if args.lora_dim > 0:
          rm_model = convert_linear_layer_to_lora(
              rm_model, args.lora_module_name, args.lora_dim
          )
          if args.only_optimize_lora:
              force_optimize_params.append("v_head.weight")
              rm_model = only_optimize_lora_parameters(
                  rm_model, force_optimize_params
              )
              rm_model = make_model_gradient_checkpointing_compatible(
                  rm_model
              )
  
      train_phase = 2
      train_dataset, eval_dataset = create_prompt_dataset(
          args.local_rank,
          args.data_path,
          args.data_split,
          args.data_output_path,
          train_phase,
          args.seed,
          tokenizer,
          args.max_seq_len,
      )
  
      # DataLoaders creation:
      data_collator = DataCollatorReward()
      if args.local_rank == -1:
          train_sampler = RandomSampler(train_dataset)
          eval_sampler = SequentialSampler(eval_dataset)
      else:
          train_sampler = DistributedSampler(train_dataset)
          eval_sampler = DistributedSampler(eval_dataset)
      train_dataloader = DataLoader(
          train_dataset,
          collate_fn=data_collator,
          sampler=train_sampler,
          batch_size=args.per_device_train_batch_size,
      )
      eval_dataloader = DataLoader(
          eval_dataset,
          collate_fn=data_collator,
          sampler=eval_sampler,
          batch_size=args.per_device_eval_batch_size,
      )
  
      def evaluation_reward(model, dataloader, eval_iters):
          model.eval()
          correct_predictions = 0
          total_predictions = 0
          chosen_scores = 0.0
          rejected_scores = 0.0
          for _step, _batch in enumerate(dataloader):
              _batch = to_device(_batch, device)
              with torch.no_grad():
                  _outputs = model(**_batch)
  
              chosen = _outputs["chosen_mean_scores"]
              rejected = _outputs["rejected_mean_scores"]
              correct_predictions += (chosen > rejected).sum()
              total_predictions += chosen.shape[0]
              chosen_scores += (
                  _outputs["chosen_mean_scores"].mean().float()
              )
              rejected_scores += (
                  _outputs["rejected_mean_scores"].mean().float()
              )
              if (_step + 1) == eval_iters:
                  break
          _acc = correct_predictions / total_predictions
          chosen_scores = chosen_scores / (_step + 1)
          rejected_scores = rejected_scores / (_step + 1)
          try:
              _acc = get_all_reduce_mean(_acc).item()
              chosen_scores = get_all_reduce_mean(chosen_scores).item()
              rejected_scores = get_all_reduce_mean(
                  rejected_scores
              ).item()
          except:
              pass
          return chosen_scores, rejected_scores, _acc
  
      # Split weights in two groups, one with weight decay and the other not.
      optimizer_grouped_parameters = get_optimizer_grouped_parameters(
          rm_model, args.weight_decay, args.lora_learning_rate
      )
  
      AdamOptimizer = DeepSpeedCPUAdam if args.offload else FusedAdam
      optimizer = AdamOptimizer(
          optimizer_grouped_parameters,
          lr=args.learning_rate,
          betas=(0.9, 0.95),
      )
  
      num_update_steps_per_epoch = math.ceil(
          len(train_dataloader) / args.gradient_accumulation_steps
      )
  
      lr_scheduler = get_scheduler(
          name=args.lr_scheduler_type,
          optimizer=optimizer,
          num_warmup_steps=args.num_warmup_steps,
          num_training_steps=args.num_train_epochs
          * num_update_steps_per_epoch,
      )
  
      rm_model, optimizer, _, lr_scheduler = deepspeed.initialize(
          model=rm_model,
          optimizer=optimizer,
          args=args,
          config=ds_config,
          lr_scheduler=lr_scheduler,
          dist_init_required=True,
      )
  
      if args.gradient_checkpointing:
          rm_model.gradient_checkpointing_enable()
  
      # Train!
      print_rank_0("***** Running training *****", args.global_rank)
  
      print_rank_0(
          f"***** Evaluating reward, Epoch {0}/{args.num_train_epochs} *****",
          args.global_rank,
      )
      reward_score, reject_score, acc = evaluation_reward(
          rm_model, eval_dataloader, args.eval_iters
      )
      print_rank_0(
          f"chosen_last_scores (higher is better) : {reward_score}, "
          f"rejected_last_scores (lower is better) : {reject_score}, "
          f"acc (higher is better) : {acc}",
          args.global_rank,
      )
  
      total_micro_steps = 0
      for epoch in range(args.num_train_epochs):
          print_rank_0(
              f"Beginning of Epoch {epoch+1}/{args.num_train_epochs}, Total Micro Batches {len(train_dataloader)}",
              args.global_rank,
          )
          rm_model.train()
          mean_loss = 0
          for step, batch in enumerate(train_dataloader):
              batch = to_device(batch, device)
              outputs = rm_model(**batch, use_cache=False)
              loss = outputs["loss"]
              rm_model.backward(loss)
              rm_model.step()
              mean_loss += loss.item()
              total_micro_steps += 1
              gas_boundary = (
                  total_micro_steps % args.gradient_accumulation_steps
                  == 0
              )
              total_steps = (
                  total_micro_steps // args.gradient_accumulation_steps
              )
              if (
                  args.eval_interval
                  and gas_boundary
                  and (total_steps % args.eval_interval == 0)
              ):
                  print_rank_0(
                      f"Iter {total_steps}: Evaluating reward",
                      args.global_rank,
                  )
                  reward_score, reject_score, acc = evaluation_reward(
                      rm_model, eval_dataloader, args.eval_iters
                  )
                  print_rank_0(
                      f"Iter {total_steps}: c_scores: {reward_score}, r_scores: {reject_score}, "
                      f"diff: {reward_score - reject_score}, acc: {acc}",
                      args.global_rank,
                  )
                  rm_model.train()
  
          print_rank_0(
              f"Epoch {epoch+1}/{args.num_train_epochs} with loss {mean_loss/(step+1)}",
              args.global_rank,
          )
          # Evaluate reward_loss on the validation set.
          print_rank_0(
              f"***** Evaluating reward, Epoch {epoch+1}/{args.num_train_epochs} *****",
              args.global_rank,
          )
          reward_score, reject_score, acc = evaluation_reward(
              rm_model, eval_dataloader, args.eval_iters
          )
          print_rank_0(
              f"chosen_last_scores (higher is better) : {reward_score}, "
              f"rejected_last_scores (lower is better) : {reject_score}, "
              f"acc (higher is better) : {acc}",
              args.global_rank,
          )
          rm_model.tput_timer.update_epoch_count()
  
      if args.output_dir is not None:
          print_rank_0("saving model ...", args.global_rank)
          rm_model = convert_lora_to_linear_layer(rm_model)
  
          if args.global_rank == 0:
              save_hf_format(rm_model, tokenizer, args)
          if args.zero_stage == 3:
              # for zero stage 3, each gpu only has a part of the model, so we need to save the model on each gpu by using DS-Engine
              save_zero_three_model(
                  rm_model,
                  args.global_rank,
                  args.output_dir,
                  zero_stage=args.zero_stage,
              )
  
file_path: deepspeedexamples/applications/DeepSpeed-Chat/training/step2_reward_model_finetuning/main.py
imports: import argparse
  import math
  import torch
  from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
  from torch.utils.data.distributed import DistributedSampler
  from transformers import SchedulerType, get_scheduler
  import deepspeed
  from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
  from deepspeed.accelerator import get_accelerator
  from dschat.utils.model.model_utils import create_critic_model
  from dschat.utils.data.data_utils import create_prompt_dataset, DataCollatorReward
  from dschat.utils.utils import print_rank_0, to_device, save_hf_format, set_random_seed, get_all_reduce_mean, get_optimizer_grouped_parameters, save_zero_three_model, load_hf_tokenizer
  from dschat.utils.ds_utils import get_train_ds_config
  from dschat.utils.module.lora import convert_linear_layer_to_lora, convert_lora_to_linear_layer, only_optimize_lora_parameters, make_model_gradient_checkpointing_compatible
label_desc: Use image node_id 2 for calling a global function with example usage: main() and returns: chosen_scores,  rejected_scores,  _acc