code: def parse_optim_states(files, ds_checkpoint_dir):
      total_files = len(files)
      state_dicts = []
      for f in files:
          state_dict = torch.load(f, map_location=device)
          # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
          # and also handle the case where it was already removed by another helper script
          state_dict["optimizer_state_dict"].pop(
              "optimizer_state_dict", None
          )
          state_dicts.append(state_dict)
  
      if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
          raise ValueError(f"{files[0]} is not a zero checkpoint")
      zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
      world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
  
      # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
      # parameters can be different from data parallelism for non-expert parameters. So we can just
      # use the max of the partition_count to get the dp world_size.
  
      if type(world_size) is list:
          world_size = max(world_size)
  
      if world_size != total_files:
          raise ValueError(
              f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
              "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
          )
  
      # the groups are named differently in each stage
      if zero_stage <= 2:
          fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
      elif zero_stage == 3:
          fp32_groups_key = FP32_FLAT_GROUPS
      else:
          raise ValueError(f"unknown zero stage {zero_stage}")
  
      if zero_stage <= 2:
          fp32_flat_groups = [
              state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key]
              for i in range(len(state_dicts))
          ]
      elif zero_stage == 3:
          # if there is more than one param group, there will be multiple flattened tensors - one
          # flattened tensor per group - for simplicity merge them into a single tensor
          #
          # XXX: could make the script more memory efficient for when there are multiple groups - it
          # will require matching the sub-lists of param_shapes for each param group flattened tensor
  
          fp32_flat_groups = [
              torch.cat(
                  state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key],
                  0,
              )
              for i in range(len(state_dicts))
          ]
  
      return zero_stage, world_size, fp32_flat_groups
  
imports: import argparse
  import torch
  import glob
  import math
  import os
  import re
  from collections import OrderedDict
  from dataclasses import dataclass
  from deepspeed.utils import logger
  from deepspeed.checkpoint.constants import DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS, FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES, FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS
label_desc: Use image node_id 8 for calling a global function with example usage: parse_optim_states(files, ds_checkpoint_dir) and returns: zero_stage,  world_size,  fp32_flat_groups