code: def _save_optimizer_state(args, ds_checkpoint):
      sharded_states = [
          BASE_OPTIMIZER_STATE,
          PARAM_SLICE_MAPPINGS,
          SINGLE_PARTITION_OF_FP32_GROUPS,
      ]
      sd = ds_checkpoint.get_zero_checkpoint_state(
          pp_index=0, tp_index=0, dp_index=0
      )
  
      optim_sd = sd[OPTIMIZER_STATE_DICT]
      output_sd = {
          k: v for k, v in optim_sd.items() if k not in sharded_states
      }
      zero_output_folder = os.path.join(args.output_folder, "zero")
      output_file_path = os.path.join(
          zero_output_folder, f"optimizer_state.pt"
      )
      _save_checkpoint(output_file_path, output_sd)
  
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 12 for calling a global function with example usage: _save_optimizer_state(args, ds_checkpoint) without return types