code: def main():
      print(f"Convert DeepSpeed Checkpoint to Universal Checkpoint")
  
      args = parse_arguments()
      print(
          f"Converting DeepSpeed checkpoint in {args.input_folder} to Universal checkpoint in {args.output_folder}"
      )
  
      ds_checkpoint = DeepSpeedCheckpoint(args.input_folder)
      _check_for_required_state(ds_checkpoint)
  
      iteration = ds_checkpoint.get_iteration()
      # _create_latest_file(args.output_folder, iteration)
      checkpoint_paths = _create_checkpoint_paths(
          args.output_folder,
          iteration,
          ds_checkpoint.tp_degree,
          ds_checkpoint.pp_degree,
      )
  
      slice_shapes = []
      for mp_rank_file in ds_checkpoint.mp_rank_files:
          mp_sd = torch.load(
              mp_rank_file, map_location=torch.device("cpu")
          )
          slice_shapes += mp_sd[PARAM_SHAPES]
  
      # fix back to normal flat dict, merge duplicates for tp>1
      slice_shapes = dict(
          (k, v) for d in slice_shapes for k, v in d.items()
      )
      temp_dir = os.path.join(args.output_folder, "tmp")
  
      print("*** 1. Extracting ZeRO fragments")
      _extract_zero_shard_files(args, ds_checkpoint, temp_dir)
  
      print("*** 2. Merging slices .....")
      _merge_tp_slice_files(args, ds_checkpoint, slice_shapes, temp_dir)
  
      print("*** 3. Saving common optimizer states")
      _save_optimizer_state(args, ds_checkpoint)
  
      if not args.keep_temp_folder:
          shutil.rmtree(temp_dir, ignore_errors=True)
  
      # Copy mp* files into output folder
      for f in glob.glob(os.path.join(args.input_folder, "mp*")):
          shutil.copy2(f, args.output_folder)
  
      # Update latest to output folder
      checkpoint_root_folder, step_folder = os.path.split(
          args.output_folder
      )
      latest_file = os.path.join(
          checkpoint_root_folder, "latest_universal"
      )
      with open(latest_file, "w") as f:
          f.write(step_folder)
  
      print("*** Done!")
  
imports: from functools import partial
  import argparse
  import glob
  import itertools
  import multiprocessing
  import os
  import re
  import shutil
  import torch
  import tqdm
  from deepspeed.checkpoint import DeepSpeedCheckpoint
  from deepspeed.checkpoint import OPTIMIZER_STATE_DICT, BASE_OPTIMIZER_STATE, SINGLE_PARTITION_OF_FP32_GROUPS, PARAM_SLICE_MAPPINGS, PARAM_SHAPES, PARAM, CAT_DIM, PARAM_N_SUB_PARAMS, VOCAB_TENSOR, UNIVERSAL_CHECKPOINT_INFO, VOCABULARY_PARAMETER_PATTERNS, PIPELINE_REPLICATED_PARAMETER_PATTERNS, TP_REPLICATED_PARAMETER_PATTERNS, PARAMETER_TO_AVERAGE_PATTERNS, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0
label_desc: Use image node_id 14 for calling a global function with example usage: main() without return types