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="step1_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,
      )
  
      model = create_hf_model(
          AutoModelForCausalLM,
          args.model_name_or_path,
          tokenizer,
          ds_config,
          dropout=args.dropout,
      )
  
      if args.compute_fp32_loss:
          print_rank_0(
              f"Using model {model.__class__.__name__} with loss in fp32",
              args.global_rank,
          )
          causal_lm_model_to_fp32_loss(model)
  
      if args.lora_dim > 0:
          model = convert_linear_layer_to_lora(
              model, args.lora_module_name, args.lora_dim
          )
          if args.only_optimize_lora:
              model = only_optimize_lora_parameters(model)
              model = make_model_gradient_checkpointing_compatible(
                  model
              )
  
      # Prepare the data
      train_phase = 1
      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,
          sft_only_data_path=args.sft_only_data_path,
      )
      # DataLoaders creation:
      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=default_data_collator,
          sampler=train_sampler,
          batch_size=args.per_device_train_batch_size,
      )
      eval_dataloader = DataLoader(
          eval_dataset,
          collate_fn=default_data_collator,
          sampler=eval_sampler,
          batch_size=args.per_device_eval_batch_size,
      )
  
      def evaluation(model, eval_dataloader):
          model.eval()
          losses = 0
          for step, batch in enumerate(eval_dataloader):
              batch = to_device(batch, device)
              with torch.no_grad():
                  outputs = model(**batch)
  
              loss = outputs.loss
              losses += loss.float()
          losses = losses / (step + 1)
          try:
              losses = get_all_reduce_mean(losses)
          except:
              pass
          try:
              perplexity = torch.exp(losses).item()
          except OverflowError:
              perplexity = float("inf")
          return perplexity, losses.item()
  
      # Split weights in two groups, one with weight decay and the other not.
      optimizer_grouped_parameters = get_optimizer_grouped_parameters(
          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,
      )
  
      model, optimizer, _, lr_scheduler = deepspeed.initialize(
          model=model,
          optimizer=optimizer,
          args=args,
          config=ds_config,
          lr_scheduler=lr_scheduler,
          dist_init_required=True,
      )
  
      if args.gradient_checkpointing:
          model.gradient_checkpointing_enable()
  
      # Train!
      print_rank_0("***** Running training *****", args.global_rank)
      print_rank_0(
          f"***** Evaluating perplexity, Epoch {0}/{args.num_train_epochs} *****",
          args.global_rank,
      )
      perplexity, eval_loss = evaluation(model, eval_dataloader)
      print_rank_0(
          f"ppl: {perplexity}, loss: {eval_loss}", args.global_rank
      )
  
      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,
          )
          model.train()
          import time
  
          for step, batch in enumerate(train_dataloader):
              start = time.time()
              batch = to_device(batch, device)
              outputs = model(**batch, use_cache=False)
              loss = outputs.loss
              if args.print_loss:
                  print(
                      f"Epoch: {epoch}, Step: {step}, Rank: {torch.distributed.get_rank()}, loss = {loss}"
                  )
              model.backward(loss)
              model.step()
              end = time.time()
              if torch.distributed.get_rank() == 0:
                  print_throughput(
                      model.model, args, end - start, args.global_rank
                  )
  
          # Evaluate perplexity on the validation set.
          print_rank_0(
              f"***** Evaluating perplexity, Epoch {epoch+1}/{args.num_train_epochs} *****",
              args.global_rank,
          )
          perplexity, eval_loss = evaluation(model, eval_dataloader)
          print_rank_0(
              f"ppl: {perplexity}, loss: {eval_loss}", args.global_rank
          )
          model.tput_timer.update_epoch_count()
  
      if args.output_dir is not None:
          print_rank_0("saving the final model ...", args.global_rank)
          model = convert_lora_to_linear_layer(model)
  
          if args.global_rank == 0:
              save_hf_format(model, tokenizer, args)
  
          if args.zero_stage == 3:
              # For zero stage 3, each gpu only has a part of the model, so we need a special save function
              save_zero_three_model(
                  model,
                  args.global_rank,
                  args.output_dir,
                  zero_stage=args.zero_stage,
              )
  
file_path: deepspeedexamples/applications/DeepSpeed-Chat/training/step1_supervised_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 AutoModelForCausalLM, SchedulerType, default_data_collator, get_scheduler
  import deepspeed
  from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
  from deepspeed import get_accelerator
  from dschat.utils.data.data_utils import create_prompt_dataset
  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
  from dschat.utils.model.model_utils import create_hf_model, causal_lm_model_to_fp32_loss
  from dschat.utils.perf import print_throughput
label_desc: Use image node_id 2 for calling a global function with example usage: main() and returns: perplexity,  losses