code: def train(
      model,
      train_dataloader,
      eval_dataloader,
      tokenizer,
      optimizer,
      lr_scheduler,
      gradient_accumulation_steps,
      train_config,
      fsdp_config=None,
      local_rank=None,
      rank=None,
  ):
      """
      Trains the model on the given dataloader
  
      Args:
          model: The model to be trained
          train_dataloader: The dataloader containing the training data
          optimizer: The optimizer used for training
          lr_scheduler: The learning rate scheduler
          gradient_accumulation_steps: The number of steps to accumulate gradients before performing a backward/update operation
          num_epochs: The number of epochs to train for
          local_rank: The rank of the current node in a distributed setting
          train_config: The training configuration
          eval_dataloader: The dataloader containing the eval data
          tokenizer: tokenizer used in the eval for decoding the predicitons
  
      Returns: results dictionary containing average training and validation perplexity and loss
      """
      # Create a gradient scaler for fp16
      if train_config.use_fp16 and train_config.enable_fsdp:
          scaler = ShardedGradScaler()
      elif train_config.use_fp16 and not train_config.enable_fsdp:
          scaler = torch.cuda.amp.GradScaler()
      if train_config.enable_fsdp:
          world_size = int(os.environ["WORLD_SIZE"])
  
      autocast = (
          torch.cuda.amp.autocast
          if train_config.use_fp16
          else nullcontext
      )
  
      train_prep = []
      train_loss = []
      val_prep = []
      val_loss = []
  
      if train_config.save_metrics:
          metrics_filename = f"{train_config.output_dir}/metrics_data_{local_rank}-{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.json"
          train_step_perplexity = []
          train_step_loss = []
          val_step_loss = []
          val_step_perplexity = []
  
      epoch_times = []
      checkpoint_times = []
      results = {}
      best_val_loss = float("inf")
      for epoch in range(train_config.num_epochs):
          epoch_start_time = time.perf_counter()
          with MemoryTrace() as memtrace:  # track the memory usage
              model.train()
              total_loss = 0.0
              total_length = (
                  len(train_dataloader) // gradient_accumulation_steps
              )
              pbar = tqdm(
                  colour="blue",
                  desc=f"Training Epoch: {epoch+1}",
                  total=total_length,
                  dynamic_ncols=True,
              )
              for step, batch in enumerate(train_dataloader):
                  for key in batch.keys():
                      if train_config.enable_fsdp:
                          if is_xpu_available():
                              batch[key] = batch[key].to(
                                  torch.device(f"xpu:{local_rank}")
                              )
                          else:
                              batch[key] = batch[key].to(local_rank)
                      else:
                          if is_xpu_available():
                              batch[key] = batch[key].to("xpu:0")
                          else:
                              batch[key] = batch[key].to("cuda:0")
                  with autocast():
                      loss = model(**batch).loss
                  loss = loss / gradient_accumulation_steps
                  if train_config.save_metrics:
                      train_step_loss.append(
                          loss.detach().float().item()
                      )
                      train_step_perplexity.append(
                          float(torch.exp(loss.detach().float()))
                      )
                  total_loss += loss.detach().float()
                  if train_config.use_fp16:
                      # if fp16 is enabled, use gradient scaler to handle gradient update
                      scaler.scale(loss).backward()
                      if (
                          (step + 1) % gradient_accumulation_steps == 0
                          or step == len(train_dataloader) - 1
                      ):
                          if (
                              train_config.gradient_clipping
                              and train_config.gradient_clipping_threshold
                              > 0.0
                          ):
                              scaler.unscale_(optimizer)
                              if train_config.enable_fsdp:
                                  model.clip_grad_norm_(
                                      train_config.gradient_clipping_threshold
                                  )
                              else:
                                  torch.nn.utils.clip_grad_norm_(
                                      model.parameters(),
                                      train_config.gradient_clipping_threshold,
                                  )
                          scaler.step(optimizer)
                          scaler.update()
                          optimizer.zero_grad()
                          pbar.update(1)
                  else:
                      # regular backpropagation when fp16 is not used
                      loss.backward()
                      if (
                          (step + 1) % gradient_accumulation_steps == 0
                          or step == len(train_dataloader) - 1
                      ):
                          if (
                              train_config.gradient_clipping
                              and train_config.gradient_clipping_threshold
                              > 0.0
                          ):
                              if train_config.enable_fsdp:
                                  model.clip_grad_norm_(
                                      train_config.gradient_clipping_threshold
                                  )
                              else:
                                  torch.nn.utils.clip_grad_norm_(
                                      model.parameters(),
                                      train_config.gradient_clipping_threshold,
                                  )
                          optimizer.step()
                          optimizer.zero_grad()
                          pbar.update(1)
  
                  pbar.set_description(
                      f"Training Epoch: {epoch+1}/{train_config.num_epochs}, step {step}/{len(train_dataloader)} completed (loss: {loss.detach().float()})"
                  )
  
                  if train_config.save_metrics:
                      save_to_json(
                          metrics_filename,
                          train_step_loss,
                          train_loss,
                          train_step_perplexity,
                          train_prep,
                          val_step_loss,
                          val_loss,
                          val_step_perplexity,
                          val_prep,
                      )
              pbar.close()
  
          epoch_end_time = time.perf_counter() - epoch_start_time
          epoch_times.append(epoch_end_time)
          # Reducing total_loss across all devices if there's more than one CUDA device
          if is_xpu_available() and (
              torch.xpu.device_count() > 1 and train_config.enable_fsdp
          ):
              dist.all_reduce(total_loss, op=dist.ReduceOp.SUM)
          elif (
              torch.cuda.device_count() > 1 and train_config.enable_fsdp
          ):
              dist.all_reduce(total_loss, op=dist.ReduceOp.SUM)
          train_epoch_loss = total_loss / len(train_dataloader)
          if train_config.enable_fsdp:
              train_epoch_loss = train_epoch_loss / world_size
          train_perplexity = torch.exp(train_epoch_loss)
  
          train_prep.append(float(train_perplexity))
          train_loss.append(float(train_epoch_loss))
  
          if train_config.enable_fsdp:
              if rank == 0:
                  if is_xpu_available():
                      print(
                          f"Max XPU memory allocated was {memtrace.peak} GB"
                      )
                      print(
                          f"Max XPU memory reserved was {memtrace.max_reserved} GB"
                      )
                      print(
                          f"Peak active XPU memory was {memtrace.peak_active_gb} GB"
                      )
                      print(
                          f"Xpu Malloc retires : {memtrace.xpu_malloc_retires}"
                      )
                  else:
                      print(
                          f"Max CUDA memory allocated was {memtrace.peak} GB"
                      )
                      print(
                          f"Max CUDA memory reserved was {memtrace.max_reserved} GB"
                      )
                      print(
                          f"Peak active CUDA memory was {memtrace.peak_active_gb} GB"
                      )
                      print(
                          f"Cuda Malloc retires : {memtrace.cuda_malloc_retires}"
                      )
                  print(
                      f"CPU Total Peak Memory consumed during the train (max): {memtrace.cpu_peaked + memtrace.cpu_begin} GB"
                  )
          else:
              if is_xpu_available():
                  print(
                      f"Max XPU memory allocated was {memtrace.peak} GB"
                  )
                  print(
                      f"Max XPU memory reserved was {memtrace.max_reserved} GB"
                  )
                  print(
                      f"Peak active XPU memory was {memtrace.peak_active_gb} GB"
                  )
                  print(
                      f"Xpu Malloc retires : {memtrace.xpu_malloc_retires}"
                  )
              else:
                  print(
                      f"Max CUDA memory allocated was {memtrace.peak} GB"
                  )
                  print(
                      f"Max CUDA memory reserved was {memtrace.max_reserved} GB"
                  )
                  print(
                      f"Peak active CUDA memory was {memtrace.peak_active_gb} GB"
                  )
                  print(
                      f"Cuda Malloc retires : {memtrace.cuda_malloc_retires}"
                  )
              print(
                  f"CPU Total Peak Memory consumed during the train (max): {memtrace.cpu_peaked + memtrace.cpu_begin} GB"
              )
  
          # Update the learning rate as needed
          lr_scheduler.step()
  
          if train_config.run_validation:
              (
                  eval_ppl,
                  eval_epoch_loss,
                  temp_val_loss,
                  temp_step_perplexity,
              ) = evaluation(
                  model,
                  train_config,
                  eval_dataloader,
                  local_rank,
                  tokenizer,
              )
              if train_config.save_metrics:
                  val_step_loss.extend(temp_val_loss)
                  val_step_perplexity.extend(temp_step_perplexity)
  
              checkpoint_start_time = time.perf_counter()
              if (
                  train_config.save_model
                  and eval_epoch_loss < best_val_loss
              ):
                  if train_config.enable_fsdp:
                      dist.barrier()
                  if train_config.use_peft:
                      if train_config.enable_fsdp:
                          if rank == 0:
                              print(
                                  f"we are about to save the PEFT modules"
                              )
                      else:
                          print(
                              f"we are about to save the PEFT modules"
                          )
                      model.save_pretrained(train_config.output_dir)
                      if train_config.enable_fsdp:
                          if rank == 0:
                              print(
                                  f"PEFT modules are saved in {train_config.output_dir} directory"
                              )
                      else:
                          print(
                              f"PEFT modules are saved in {train_config.output_dir} directory"
                          )
  
                  else:
                      if (
                          not train_config.use_peft
                          and fsdp_config.checkpoint_type
                          == StateDictType.FULL_STATE_DICT
                      ):
                          save_model_checkpoint(
                              model,
                              optimizer,
                              rank,
                              train_config,
                              epoch=epoch,
                          )
                      elif (
                          not train_config.use_peft
                          and fsdp_config.checkpoint_type
                          == StateDictType.SHARDED_STATE_DICT
                      ):
                          print(
                              " Saving the FSDP model checkpoints using SHARDED_STATE_DICT"
                          )
                          print(
                              "====================================================="
                          )
  
                          save_model_and_optimizer_sharded(
                              model, rank, train_config
                          )
                          if train_config.save_optimizer:
                              save_model_and_optimizer_sharded(
                                  model,
                                  rank,
                                  train_config,
                                  optim=optimizer,
                              )
                              print(
                                  " Saving the FSDP model checkpoints and optimizer using SHARDED_STATE_DICT"
                              )
                              print(
                                  "====================================================="
                              )
  
                      if (
                          not train_config.use_peft
                          and train_config.save_optimizer
                      ):
                          save_optimizer_checkpoint(
                              model,
                              optimizer,
                              rank,
                              train_config,
                              epoch=epoch,
                          )
                          print(
                              " Saving the FSDP model checkpoints and optimizer using FULL_STATE_DICT"
                          )
                          print(
                              "====================================================="
                          )
                  if train_config.enable_fsdp:
                      dist.barrier()
              checkpoint_end_time = (
                  time.perf_counter() - checkpoint_start_time
              )
              checkpoint_times.append(checkpoint_end_time)
              if eval_epoch_loss < best_val_loss:
                  best_val_loss = eval_epoch_loss
                  if train_config.enable_fsdp:
                      if rank == 0:
                          print(
                              f"best eval loss on epoch {epoch+1} is {best_val_loss}"
                          )
                  else:
                      print(
                          f"best eval loss on epoch {epoch+1} is {best_val_loss}"
                      )
              val_loss.append(float(best_val_loss))
              val_prep.append(float(eval_ppl))
          if train_config.enable_fsdp:
              if rank == 0:
                  print(
                      f"Epoch {epoch+1}: train_perplexity={train_perplexity:.4f}, train_epoch_loss={train_epoch_loss:.4f}, epoch time {epoch_end_time}s"
                  )
          else:
              print(
                  f"Epoch {epoch+1}: train_perplexity={train_perplexity:.4f}, train_epoch_loss={train_epoch_loss:.4f}, epoch time {epoch_end_time}s"
              )
  
          # Saving the results every epoch to plot later
          if train_config.save_metrics:
              save_to_json(
                  metrics_filename,
                  train_step_loss,
                  train_loss,
                  train_step_perplexity,
                  train_prep,
                  val_step_loss,
                  val_loss,
                  val_step_perplexity,
                  val_prep,
              )
  
      avg_epoch_time = sum(epoch_times) / len(epoch_times)
      avg_checkpoint_time = (
          sum(checkpoint_times) / len(checkpoint_times)
          if len(checkpoint_times) > 0
          else 0
      )
      avg_train_prep = sum(train_prep) / len(train_prep)
      avg_train_loss = sum(train_loss) / len(train_loss)
      if train_config.run_validation:
          avg_eval_prep = sum(val_prep) / len(val_prep)
          avg_eval_loss = sum(val_loss) / len(val_loss)
  
      results["avg_train_prep"] = avg_train_prep
      results["avg_train_loss"] = avg_train_loss
      if train_config.run_validation:
          results["avg_eval_prep"] = avg_eval_prep
          results["avg_eval_loss"] = avg_eval_loss
      results["avg_epoch_time"] = avg_epoch_time
      results["avg_checkpoint_time"] = avg_checkpoint_time
      if train_config.save_metrics:
          results["metrics_filename"] = metrics_filename
  
      # saving the training params including fsdp setting for reference.
      if train_config.enable_fsdp and not train_config.use_peft:
          save_train_params(train_config, fsdp_config, rank)
  
      return results
  
imports: import os
  import time
  import yaml
  from contextlib import nullcontext
  from pathlib import Path
  from pkg_resources import packaging
  from datetime import datetime
  import torch
  import torch.cuda.nccl
  import torch.distributed
  from torch.distributed.fsdp import StateDictType
  from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
  from tqdm import tqdm
  from transformers import LlamaTokenizer
  import json
  from llama_recipes.model_checkpointing import save_model_checkpoint, save_model_and_optimizer_sharded, save_optimizer_checkpoint
  from llama_recipes.policies import fpSixteen, bfSixteen, get_llama_wrapper
  from llama_recipes.utils.memory_utils import MemoryTrace
  from accelerate.utils import is_xpu_available, is_ccl_available
label_desc: Use image node_id 3 for calling a global function with example usage: train(model, train_dataloader, eval_dataloader, tokenizer, optimizer, lr_scheduler, gradient_accumulation_steps, train_config, fsdp_config, local_rank, rank) and returns: results