code: def save_checkpoint(iteration, model, optimizer, opt_param_scheduler):
      """Save a model checkpoint."""
      args = get_args()
  
      # Only rank zero of the data parallel writes to the disk.
      model = unwrap_model(model)
  
      print_rank_0(
          "saving checkpoint at iteration {:7d} to {}".format(
              iteration, args.save
          )
      )
  
      # Collect rng state across data parallel ranks.
      rng_state = get_rng_state()
  
      # Checkpoint name.
      checkpoint_name = get_checkpoint_name(args.save, iteration)
  
      # Save distributed optimizer's custom parameter state.
      if (
          args.use_distributed_optimizer
          and not args.no_save_optim
          and optimizer is not None
      ):
          optim_checkpoint_name = (
              get_distributed_optimizer_checkpoint_name(checkpoint_name)
          )
          ensure_directory_exists(optim_checkpoint_name)
          optimizer.save_parameter_state(optim_checkpoint_name)
  
      # Collect args, model, RNG.
      if (
          not torch.distributed.is_initialized()
          or mpu.get_data_modulo_expert_parallel_rank() == 0
      ):
          # Arguments, iteration, and model.
          state_dict = {}
          state_dict["args"] = args
          state_dict["checkpoint_version"] = 3.0
          state_dict["iteration"] = iteration
          if len(model) == 1:
              state_dict["model"] = model[
                  0
              ].state_dict_for_save_checkpoint()
          else:
              for i in range(len(model)):
                  mpu.set_virtual_pipeline_model_parallel_rank(i)
                  state_dict["model%d" % i] = model[
                      i
                  ].state_dict_for_save_checkpoint()
  
          # Optimizer stuff.
          if not args.no_save_optim:
              if optimizer is not None:
                  state_dict["optimizer"] = optimizer.state_dict()
              if opt_param_scheduler is not None:
                  state_dict[
                      "opt_param_scheduler"
                  ] = opt_param_scheduler.state_dict()
  
          # RNG states.
          if not args.no_save_rng:
              state_dict["rng_state"] = rng_state
  
          # Save.
          ensure_directory_exists(checkpoint_name)
          torch.save(state_dict, checkpoint_name)
  
      # Wait so everyone is done (necessary)
      if torch.distributed.is_initialized():
          torch.distributed.barrier()
  
      print_rank_0(
          "  successfully saved checkpoint at iteration {:7d} to {}".format(
              iteration, args.save
          )
      )
  
      # And update the latest iteration
      if (
          not torch.distributed.is_initialized()
          or torch.distributed.get_rank() == 0
      ):
          tracker_filename = get_checkpoint_tracker_filename(args.save)
          with open(tracker_filename, "w") as f:
              f.write(str(iteration))
  
      # Wait so everyone is done (not necessary)
      if torch.distributed.is_initialized():
          torch.distributed.barrier()
  
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 11 for calling a global function with example usage: save_checkpoint(iteration, model, optimizer, opt_param_scheduler) without return types