code: def load_checkpoint(
      model,
      optimizer,
      opt_param_scheduler,
      load_arg="load",
      strict=True,
  ):
      """Load a model checkpoint and return the iteration.
      strict (bool): whether to strictly enforce that the keys in
          :attr:`state_dict` of the checkpoint match the names of
          parameters and buffers in model.
      """
      args = get_args()
      load_dir = getattr(args, load_arg)
  
      model = unwrap_model(model)
  
      state_dict, checkpoint_name, release = _load_base_checkpoint(
          load_dir, rank0=False
      )
  
      # Checkpoint not loaded.
      if state_dict is None:
          # Conditionally exit at this point.
          if args.exit_on_missing_checkpoint:
              print_rank_0(
                  ">> '--exit-on-missing-checkpoint' set ... exiting. <<"
              )
              torch.distributed.barrier()
              sys.exit()
  
          # Iteration defaults to 0.
          return 0
  
      # Set checkpoint version.
      set_checkpoint_version(state_dict.get("checkpoint_version", 0))
  
      # Set iteration.
      if args.finetune or release:
          iteration = 0
      else:
          try:
              iteration = state_dict["iteration"]
          except KeyError:
              try:  # Backward compatible with older checkpoints
                  iteration = state_dict["total_iters"]
              except KeyError:
                  print_rank_0(
                      "A metadata file exists but unable to load "
                      "iteration from checkpoint {}, exiting".format(
                          checkpoint_name
                      )
                  )
                  sys.exit()
  
      # Check arguments.
      assert args.consumed_train_samples == 0
      assert args.consumed_valid_samples == 0
      if "args" in state_dict and not args.finetune:
          checkpoint_args = state_dict["args"]
          check_checkpoint_args(checkpoint_args)
          args.consumed_train_samples = getattr(
              checkpoint_args, "consumed_train_samples", 0
          )
          update_num_microbatches(
              consumed_samples=args.consumed_train_samples
          )
          args.consumed_valid_samples = getattr(
              checkpoint_args, "consumed_valid_samples", 0
          )
      else:
          print_rank_0("could not find arguments in the checkpoint ...")
  
      # Model.
      strict = (
          False
          if args.retro_add_retriever
          or args.transformer_impl == "transformer_engine"
          else strict
      )
      if len(model) == 1:
          model[0].load_state_dict(state_dict["model"], strict=strict)
      else:
          for i in range(len(model)):
              mpu.set_virtual_pipeline_model_parallel_rank(i)
              model[i].load_state_dict(
                  state_dict["model%d" % i], strict=strict
              )
  
      # Fix up query/key/value matrix ordering if needed.
      checkpoint_version = get_checkpoint_version()
      print_rank_0(f" checkpoint version {checkpoint_version}")
      fix_query_key_value_ordering(model, checkpoint_version)
  
      # Optimizer.
      if not release and not args.finetune and not args.no_load_optim:
          try:
              # Load state dict.
              if optimizer is not None:
                  optimizer.load_state_dict(state_dict["optimizer"])
  
              # Load distributed optimizer's custom parameter state.
              if args.use_distributed_optimizer:
                  tracker_filename = get_checkpoint_tracker_filename(
                      load_dir
                  )
                  iteration, release = read_metadata(tracker_filename)
                  model_checkpoint_name = get_checkpoint_name(
                      load_dir, iteration, release
                  )
                  optim_checkpoint_name = (
                      get_distributed_optimizer_checkpoint_name(
                          model_checkpoint_name
                      )
                  )
                  optimizer.load_parameter_state(optim_checkpoint_name)
  
              # Load scheduler.
              if opt_param_scheduler is not None:
                  if (
                      "lr_scheduler" in state_dict
                  ):  # backward compatbility
                      opt_param_scheduler.load_state_dict(
                          state_dict["lr_scheduler"]
                      )
                  else:
                      opt_param_scheduler.load_state_dict(
                          state_dict["opt_param_scheduler"]
                      )
          except KeyError:
              print_rank_0(
                  "Unable to load optimizer from checkpoint {}. "
                  "Specify --no-load-optim or --finetune to prevent "
                  "attempting to load the optimizer state, "
                  "exiting ...".format(checkpoint_name)
              )
              sys.exit()
      else:
          if (args.fp16 or args.bf16) and optimizer is not None:
              optimizer.reload_model_params()
  
      # rng states.
      if not release and not args.finetune and not args.no_load_rng:
          try:
              if "rng_state" in state_dict:
                  # access rng_state for data parallel rank
                  if args.data_parallel_random_init:
                      rng_state = state_dict["rng_state"][
                          mpu.get_data_parallel_rank()
                      ]
                  else:
                      rng_state = state_dict["rng_state"][0]
                  random.setstate(rng_state["random_rng_state"])
                  np.random.set_state(rng_state["np_rng_state"])
                  torch.set_rng_state(rng_state["torch_rng_state"])
                  torch.cuda.set_rng_state(rng_state["cuda_rng_state"])
                  # Check for empty states array
                  if not rng_state["rng_tracker_states"]:
                      raise KeyError
                  tensor_parallel.get_cuda_rng_tracker().set_states(
                      rng_state["rng_tracker_states"]
                  )
              else:  # backward compatability
                  random.setstate(state_dict["random_rng_state"])
                  np.random.set_state(state_dict["np_rng_state"])
                  torch.set_rng_state(state_dict["torch_rng_state"])
                  torch.cuda.set_rng_state(state_dict["cuda_rng_state"])
                  # Check for empty states array
                  if not state_dict["rng_tracker_states"]:
                      raise KeyError
                  tensor_parallel.get_cuda_rng_tracker().set_states(
                      state_dict["rng_tracker_states"]
                  )
          except KeyError:
              print_rank_0(
                  "Unable to load rng state from checkpoint {}. "
                  "Specify --no-load-rng or --finetune to prevent "
                  "attempting to load the rng state, "
                  "exiting ...".format(checkpoint_name)
              )
              sys.exit()
  
      # Some utilities want to load a checkpoint without distributed being initialized
      if torch.distributed.is_initialized():
          torch.distributed.barrier()
  
      print_rank_0(
          f"  successfully loaded checkpoint from {args.load} "
          f"at iteration {iteration}"
      )
  
      return iteration
  
imports: import os
  import random
  import sys
  import numpy
  import torch
  from megatron import update_num_microbatches
  from megatron.core import mpu, tensor_parallel
  from .global_vars import get_args
  from .utils import unwrap_model, print_rank_0
label_desc: Use image node_id 16 for calling a global function with example usage: load_checkpoint(model, optimizer, opt_param_scheduler, load_arg, strict) and returns: iteration, int